### 使用 YOLOv8 和 TensorRT 的 `.engine` 文件实现 USB 摄像头实时检测的代码 以下是一个完整的 Python 代码示例，展示如何利用 YOLOv8 和 TensorRT 的 `.engine` 文件对 USB 摄像头进行实时目标检测。 --- #### 1. **环境准备** 确保已安装以下依赖项： - OpenCV (`pip install opencv-python`) - PyCUDA (`pip install pycuda`) - TensorRT（通过 NVIDIA 提供的包安装） - NumPy (`pip install numpy`) --- #### 2. **代码实现** 以下是基于 Python 的完整代码： ```python import cv2 import numpy as np import tensorrt as trt import pycuda.autoinit import pycuda.driver as cuda import time # 定义辅助类用于管理主机和设备内存 class HostDeviceMem(object): def __init__(self, host_mem, device_mem): self.host = host_mem self.device = device_mem def allocate_buffers(engine): """为推理分配输入输出缓冲区""" inputs = [] outputs = [] bindings = [] stream = cuda.Stream() for binding in engine: size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size dtype = trt.nptype(engine.get_binding_dtype(binding)) # 分配主机和设备端内存 host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(host_mem.nbytes) # 添加到绑定列表 bindings.append(int(device_mem)) # 判断是输入还是输出 if engine.binding_is_input(binding): inputs.append(HostDeviceMem(host_mem, device_mem)) else: outputs.append(HostDeviceMem(host_mem, device_mem)) return inputs, outputs, bindings, stream def do_inference(context, bindings, inputs, outputs, stream, batch_size=1): """执行推理操作""" [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle) [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] stream.synchronize() # 图片预处理函数 def preprocess(image, shape=(640, 640)): resized = cv2.resize(image, shape) normalized = resized / 255.0 transposed = np.transpose(normalized, (2, 0, 1)).astype(np.float32) return transposed[np.newaxis, ...] # 结果后处理函数 def postprocess(output_data, conf_thres=0.5, iou_thres=0.45): predictions = output_data.reshape((1, -1, 85)) boxes = predictions[:, :, :4] scores = predictions[:, :, 4:5] classes = predictions[:, :, 5:] class_ids = np.argmax(classes, axis=-1).reshape(-1, 1) selected_indices = cv2.dnn.NMSBoxes( bboxes=np.squeeze(boxes), scores=np.squeeze(scores), score_threshold=conf_thres, nms_threshold=iou_thres ) results = [] if len(selected_indices) > 0: for idx in selected_indices.flatten(): box = boxes[0][idx] confidence = scores[0][idx] cls_id = class_ids[idx] results.append({ "box": box.tolist(), "confidence": float(confidence), "cls_id": int(cls_id) }) return results # 可视化结果 def draw_boxes(frame, detections, labels=None): for det in detections: x_min, y_min, x_max, y_max = map(int, det["box"]) label = f"{labels[det['cls_id']]} {det['confidence']:.2f}" if labels else str(det['cls_id']) color = (0, 255, 0) cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2) cv2.putText(frame, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2) return frame # 主程序逻辑 if __name__ == "__main__": TRT_LOGGER = trt.Logger(trt.Logger.WARNING) with open("yolov8n.trt", "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) inputs, outputs, bindings, stream = allocate_buffers(engine) context = engine.create_execution_context() cap = cv2.VideoCapture(0) # 打开默认USB摄像头 while True: ret, frame = cap.read() if not ret: break # 预处理帧数据 processed_frame = preprocess(frame) inputs[0].host = processed_frame.ravel().astype(np.float32) # 进行推理 start_time = time.time() do_inference(context, bindings, inputs, outputs, stream) end_time = time.time() fps = 1 / (end_time - start_time) print(f"Inference FPS: {fps:.2f}") # 后处理输出 output_data = outputs[0].host.reshape((1, -1, 85)) detections = postprocess(output_data) # 绘制边界框 labeled_frame = draw_boxes(frame.copy(), detections) cv2.imshow("YOLOv8 Detection", labeled_frame) if cv2.waitKey(1) & 0xFF == ord('q'): break cap.release() cv2.destroyAllWindows() ``` --- #### 3. **说明** - 上述代码实现了从读取 USB 摄像头数据到使用 TensorRT 加速的目标检测全过程。 - `preprocess()` 函数负责将原始图像缩放到模型所需的尺寸，并对其进行标准化和通道变换。 - `postprocess()` 函数解析网络输出，筛选出高置信度的检测结果，并应用 NMS 抑制冗余框[^1]。 - `draw_boxes()` 函数在原图上绘制检测框及其类别标签[^2]。 --- #### 4. **注意事项** - 确保 `.engine` 文件路径正确无误。 - 如果需要支持不同的分辨率或批量大小，请修改 `allocate_buffers()` 函数中的相关部分。 - 对于 Jetson Nano 设备，建议降低输入分辨率以减少计算负担[^5]。 --- ###