📚 课程目录第2章目标检测实战

2. 目标检测实战

任务目标

在Jetson Nano上使用TensorFlow的SSD-MobileNet预训练模型（COCO数据集，含person类别），对USB摄像头或视频文件进行实时行人检测，并在画面中绘制边界框与置信度。本案例聚焦于推理落地，不涉及训练。

环境与依赖

Jetson Nano（JetPack 4.6+，已安装CUDA/cuDNN）
Python 3.6+，TensorFlow 1.15（JetPack自带）或TensorFlow 2.x（需自行安装兼容版本）
OpenCV（用于视频读取与图像绘制）
模型文件：ssd_mobilenet_v1_coco_2018_01_28（需手动下载并解压，获取frozen_inference_graph.pb）

完整实现代码

以下代码同时支持摄像头（默认）与视频文件输入。使用TensorFlow 1.x的GraphDef方式加载模型，以兼容Jetson上最常见的预装环境。

import cv2
import numpy as np
import tensorflow as tf
import sys

# ---------- 配置 ----------
MODEL_PATH = 'ssd_mobilenet_v1_coco_2018_01_28/frozen_inference_graph.pb'
INPUT_SOURCE = 0          # 0表示USB摄像头，或填写视频文件路径如 'test.mp4'
CONF_THRESHOLD = 0.5      # 置信度阈值
PERSON_CLASS_ID = 1       # COCO中person类别ID为1

# ---------- 加载模型 ----------
def load_graph(model_path):
    graph = tf.Graph()
    with 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 graph

print('[INFO] 加载模型中...')
detection_graph = load_graph(MODEL_PATH)

# ---------- 创建会话 ----------
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.compat.v1.Session(graph=detection_graph, config=config)

# 获取输入输出张量
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
boxes_tensor = detection_graph.get_tensor_by_name('detection_boxes:0')
scores_tensor = detection_graph.get_tensor_by_name('detection_scores:0')
classes_tensor = detection_graph.get_tensor_by_name('detection_classes:0')
num_detections = detection_graph.get_tensor_by_name('num_detections:0')

# ---------- 视频源 ----------
cap = cv2.VideoCapture(INPUT_SOURCE)
if not cap.isOpened():
    print('[错误] 无法打开视频源:', INPUT_SOURCE)
    sys.exit(1)

print('[INFO] 开始实时检测，按 Q 键退出')

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 预处理：转为RGB并增加batch维度
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame_expanded = np.expand_dims(frame_rgb, axis=0)

    # 推理
    (boxes, scores, classes, num) = sess.run(
        [boxes_tensor, scores_tensor, classes_tensor, num_detections],
        feed_dict={image_tensor: frame_expanded})

    # 解析结果并绘制
    height, width, _ = frame.shape
    num_det = int(num[0])
    for i in range(num_det):
        score = float(scores[0][i])
        cls = int(classes[0][i])
        if cls == PERSON_CLASS_ID and score >= CONF_THRESHOLD:
            ymin, xmin, ymax, xmax = boxes[0][i]
            x1 = int(xmin * width)
            x2 = int(xmax * width)
            y1 = int(ymin * height)
            y2 = int(ymax * height)
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label = f'Person: {score:.2f}'
            cv2.putText(frame, label, (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    cv2.imshow('SSD-MobileNet Person Detection', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
sess.close()

运行方式

下载模型并解压：
wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2018_01_28.tar.gz
tar -xzvf ssd_mobilenet_v1_coco_2018_01_28.tar.gz
将上述代码保存为 person_detection.py，确保模型路径正确。
运行：python3 person_detection.py
如需检测视频文件，将 INPUT_SOURCE 改为视频文件路径。

验证方法与判定标准

功能验证：程序启动后，画面窗口正常显示摄像头/视频画面。当行人出现在画面中时，绿色边界框实时跟随并标注“Person: 置信度”。
性能验证：观察画面流畅度，在Jetson Nano上应达到约10～15 FPS（受分辨率与光照影响）。若帧率过低，可降低输入分辨率（如将cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)）。
正确性验证：检测框应准确框住行人主体，且不会对非人物体（如车辆、桌子）产生持续误检。按Q键可正常退出。

常见问题

CUDA内存不足：在代码中已设置allow_growth=True，若仍报错，可尝试关闭其他GUI程序或降低输入分辨率。
TensorFlow 2.x兼容：若使用TF2，需将tf.compat.v1替换为tf，并确认模型仍以GraphDef方式加载（TF2中需使用tf.compat.v1.GraphDef）。
检测不到行人：确认摄像头画面光线充足，且行人距离在2～8米范围内（SSD-MobileNet对中近距离行人效果最佳）。