1. 图像分类入门

任务目标

在Jetson Nano上构建一个基于TensorRT加速的MobileNetV2图像分类程序，读取CSI或USB摄像头实时画面，对画面中央区域物体进行分类，并在输出窗口叠加显示类别名称与置信度。本案例聚焦于TensorRT引擎构建与实时推理管线两个核心环节。

环境与依赖

  • Jetson Nano（JetPack 4.6+），已安装TensorRT 8.x、OpenCV 4.x（自带于JetPack）
  • Python 3.6+，需安装pycuda（pip3 install pycuda）
  • 预训练模型：resnet50_mobilenetv2.uff（NVIDIA官方提供，对应ImageNet 1000类）
  • 类别标签文件：imagenet_labels.txt（可从TensorRT官方样例获取）

核心实现

以下代码完整实现了从UFF模型构建TensorRT引擎、预处理、推理到结果可视化的全流程。引擎构建仅需首次运行执行，之后会缓存为.trt文件加速加载。

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

# ---------- TensorRT 引擎构建 ----------
class TRTEngine:
    def __init__(self, uff_path, trt_path, input_name="input_1", 
                 input_shape=(3, 224, 224), output_name="dense_1"):
        self.input_name = input_name
        self.output_name = output_name
        self.input_shape = input_shape
        self.trt_path = trt_path
        
        if os.path.exists(trt_path):
            self.engine = self._load_engine(trt_path)
        else:
            self.engine = self._build_engine(uff_path, trt_path)
        
        self.context = self.engine.create_execution_context()
        self._allocate_buffers()
    
    def _build_engine(self, uff_path, trt_path):
        logger = trt.Logger(trt.Logger.WARNING)
        builder = trt.Builder(logger)
        network = builder.create_network()
        parser = trt.UffParser()
        parser.register_input(self.input_name, self.input_shape)
        parser.register_output(self.output_name)
        parser.parse(uff_path, network)
        
        builder.max_workspace_size = 1 << 30  # 1GB
        builder.fp16_mode = True  # 启用FP16加速
        engine = builder.build_cuda_engine(network)
        with open(trt_path, "wb") as f:
            f.write(engine.serialize())
        return engine
    
    def _load_engine(self, trt_path):
        logger = trt.Logger(trt.Logger.WARNING)
        with open(trt_path, "rb") as f:
            engine_data = f.read()
        runtime = trt.Runtime(logger)
        return runtime.deserialize_cuda_engine(engine_data)
    
    def _allocate_buffers(self):
        self.h_input = cuda.pagelocked_empty(
            trt.volume(self.input_shape), dtype=np.float32)
        self.h_output = cuda.pagelocked_empty(
            trt.volume((1000,)), dtype=np.float32)
        self.d_input = cuda.mem_alloc(self.h_input.nbytes)
        self.d_output = cuda.mem_alloc(self.h_output.nbytes)
    
    def infer(self, img):
        # 预处理：resize -> HWC转CHW -> 归一化到[-1,1]
        img = cv2.resize(img, (self.input_shape[2], self.input_shape[1]))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 127.5 - 1.0
        img = np.transpose(img, (2, 0, 1)).ravel()
        np.copyto(self.h_input, img)
        
        cuda.memcpy_htod(self.d_input, self.h_input)
        self.context.execute(bindings=[int(self.d_input), int(self.d_output)])
        cuda.memcpy_dtoh(self.h_output, self.d_output)
        return self.h_output.copy()

# ---------- 类别标签加载 ----------
def load_labels(path):
    with open(path, "r") as f:
        return [line.strip() for line in f.readlines()]

# ---------- 主程序 ----------
def main():
    # 配置路径（请根据实际位置修改）
    UFF_PATH = "resnet50_mobilenetv2.uff"
    TRT_PATH = "mobilenetv2.trt"
    LABEL_PATH = "imagenet_labels.txt"
    
    labels = load_labels(LABEL_PATH)
    engine = TRTEngine(UFF_PATH, TRT_PATH)
    
    # 打开摄像头（CSI摄像头用 nvarguscamerasrc，USB用 0）
    cap = cv2.VideoCapture("nvarguscamerasrc ! video/x-raw(memory:NVMM), width=640, height=480, format=NV12, framerate=30/1 ! nvvidconv ! video/x-raw, format=BGR ! appsink", cv2.CAP_GSTREAMER)
    if not cap.isOpened():
        cap = cv2.VideoCapture(0)  # 回退到USB摄像头
    
    print("推理开始，按Q退出")
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # 取画面中央224x224区域
        h, w = frame.shape[:2]
        cx, cy = w//2, h//2
        crop = frame[cy-112:cy+112, cx-112:cx+112]
        
        # 推理
        start = time.time()
        output = engine.infer(crop)
        idx = np.argmax(output)
        conf = float(output[idx])
        fps = 1.0 / (time.time() - start)
        
        # 可视化
        label = f"{labels[idx]}: {conf:.2%}"
        cv2.rectangle(frame, (cx-112, cy-112), (cx+112, cy+112), (0,255,0), 2)
        cv2.putText(frame, label, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2)
        cv2.putText(frame, f"FPS: {fps:.1f}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
        
        cv2.imshow("TensorRT MobileNetV2", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

运行与验证

  1. 将上述代码保存为trt_mobilenet.py，与.uff模型和标签文件置于同一目录
  2. 首次运行会自动构建TensorRT引擎（约需1-2分钟），生成mobilenetv2.trt缓存文件
  3. 执行python3 trt_mobilenet.py启动程序
  4. 通过标准：摄像头画面出现绿色框，框内物体被正确识别（如对准杯子显示"cup"、对准人显示"person"），置信度高于60%，FPS达到15以上（FP16模式）

常见问题

  • 摄像头打不开：确认CSI摄像头排线连接牢固，或改用USB摄像头（将VideoCapture参数改为0）
  • 构建引擎报错：检查UFF模型路径是否正确，JetPack版本需包含对应TensorRT版本
  • 识别不准确：确保画面中央区域包含完整物体，光照充足；ImageNet类别偏向日常物体，极端角度或遮挡会影响效果
← 上一章 📖 返回目录 下一章 →