4. 实时图像分类

任务目标

在Jetson Nano上使用TensorRT加速的ResNet-18模型，对CSI或USB摄像头采集的实时视频流逐帧执行图像分类，并在画面中叠加显示Top-1类别名称与置信度。

环境依赖

  • Jetson Nano（JetPack 4.6+，已安装TensorRT 8.x）
  • Python 3.6+，依赖库：cv2（OpenCV 4.x）、numpy、pycuda（TensorRT Python绑定）
  • 预训练模型：ResNet-18 ONNX文件（可从PyTorch导出或使用ONNX Model Zoo）
  • ImageNet类别标签文件（imagenet_classes.txt，1000行）

核心实现

以下代码完成从ONNX到TensorRT引擎的构建、视频流捕获、预处理、推理及结果叠加的完整闭环。引擎首次运行会序列化为.trt文件，后续启动直接加载。

import cv2
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import time
import os

# ---------- TensorRT 引擎构建 ----------
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_path, trt_path, precision="fp16"):
    """构建或加载TensorRT引擎，支持FP16加速"""
    if os.path.exists(trt_path):
        with open(trt_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    
    # 构建引擎
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)
    
    with open(onnx_path, "rb") as f:
        assert parser.parse(f.read()), "ONNX解析失败"
    
    config = builder.create_builder_config()
    if precision == "fp16" and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    
    # 设置动态batch（1~8）
    profile = builder.create_optimization_profile()
    input_tensor = network.get_input(0)
    profile.set_shape(input_tensor.name, (1, 3, 224, 224), (1, 3, 224, 224), (8, 3, 224, 224))
    config.add_optimization_profile(profile)
    
    engine = builder.build_serialized_network(network, config)
    with open(trt_path, "wb") as f:
        f.write(engine)
    with trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(engine)

# ---------- 推理上下文封装 ----------
class TRTInference:
    def __init__(self, engine):
        self.engine = engine
        self.context = engine.create_execution_context()
        self.stream = cuda.Stream()
        
        # 分配输入输出缓冲区
        self.input_idx = engine.get_binding_index("input")
        self.output_idx = engine.get_binding_index("output")
        self.input_shape = engine.get_binding_shape(self.input_idx)
        self.output_shape = engine.get_binding_shape(self.output_idx)
        
        self.input_size = trt.volume(self.input_shape) * np.dtype(np.float32).itemsize
        self.output_size = trt.volume(self.output_shape) * np.dtype(np.float32).itemsize
        
        self.d_input = cuda.mem_alloc(self.input_size)
        self.d_output = cuda.mem_alloc(self.output_size)
        
        self.h_input = np.zeros(self.input_shape, dtype=np.float32)
        self.h_output = np.zeros(self.output_shape, dtype=np.float32)
    
    def infer(self, img):
        """输入预处理后的CHW float32数组，返回1000维logits"""
        self.h_input[0] = img
        cuda.memcpy_htod_async(self.d_input, self.h_input, self.stream)
        self.context.execute_async_v2(bindings=[int(self.d_input), int(self.d_output)], stream_handle=self.stream.handle)
        cuda.memcpy_dtoh_async(self.h_output, self.d_output, self.stream)
        self.stream.synchronize()
        return self.h_output[0]

# ---------- 图像预处理 ----------
def preprocess(frame):
    """BGR帧 → 归一化CHW，ResNet标准预处理"""
    img = cv2.resize(frame, (224, 224))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    img = (img - mean) / std
    img = np.transpose(img, (2, 0, 1))  # HWC → CHW
    return np.ascontiguousarray(img)

# ---------- 主程序 ----------
if __name__ == "__main__":
    # 配置路径
    ONNX_PATH = "resnet18.onnx"
    TRT_PATH = "resnet18_fp16.trt"
    LABELS_PATH = "imagenet_classes.txt"
    
    # 加载类别标签
    with open(LABELS_PATH) as f:
        labels = [line.strip() for line in f.readlines()]
    
    # 构建/加载引擎
    engine = build_engine(ONNX_PATH, TRT_PATH, precision="fp16")
    trt_infer = TRTInference(engine)
    
    # 打开摄像头（0=USB摄像头，或使用CSI: "nvarguscamerasrc ! video/x-raw(memory:NVMM), width=640, height=480, framerate=30/1 ! nvvidconv ! video/x-raw, format=BGRx ! videoconvert ! video/x-raw, format=BGR ! appsink"）
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    
    assert cap.isOpened(), "无法打开摄像头"
    
    # 预热（首次推理包含cuDNN自动调优）
    dummy = np.zeros((3, 224, 224), dtype=np.float32)
    trt_infer.infer(dummy)
    
    fps_counter = 0
    fps_time = time.time()
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # 预处理 + 推理
        input_blob = preprocess(frame)
        logits = trt_infer.infer(input_blob)
        
        # Softmax获取概率
        exp_logits = np.exp(logits - np.max(logits))
        probs = exp_logits / np.sum(exp_logits)
        top1_idx = np.argmax(probs)
        confidence = probs[top1_idx]
        
        # 显示结果
        label_text = f"{labels[top1_idx]}: {confidence:.2%}"
        cv2.putText(frame, label_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
        
        # 显示FPS
        fps_counter += 1
        if time.time() - fps_time >= 1.0:
            fps = fps_counter / (time.time() - fps_time)
            fps_counter = 0
            fps_time = time.time()
        cv2.putText(frame, f"FPS: {fps:.1f}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
        
        cv2.imshow("Real-time Classification", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()

运行与验证

  1. 将resnet18.onnx与imagenet_classes.txt置于脚本同目录，执行：
python3 realtime_classify.py
  1. 通过条件：
  • 窗口正常显示摄像头画面，画面中持续叠加类别名称与置信度（如“tabby cat: 92.3%”）；
  • 将摄像头对准常见物体（键盘、水杯、人等），类别标签随之正确变化；
  • 画面左上角FPS稳定在15以上（FP16精度下Jetson Nano典型性能）；
  • 按q键正常退出，无报错。

常见问题

  • 引擎构建失败：确认ONNX输入名为input、输出名为output，且batch维度为动态。若使用PyTorch导出，需设置dynamic_axes。
  • CSI摄像头无法打开：将cv2.VideoCapture(0)替换为GStreamer管道字符串（见代码注释），并确保已执行sudo apt install gstreamer1.0-plugins-good。
  • 首次运行较慢：TensorRT构建引擎及cuDNN自动调优需要1~2分钟，属正常现象，后续启动直接加载.trt文件。
← 上一章 📖 返回目录 下一章 →