如何用TensorFlow提取含检测目标的视频帧？已训练行人检测器

Hey there! Since you’re already up and running with your custom pedestrian detector, let’s turn your pseudo-code ideas into working code to save only frames that have detected people. There are two straightforward ways to do this—let’s break them down.

Method 1: Check Detection Results Directly in Your Main Loop #

This approach keeps your vis_util function untouched and handles the saving logic right in your video processing loop. It’s great if you want to keep utility code unmodified.

1. First, initialize a frame counter at the top of your session block:

count = 0

2. After running the detection with sess.run(), add logic to filter valid pedestrian detections (we’ll use a confidence threshold like 0.5 to skip low-confidence hits):

# Process squeezed detection results
squeezed_classes = np.squeeze(classes).astype(np.int32)
squeezed_scores = np.squeeze(scores)

# Check if any detection is a person (class 1) with confidence > 0.5
has_person = any((cls == 1 and score > 0.5) for cls, score in zip(squeezed_classes, squeezed_scores))

if has_person:
    cv2.imwrite(f"frame{count}.jpg", image_np)
    count += 1

3. Here’s how the full integrated code looks:

import cv2
import numpy as np
import tensorflow as tf
from object_detection.utils import vis_util, label_map_util

# Assume you've already loaded detection_graph and category_index
cap = cv2.VideoCapture('http://192.168.15.8:8080/video')
with detection_graph.as_default():
    with tf.Session(graph=detection_graph) as sess:
        ret = True
        count = 0  # Initialize frame counter
        while ret:
            ret, image_np = cap.read()
            if not ret:
                break  # Exit loop when no more frames
            
            # Expand dimensions for the model
            image_np_expanded = np.expand_dims(image_np, axis=0)
            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 detection
            (boxes, scores, classes, num_detections) = sess.run(
                [boxes, scores, classes, num_detections],
                feed_dict={image_tensor: image_np_expanded})
            
            # Check for valid pedestrian detections
            squeezed_classes = np.squeeze(classes).astype(np.int32)
            squeezed_scores = np.squeeze(scores)
            has_person = any((cls == 1 and score > 0.5) for cls, score in zip(squeezed_classes, squeezed_scores))
            
            # Save frame if person is detected
            if has_person:
                cv2.imwrite(f"frame{count}.jpg", image_np)
                count += 1
            
            # Visualize detection results
            vis_util.visualize_boxes_and_labels_on_image_array(
                image_np,
                np.squeeze(boxes),
                squeezed_classes,
                squeezed_scores,
                category_index,
                use_normalized_coordinates=True,
                line_thickness=8)
            
            cv2.imshow('image', cv2.resize(image_np,(600,400)))
            if cv2.waitKey(25) & 0xFF == ord('q'):
                cv2.destroyAllWindows()
                cap.release()
                break

Method 2: Modify the vis_util Visualization Function #

This is the approach you mentioned trying, and it ties frame saving directly to when the visualization draws a person label. Here’s how to implement it cleanly:

1. Locate visualize_boxes_and_labels_on_image_array in object_detection/utils/visualization_utils.py, then add a count parameter and saving logic inside the detection loop:

def visualize_boxes_and_labels_on_image_array(
    image,
    boxes,
    classes,
    scores,
    category_index,
    # ... keep existing parameters ...
    count=None  # Add this parameter to pass the counter
):
    # ... keep existing code above ...
    
    for i in range(min(max_boxes_to_draw, boxes.shape[0])):
        if scores is None or scores[i] > min_score_thresh:
            # ... keep existing code for box/mask handling ...
            
            if not agnostic_mode:
                class_id = int(classes[i])
                if class_id not in category_index.keys():
                    continue
                class_name = category_index[class_id]['name']
                # Add your frame-saving logic here
                if class_name == 'person' and count is not None:
                    cv2.imwrite(f"frame{count[0]}.jpg", image)
                    count[0] += 1  # Use a list to modify the counter in-place
            # ... keep remaining visualization code ...

2. In your main script, initialize the counter as a list (since integers are immutable in Python, a list lets the function update its value):

count = [0]

3. Pass the counter when calling the visualization function:

vis_util.visualize_boxes_and_labels_on_image_array(
    image_np,
    np.squeeze(boxes),
    squeezed_classes,
    squeezed_scores,
    category_index,
    use_normalized_coordinates=True,
    line_thickness=8,
    count=count  # Pass the counter here
)

Both methods work perfectly—Method 1 is cleaner if you want to avoid modifying TensorFlow’s utility code, while Method 2 directly links saving to successful visualization. Since you mentioned you already got Method 2 working, that’s a totally valid solution!

内容的提问来源于stack exchange，提问作者Ryan Christopher Martinez Fede