YOLOv5 中使用 PyTorch 实现自定义模型的实时检测
在智能制造、智能安防和边缘计算快速发展的今天,如何让训练好的目标检测模型走出实验室,真正“看见”现实世界,是每一个视觉工程师必须面对的问题。尤其是在产线质检、行为识别或无人值守监控等场景中,低延迟、高鲁棒性的实时检测能力,往往直接决定了系统的可用性。
而 YOLOv5 凭借其简洁的架构设计、出色的推理速度与精度平衡,以及强大的社区支持,已成为工业部署中最受欢迎的目标检测框架之一。更关键的是,它通过torch.hub.load提供了一种极为轻量级的模型集成方式——无需复杂的封装,几行代码即可将.pt权重加载为可调用的推理对象。
本文将带你从零开始,构建一个基于摄像头输入的 YOLOv5 实时检测系统。我们将聚焦于本地自定义模型的加载逻辑、OpenCV 视频流处理、结果解析与可视化渲染,并深入剖析常见问题背后的原理,帮助你绕过那些看似简单却令人头疼的“坑”。
假设你已经完成了 YOLOv5 自定义数据集的训练,并得到了位于runs/train/exp/weights/best.pt的最优权重文件。现在的问题是:如何不依赖detect.py脚本,在自己的项目中独立调用这个模型?
答案就是 PyTorch Hub 的load方法。但这里有个陷阱:即使你是从本地加载模型,也必须保留官方仓库的路径标识。也就是说,下面这行代码中的'ultralytics/yolov5'并不是让你去联网下载,而是告诉 PyTorch Hub:“我要加载的是符合 YOLOv5 项目结构的模型”。
model = torch.hub.load( 'ultralytics/yolov5', # 必须写这个字符串,否则无法识别 'custom', path='/path/to/your/yolov5/runs/train/exp/weights/best.pt', source='local' # 明确指定本地加载 )但这还不够。当你运行这段代码时,大概率会遇到这样的报错:
ModuleNotFoundError: No module named 'models.common'为什么?因为torch.hub.load在初始化模型时,会尝试导入 YOLOv5 项目内部的模块(如models,utils.datasets等)。如果你只是把best.pt单独拿出来,而没有包含整个项目结构,这些依赖就找不到。
解决办法也很直接:将你的 YOLOv5 项目根目录加入 Python 的系统路径。
import sys sys.path.insert(0, '/path/to/your/yolov5') # 替换为实际路径这一行代码看似不起眼,却是能否成功加载的关键。它确保了所有内部模块都能被正确导入。建议使用绝对路径,避免相对路径带来的歧义,特别是在不同工作目录下运行脚本时。
接下来是视频流的采集。OpenCV 是最常用的工具,但要注意一点:OpenCV 默认读取的是 BGR 格式的图像,而 YOLOv5 模型期望的是 RGB 输入。如果不做转换,虽然模型仍能输出结果,但颜色通道错乱可能导致某些预处理逻辑异常(尤其是在后续接入其他视觉任务时)。
因此,在送入模型前,务必进行色彩空间转换:
ret, frame = cap.read() img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) results = model(img_rgb) # 推理推理完成后,results对象包含了丰富的信息。YOLOv5 官方提供了一个非常实用的方法.pandas(),它可以将原始张量输出转换为结构化的 Pandas DataFrame,极大提升了后处理的灵活性。
detections = results.pandas().xyxy[0]这里的xyxy表示边界框以(xmin, ymin, xmax, ymax)的格式返回,而[0]表示第一帧的结果(因为你可以一次性传入多张图片进行批量推理)。返回的 DataFrame 包含以下字段:
| 字段名 | 含义 |
|---|---|
xmin,ymin | 检测框左上角坐标(像素) |
xmax,ymax | 右下角坐标 |
confidence | 置信度分数(0~1) |
name | 预测类别标签(来自训练时的names.yaml) |
举个例子,如果你训练的是一个缺陷检测模型,那么name可能是"scratch"或"crack";如果是安防场景,则可能是"person"、"phone"等。
这个 DataFrame 的强大之处在于,你可以像操作数据库一样对其进行筛选、排序和统计。比如只保留置信度高于 0.5 的检测结果:
high_conf = detections[detections.confidence > 0.5]或者统计当前画面中各类别的数量:
counts = detections['name'].value_counts() print(counts.to_dict()) # {'person': 2, 'bottle': 1}当然,最终我们还是要回到画面上来——把检测结果可视化地展示出来。此时需要将标注绘制回原本的 BGR 图像上,以便用 OpenCV 显示。
annotated_frame = frame.copy() for _, row in detections.iterrows(): x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax']) label = f"{row['name']} {row['confidence']:.2f}" color = (0, 255, 0) # 绿色边框 cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), color, 2) cv2.putText(annotated_frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)注意坐标的类型转换:模型输出的是浮点数,但cv2.rectangle要求整数坐标,所以要用int()强制转换。
如果你想进一步利用这些检测结果来做业务逻辑判断,比如“发现手机就报警”,可以这样写:
if 'phone' in detections['name'].values: phone_count = len(detections[detections['name'] == 'phone']) max_conf = detections[detections['name'] == 'phone']['confidence'].max() if max_conf > 0.7: print(f"⚠️ 检测到 {phone_count} 部手机!最高置信度:{max_conf:.3f}") # 这里可以触发警报、拍照上传、发送通知等这种模式非常适合嵌入到更大的系统中,比如 Flask Web 服务或 ROS 节点,作为前端感知模块输出结构化事件。
关于设备选择,YOLOv5 支持 CPU 和 GPU 推理。默认情况下会自动使用 CUDA(如果可用),但你也可以手动指定:
model = torch.hub.load(..., device='cpu') # 强制使用 CPU model = torch.hub.load(..., device='cuda:0') # 指定 GPU 设备性能方面差异显著:在 RTX 3060 上,YOLOv5s 可达 80+ FPS;而在树莓派 4B 的 CPU 上,可能只有 5~8 FPS。对于低功耗设备,建议使用yolov5n或yolov5s这类轻量模型,并适当降低输入分辨率(修改model.img_size)。
还有一点容易被忽略:当画面中没有任何检测结果时,detections是一个空的 DataFrame。如果你直接遍历它:
for _, row in detections.iterrows(): # 如果为空,不会进入循环好消息是,Pandas 的iterrows()在空 DataFrame 上不会抛错,而是安静地跳过循环。但从工程健壮性角度出发,最好还是加个判断:
if len(detections) == 0: print("📭 当前帧未检测到任何目标") else: for _, row in detections.iterrows(): ...这样便于调试和日志记录。
至于部署环境,这套方案已在多种平台上验证通过:NVIDIA Jetson Nano、Intel NUC、树莓派 4B(配合 Ubuntu Core)、甚至 Windows 工控机。只要安装好 PyTorch、OpenCV 和 Pandas,基本无需额外配置。
不过要提醒一句:PyTorch 的启动开销较大,首次加载模型可能会有 2~5 秒的延迟(尤其是大模型)。这不是 bug,而是模型加载和 JIT 编译的过程。生产环境中建议在服务启动时就完成模型加载,而不是每次请求都重新加载。
最后,虽然我们用了torch.hub.load,但它本质上只是一个便利接口。它的背后仍然是标准的torch.load和模型重建流程。这也意味着你可以进一步定制,比如:
修改非极大抑制(NMS)阈值:
python model.iou = 0.3 model.conf = 0.5更改输入尺寸:
python model.img_size = [640] # 影响推理速度与小物体检测能力添加自定义后处理钩子,实现跟踪、计数或多阶段推理。
正是这种“简单而不简陋”的设计哲学,让 YOLOv5 成为了连接研究与落地的理想桥梁。
技术总是在演进,YOLO 系列也已发展到 v8、v10,甚至出现了 YOLO-NAS 等新架构。但无论形式如何变化,“一次前向传播完成检测”的核心思想始终未变。它所追求的,不只是更高的 mAP,更是更低的延迟、更强的泛化能力和更简单的部署路径。
而我们作为开发者,真正要掌握的,从来都不是某一行代码,而是如何在真实世界的约束条件下,把一个训练好的模型变成看得见、用得上的智能系统。
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考