如何从TensorFlow Object Detection API提取标签并存入数组

Got it, let's adjust that detect_objects function to capture detected labels into an array instead of rendering them on video. Here's a straightforward implementation with explanations:

def detect_objects(image_np, sess, detection_graph, category_index):
    # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
    image_np_expanded = np.expand_dims(image_np, axis=0)
    
    # Fetch core detection tensors from the graph
    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
    boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
    scores = detection_graph.get_tensor_by_name('detection_scores:0')
    classes = detection_graph.get_tensor_by_name('detection_classes:0')
    num_detections = detection_graph.get_tensor_by_name('num_detections:0')
    
    # Run inference to get detection results
    (boxes, scores, classes, num_detections) = sess.run(
        [boxes, scores, classes, num_detections],
        feed_dict={image_tensor: image_np_expanded})
    
    # Initialize empty array to store valid detected labels
    detected_labels = []
    
    # Filter detections by confidence threshold (tweak this value as needed)
    confidence_threshold = 0.5
    for i in range(int(num_detections[0])):
        if scores[0][i] > confidence_threshold:
            # Map class ID to human-readable label using category index
            class_id = int(classes[0][i])
            label = category_index[class_id]['name']
            # Add label to our array
            detected_labels.append(label)
            # Optional: Store labels with their confidence scores if needed
            # detected_labels.append({"label": label, "confidence": round(scores[0][i], 2)})
    
    return detected_labels

Key Details to Note #

• category_index Parameter: This is the label mapping (like the COCO category index included in TensorFlow Object Detection API samples) that converts numeric class IDs to human-readable labels (e.g., "person", "car").
• Confidence Threshold: We filter out low-confidence detections (set to 0.5 here) to avoid noisy results—you can raise or lower this based on how strict you want the detection to be.
• No Visualization Code: All code related to drawing boxes/text on the image is removed, since we're only focused on collecting labels.

Example Usage #

Once you've loaded your model and category index, call the function like this:

# After loading detection_graph, sess, and category_index...
detected_labels = detect_objects(image_np, sess, detection_graph, category_index)
print("Detected objects:", detected_labels)
# Sample output: ["person", "dog", "bicycle"]

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