如何在TensorFlow Object Detection API中同时使用两个检测模型

Great question! To detect both the 90 COCO classes and your custom objects in a single workflow, you'll need to load both models, run inference with each, combine their results (making sure class IDs don't overlap), and visualize everything together. Here's a complete, working code snippet to do that:

Step-by-Step Code Implementation #

First, ensure you have the TensorFlow Object Detection API utilities (like label_map_util and visualization_utils) available in your environment. Then follow this code:

1. Import Required Libraries #

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

2. Helper Functions #

def load_detection_model(model_path):
    """Load a frozen TensorFlow detection graph"""
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.compat.v1.GraphDef()
        with tf.io.gfile.GFile(model_path, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
    return detection_graph

def run_inference(image, sess, image_tensor, boxes_tensor, scores_tensor, classes_tensor, num_detections_tensor, class_id_offset=0):
    """Run inference on an image and return filtered high-confidence detections"""
    # Expand dimensions to match model input shape [1, None, None, 3]
    image_expanded = np.expand_dims(image, axis=0)
    
    # Run inference
    (boxes, scores, classes, num) = sess.run(
        [boxes_tensor, scores_tensor, classes_tensor, num_detections_tensor],
        feed_dict={image_tensor: image_expanded})
    
    # Clean up output tensors by removing extra dimensions
    boxes = np.squeeze(boxes)
    scores = np.squeeze(scores)
    classes = np.squeeze(classes).astype(np.int32)
    
    # Offset class IDs to avoid overlap with COCO's 1-90 range
    if class_id_offset > 0:
        classes += class_id_offset
    
    # Filter out low-confidence detections (adjust threshold as needed)
    threshold = 0.5
    valid_mask = scores >= threshold
    return boxes[valid_mask], scores[valid_mask], classes[valid_mask]

3. Load Models and Label Maps #

# --------------------------
# Update these paths to match your setup!
# --------------------------
SSD_MODEL_PATH = 'path/to/ssd_mobilenet_v1_coco_2017_11_17/frozen_inference_graph.pb'
FASTER_RCNN_MODEL_PATH = 'path/to/faster_rcnn_inception_v2_coco_2018_01_28/frozen_inference_graph.pb'
COCO_LABEL_MAP = 'object_detection/data/mscoco_label_map.pbtxt'
CUSTOM_LABEL_MAP = 'path/to/your/custom_label_map.pbtxt'
INPUT_IMAGE_PATH = 'path/to/your/input_image.jpg'
OUTPUT_IMAGE_PATH = 'path/to/save/output_image.jpg'

# Load both detection models
ssd_graph = load_detection_model(SSD_MODEL_PATH)
frcnn_graph = load_detection_model(FASTER_RCNN_MODEL_PATH)

# Create TensorFlow sessions for each model
ssd_sess = tf.compat.v1.Session(graph=ssd_graph)
frcnn_sess = tf.compat.v1.Session(graph=frcnn_graph)

# Get input/output tensors for SSD model
ssd_image_tensor = ssd_graph.get_tensor_by_name('image_tensor:0')
ssd_boxes = ssd_graph.get_tensor_by_name('detection_boxes:0')
ssd_scores = ssd_graph.get_tensor_by_name('detection_scores:0')
ssd_classes = ssd_graph.get_tensor_by_name('detection_classes:0')
ssd_num_detections = ssd_graph.get_tensor_by_name('num_detections:0')

# Get input/output tensors for Faster RCNN model
frcnn_image_tensor = frcnn_graph.get_tensor_by_name('image_tensor:0')
frcnn_boxes = frcnn_graph.get_tensor_by_name('detection_boxes:0')
frcnn_scores = frcnn_graph.get_tensor_by_name('detection_scores:0')
frcnn_classes = frcnn_graph.get_tensor_by_name('detection_classes:0')
frcnn_num_detections = frcnn_graph.get_tensor_by_name('num_detections:0')

# Load and combine label maps
# COCO uses IDs 1-90, so we offset custom classes to start at 91
coco_categories = label_map_util.create_category_index_from_labelmap(COCO_LABEL_MAP, use_display_name=True)
custom_categories = label_map_util.create_category_index_from_labelmap(CUSTOM_LABEL_MAP, use_display_name=True)

combined_categories = {}
# Add COCO classes first
combined_categories.update(coco_categories)
# Add custom classes with offset IDs
offset = 90
for idx, cat in custom_categories.items():
    new_idx = idx + offset
    combined_categories[new_idx] = {'id': new_idx, 'name': cat['name']}

4. Run Inference and Combine Results #

# Load and preprocess input image
image = cv2.imread(INPUT_IMAGE_PATH)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# Run inference with SSD (detects COCO classes)
ssd_det_boxes, ssd_det_scores, ssd_det_classes = run_inference(
    image_rgb, ssd_sess, ssd_image_tensor, ssd_boxes, ssd_scores, ssd_classes, ssd_num_detections
)

# Run inference with Faster RCNN (detects custom classes, IDs offset by 90)
frcnn_det_boxes, frcnn_det_scores, frcnn_det_classes = run_inference(
    image_rgb, frcnn_sess, frcnn_image_tensor, frcnn_boxes, frcnn_scores, frcnn_classes, frcnn_num_detections,
    class_id_offset=offset
)

# Combine all valid detections into single arrays
combined_boxes = np.concatenate([ssd_det_boxes, frcnn_det_boxes], axis=0)
combined_scores = np.concatenate([ssd_det_scores, frcnn_det_scores], axis=0)
combined_classes = np.concatenate([ssd_det_classes, frcnn_det_classes], axis=0)

5. Visualize and Save Results #

# Draw combined detections on the image
vis_util.visualize_boxes_and_labels_on_image_array(
    image_rgb,
    combined_boxes,
    combined_classes,
    combined_scores,
    combined_categories,
    use_normalized_coordinates=True,
    line_thickness=4,
    min_score_thresh=0.5
)

# Convert back to BGR for OpenCV compatibility
image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)

# Display the result
cv2.imshow('Combined Object Detection', image_bgr)
cv2.waitKey(0)
cv2.destroyAllWindows()

# Save the output image
cv2.imwrite(OUTPUT_IMAGE_PATH, image_bgr)

# Clean up sessions to free resources
ssd_sess.close()
frcnn_sess.close()

Key Notes #

• Class ID Offset: The most important step is offsetting your custom model's class IDs by 90. This ensures no overlap between COCO's 1-90 IDs and your custom class IDs, preventing labeling confusion.
• Threshold Adjustment: The 0.5 confidence threshold can be tuned up (for stricter detections) or down (for more detections) based on your needs.
• TF Version: This code uses TensorFlow 1.x compatibility mode because the models you referenced are frozen graphs from older TF versions. If you're using TF 2.x, consider converting these models to SavedModel format for better performance.
• Performance: Running two models sequentially will take longer than a single model, but it's the simplest way to combine their outputs without retraining a new multi-class model.

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