YOLOv8安全帽检测实战：从训练到部署，在Jetson Nano上跑起来

YOLOv8安全帽检测实战从训练到部署在Jetson Nano上跑起来在工业安全领域实时检测工人是否佩戴安全帽是一项至关重要的任务。传统的人工巡检方式不仅效率低下而且难以实现全天候监控。基于深度学习的计算机视觉技术为解决这一问题提供了新的可能性。本文将聚焦于如何将训练好的YOLOv8安全帽检测模型部署到NVIDIA Jetson Nano这样的边缘设备上实现真正的端到端解决方案。1. 模型转换与优化1.1 从PyTorch到TensorRT的转换YOLOv8模型训练完成后通常会得到一个.pt格式的权重文件。为了在Jetson Nano上获得最佳性能我们需要将其转换为TensorRT格式。TensorRT是NVIDIA推出的高性能深度学习推理框架能够显著提升模型在边缘设备上的运行速度。转换过程主要分为以下几个步骤将PyTorch模型导出为ONNX格式使用TensorRT的ONNX解析器加载模型针对Jetson Nano进行优化并生成TensorRT引擎from ultralytics import YOLO # 加载训练好的模型 model YOLO(best.pt) # 导出为ONNX格式 model.export(formatonnx, imgsz[640, 640], dynamicTrue)注意导出ONNX时建议启用动态维度(dynamicTrue)以便后续处理不同尺寸的输入图像。1.2 TensorRT优化技巧在Jetson Nano上运行TensorRT模型时有几个关键优化点需要考虑精度选择FP16模式通常能在保持较好精度的同时显著提升速度批处理大小根据内存限制选择合适的批处理大小输入尺寸固定输入尺寸可以避免动态形状带来的性能损失# 使用trtexec工具转换ONNX到TensorRT trtexec --onnxbest.onnx --saveEnginebest.engine --fp16 --workspace20482. Jetson Nano环境配置2.1 系统准备与依赖安装Jetson Nano运行的是基于ARM架构的Ubuntu系统需要特别注意软件包的兼容性。以下是关键依赖的安装步骤# 更新系统 sudo apt-get update sudo apt-get upgrade -y # 安装基础依赖 sudo apt-get install -y python3-pip cmake libpython3-dev # 安装PyTorch for Jetson wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl pip3 install torch-1.8.0-cp36-cp36m-linux_aarch64.whl # 安装其他Python依赖 pip3 install numpy opencv-python ultralytics2.2 TensorRT环境验证安装完成后可以通过以下代码验证TensorRT是否正常工作import tensorrt as trt TRT_LOGGER trt.Logger(trt.Logger.WARNING) runtime trt.Runtime(TRT_LOGGER) with open(best.engine, rb) as f: engine runtime.deserialize_cuda_engine(f.read()) print(f成功加载TensorRT引擎输入形状: {engine.get_binding_shape(0)})3. 推理脚本开发3.1 基础推理流程一个完整的推理脚本需要处理以下环节图像预处理尺寸调整、归一化等模型推理后处理非极大值抑制、置信度过滤等结果可视化import cv2 import numpy as np import pycuda.autoinit import pycuda.driver as cuda import tensorrt as trt class HostDeviceMem: def __init__(self, host_mem, device_mem): self.host host_mem self.device device_mem def load_engine(engine_path): TRT_LOGGER trt.Logger(trt.Logger.WARNING) with open(engine_path, rb) as f, trt.Runtime(TRT_LOGGER) as runtime: return runtime.deserialize_cuda_engine(f.read()) 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, stream3.2 实时视频处理对于工地监控场景实时视频处理能力至关重要。以下代码展示了如何使用OpenCV捕获视频流并进行实时分析def detect_video(engine_path, video_source0, conf_threshold0.5): # 初始化TensorRT引擎 engine load_engine(engine_path) context engine.create_execution_context() inputs, outputs, bindings, stream allocate_buffers(engine) # 打开视频源 cap cv2.VideoCapture(video_source) if not cap.isOpened(): print(无法打开视频源) return while True: ret, frame cap.read() if not ret: break # 预处理 input_image preprocess(frame) # 执行推理 np.copyto(inputs[0].host, input_image.ravel()) [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] context.execute_async_v2(bindingsbindings, stream_handlestream.handle) [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] stream.synchronize() # 后处理 detections postprocess(outputs, frame.shape) # 可视化 visualize(frame, detections, conf_threshold) cv2.imshow(Safety Helmet Detection, frame) if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows()4. 性能优化与调试4.1 Jetson Nano性能调优Jetson Nano的算力有限需要通过多种手段优化性能电源模式设置启用最大性能模式sudo nvpmodel -m 0 sudo jetson_clocks内存管理使用jetson_stats工具监控资源使用情况线程优化合理设置OpenCV和TensorRT的线程数4.2 推理速度与精度平衡在实际部署中需要在速度和精度之间找到平衡点。以下是几个关键参数的影响参数速度影响精度影响推荐值输入尺寸尺寸越小越快尺寸越大精度越高640x640FP16模式显著提升速度轻微精度损失启用置信度阈值影响后处理速度过滤低质量检测0.5-0.7IOU阈值影响NMS速度控制重复检测0.4-0.64.3 常见问题排查在部署过程中可能会遇到以下问题模型加载失败检查TensorRT版本兼容性确认模型导出时使用的opset版本推理结果异常验证预处理/后处理与训练时一致检查输入数据范围0-1或0-255性能不达标使用tegrastats监控硬件利用率尝试减少并发任务# 性能测试代码示例 import time def benchmark(engine_path, num_runs100): engine load_engine(engine_path) context engine.create_execution_context() inputs, outputs, bindings, stream allocate_buffers(engine) # 准备测试图像 test_image np.random.rand(640, 640, 3).astype(np.float32) start time.time() for _ in range(num_runs): np.copyto(inputs[0].host, test_image.ravel()) [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] context.execute_async_v2(bindingsbindings, stream_handlestream.handle) [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] stream.synchronize() avg_time (time.time() - start) / num_runs print(f平均推理时间: {avg_time*1000:.2f}ms, FPS: {1/avg_time:.2f})在实际项目中我们发现Jetson Nano运行YOLOv8n模型输入尺寸640x640可以达到15-20FPS的推理速度足以满足大多数工地监控场景的需求。对于更复杂的模型或更高分辨率的输入可能需要考虑使用Jetson Xavier NX等性能更强的设备。