如何修改YOLO目标检测代码以将未检测到目标的图片保存至"unclassified"文件夹

Got it, let's sort out this problem! The core issue here is that your current code only saves images when it finds valid detections meeting the confidence threshold. For images with no detected objects, there's no logic to handle them—so they get ignored. Here's how to fix this step by step:

Key Fixes to Implement #

1. Track valid detections per image: Add a flag to check if any valid targets were detected in the current image.
2. Handle unclassified case: After processing all detections for an image, if no valid targets were found, save it to an unclassified folder.
3. Clean up state per image: Move the initialization of boxes, confidences, and classIDs inside the image loop—your original code might be carrying over data from previous images, which is a hidden bug!

Modified Full Code #

import glob
import cv2
import numpy as np
import os

# Assume these variables are already defined elsewhere:
# net, ln, COLORS, LABELS, args (with "confidence" key), data_path

files = glob.glob(data_path)
for f1 in files:
    # Initialize per-image state (fixes hidden bug of accumulating old data)
    boxes = []
    confidences = []
    classIDs = []
    has_valid_detection = False  # Flag to track if we found any valid targets
    
    image = cv2.imread(f1)
    blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (416, 416), swapRB=True, crop=False)
    net.setInput(blob)
    layerOutputs = net.forward(ln)

    # Loop over each layer output
    for output in layerOutputs:
        # Loop over each detection
        for detection in output:
            scores = detection[5:]
            classID = np.argmax(scores)
            confidence = scores[classID]
            box = detection[0:4]
            (centerX, centerY, width, height) = box.astype("int")
            x = int(centerX - (width / 2))
            y = int(centerY - (height / 2))

            boxes.append([x, y, int(width), int(height)])
            confidences.append(float(confidence))
            classIDs.append(classID)

            # Filter weak predictions
            if confidence > args["confidence"]:
                has_valid_detection = True  # Mark this image as having valid detections
                i = len(confidences) - 1  # Get index of current detection
                (x, y) = (boxes[i][0], boxes[i][1])
                (w, h) = (boxes[i][2], boxes[i][3])
                color = [int(c) for c in COLORS[classIDs[i]]]
                cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
                text = "{}: {:.4f}".format(LABELS[classIDs[i]], confidences[i])
                cv2.putText(image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

                # Save image and label to class-specific folder
                class_dir = os.path.join('/path to folder/', LABELS[classID])
                if not os.path.exists(class_dir):
                    os.makedirs(class_dir)
                base_name = f1.split('/')[-1][:-4]
                path = os.path.join(class_dir, base_name)
                cv2.imwrite(path + '.jpeg', image)
                with open((path + '.txt'), 'a+') as f:
                    f.write(f"{classID} {box[0]} {box[1]} {box[2]} {box[3]}\n")

    # Handle images with no valid detections
    if not has_valid_detection:
        unclassified_dir = os.path.join('/path to folder/', 'unclassified')
        if not os.path.exists(unclassified_dir):
            os.makedirs(unclassified_dir)
        base_name = f1.split('/')[-1][:-4]
        save_path = os.path.join(unclassified_dir, f"{base_name}.jpeg")
        cv2.imwrite(save_path, image)
        # Optional: Create an empty txt file or add a note
        with open(os.path.join(unclassified_dir, f"{base_name}.txt"), 'w') as f:
            f.write("No objects detected\n")

What Changed? #

• Added has_valid_detection flag to track if any valid targets were found in the image.
• Moved boxes, confidences, classIDs initialization inside the image loop to avoid cross-image data contamination.
• Added a post-processing check: if has_valid_detection remains False, save the original (or unannotated) image to the unclassified folder, along with an optional txt file noting no detections.

内容的提问来源于stack exchange，提问作者happypumpkin