简介在计算机视觉领域目标检测是图像理解的核心任务之一其目标是在图像或视频中定位并分类特定对象广泛应用于安全监管与智能运维。YOLO系列作为单阶段检测器的代表凭借端到端的设计和出色的实时性能成为工业界部署的首选算法之一。PyTorch作为主流的深度学习框架为模型训练与迁移学习提供了高效的工程支持。将两者结合可以构建一套完整的实时检测系统。本文以工地安全帽佩戴检测为例从数据集准备、YOLOv5模型训练、PyTorch环境配置到摄像头视频流推理和边缘设备部署详细梳理了目标检测技术落地的完整链路帮助开发者理解从算法到工程的关键环节。 先说结论这个项目本身没有多玄乎但它是把深度学习、目标检测、模型部署串起来的一条完整链路认真跑通一遍你对整个CV流程的理解会上一个台阶。我去年帮一个工地做安全巡检方案时就把头盔检测作为第一块试验田。当时前后折腾了三周踩了不少坑从数据集标注、YOLOv5训练到摄像头实时推流每一步都有一些文档里不会写的东西。这篇就按照当时的实操路径把整个流程拆开讲清楚从环境搭建到部署上线尽量让一个刚接触深度学习的人也能照着跑通。1. 需求与技术选型为什么是YOLOv5加PyTorch做头盔检测1.1 先搞清楚头盔检测到底要解决什么问题工地、工厂、物流园区这些场景里安全帽、防砸头盔的佩戴监管是一件高频但容易疲劳的事情。靠人工盯监控盯着盯着就走神了靠传统图像处理颜色分割、形状匹配来做又扛不住光照变化、角度变化、遮挡这些现实干扰。头盔检测的核心需求可以拆成三点一是要能实时处理视频流二是要在复杂场景下保持较高准确率三是能够部署到普通GPU服务器甚至边缘设备上不能太吃算力。目标检测这个领域主流方案有两条路线两阶段检测器Faster R-CNN、Mask R-CNN这类精度高但速度慢适合离线分析单阶段检测器YOLO系列、SSD速度快、端到端适合实时场景。头盔检测属于施工安全类应用实时性要求排在第一位所以单阶段检测器是更合理的选择。YOLOv5在这个方向上是目前工程落地最成熟的选项训练流程完善、文档充足、社区体量大遇到问题几乎都能搜到现成解法。1.2 YOLOv5的版本与模型大小的选择YOLOv5官方仓库里提供了n/s/m/l/x五个尺寸的模型对应的模型深度和宽度逐级递增参数量也从几百万到几千万不等。做头盔检测这类目标不算太小、场景相对固定的任务我个人建议从YOLOv5s入手——它比n型精度高不少又没有m型那么吃显存训练速度和推理速度都比较均衡。选型上还要考虑一个问题头部目标在监控画面里通常只有几十到一百多像素的大小属于中小尺寸目标。YOLOv5的检测头分三层分别负责大中小目标的预测对头盔这种尺度范围的目标覆盖是比较友好的。如果以后想提升小目标检测能力可以在数据增强阶段增加小目标裁剪与复制也可以把输入图像分辨率从640提到1280但代价是训练和推理速度都会明显下降。如果你是在嵌入式设备上做部署比如Jetson Nano、RK3588这类建议直接用YOLOv5n或者YOLOv5s的FP16量化版在精度损失可接受的范围内推理速度能快一倍以上。2. 搭建PyTorch与YOLOv5环境最容易翻车的一步2.1 深度学习环境的总体规划头盔检测项目要跑起来底层需要三样东西Python环境、PyTorch深度学习框架、YOLOv5代码仓库。我强烈建议用Anaconda来管理Python环境不要直接往系统Python里装包因为YOLOv5对依赖库的版本有要求不同项目之间很容易出现版本冲突虚拟环境可以隔离掉这些麻烦。整个环境的层级关系是操作系统我用的Ubuntu 20.04→ NVIDIA驱动 → CUDA工具包 → cuDNN → PyTorch → YOLOv5依赖。如果只是做推理Post训练不一定非要自己编译CUDA但训练阶段没有GPU加速会慢到你怀疑人生。2.2 PyTorch与CUDA版本匹配的问题PyTorch的安装有一个很多人反复踩坑的点CUDA驱动版本和PyTorch默认编译的CUDA版本是两回事。你不需要自己去官网装一个完整的CUDA Toolkit如果你已经装了也没关系版本别太老就行PyTorch的pip安装包会把运行所需的CUDA运行库一并带过来只要NVIDIA驱动的版本能支撑对应的CUDA运行库版本就行。具体操作上给出一个稳定的组合2024年中这个组合非常稳# 创建虚拟环境 conda create -n helmet python3.10 -y conda activate helmet # 安装PyTorch 2.1.2对应CUDA 12.1 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121装完之后务必先验证一下GPU是否可用import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果输出是True和你的显卡型号环境这关就过了。如果打印False大概率是驱动装得不对或者版本太老用nvidia-smi查一下驱动版本低于470的话建议先升级驱动再回来装PyTorch。2.3 下载并安装YOLOv5YOLOv5的官方仓库在GitHub上可以直接用git拉取。git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt这里有个建议安装requirements.txt的时候用官方版本就行不用手动去升级或降级。我第一次跑项目时自作聪明把numpy、opencv都升到了最新版结果yolov5报了一堆兼容性错误后来重置环境重装才解决。YOLOv5对依赖的版本范围卡得很死官方给定的版本组合是最稳的。3. 头盔检测数据集准备、标注与增强3.1 数据从哪来做头盔检测数据集的获取有三条路一是用公开数据集比如SHWDSafety Helmet Wearing Dataset安全头盔佩戴数据集里面有7000多张已经标注好的图片类别分为helmet和head两类二是自己拍视频抽帧标注适合项目现场针对性强的情况三是两者混合先用公开数据集做预训练再采集现场数据做微调。我那次做的工地巡检项目,开始只用SHWD数据集训练模型在公共区域测着还行但一旦到了现场摄像头从高处往下俯拍的视角检测率掉得很厉害。后来加了两千多张现场摄像头抓拍的图片重新标注效果才恢复到可用的水平。这里想强调一个容易被忽略的点目标检测模型的训练集分布要贴合部署现场的视角和光照否则泛化能力很难保证。3.2 标注格式与目录结构YOLOv5使用的标注格式是每张图片对应一个同名的txt文本文件每一行表示一个目标内容包括类别ID和归一化后的目标框坐标。标注格式示例# classes: 0helmet 1head 0 0.512345 0.483210 0.213456 0.283456 1 0.621234 0.635432 0.145678 0.196543坐标的含义分别是类别ID、x_center、y_center、width、height都除以了图片宽高做了归一化。我习惯用LabelImg做标注选YOLO格式导出标注完自动生成同名txt文件。数据集的目录结构按YOLOv5要求组织:helmet_dataset/ ├── images/ │ ├── train/ │ └── val/ └── labels/ ├── train/ └── val/图片和标签文件必须同名train和val的目录对应存放。另外注意路径中不要出现中文否则yolov5读取数据时可能报编码错误。3.3 数据增强策略与经验YOLOv5自带了一套数据增强管线包括Mosaic把四张图拼成一张、随机仿射变换、HSV色彩空间扰动、水平翻转等。Mosaic增强在很大程度上提升了小目标的检测能力因为拼接后每个目标在最终图像中的尺寸更小相当于均衡了训练样本中的目标尺度分布。我实际操作中的体会是对于头盔检测场景Augment中的HSV扰动参数值得适当加大。施工现场的摄像头会随着时间变化出现各种色温变化从清晨的偏蓝、正午的偏白到傍晚的偏黄模型如果对色彩太敏感部署后就容易出现误检。把HSV-H、HSV-S、HSV-V三个参数各加一点模型对不同光照的鲁棒性会好很多。4. 模型训练配置、参数调优与训练监控4.1 准备数据配置文件YOLOv5训练前需要一个yaml文件来声明数据集路径、类别数和类别名。这个文件非常关键路径写错会让训练阶段直接报错。# helmet.yaml train: /home/yourname/helmet_dataset/images/train val: /home/yourname/helmet_dataset/images/val nc: 2 names: [helmet, head]这里的train和val路径建议都写成绝对路径。我之前用相对路径跑训练jupyter里能跑通换到终端里就报找不到图片排查了很久才发现是当前工作目录不一样导致的。4.2 训练命令与超参数选择YOLOv5的训练入口是train.py我用的训练命令如下python train.py \ --img 640 \ --batch 16 \ --epochs 100 \ --data helmet.yaml \ --weights yolov5s.pt \ --cache \ --device 0几个关键参数简单说一下--img 640输入图像大小默认640。头盔目标比较小可以考虑用1280但显存占用会增加许多batch大小不变的情况下显存变为约4倍先按640起步比较稳妥。--batch 16batch size越大训练越稳定但显存吃紧的话可以调低到8或4。--epochs 100首次训练建议100轮起步我用100轮训出来的模型基本已经收敛。--weights yolov5s.pt加载在COCO数据集上预训练好的权重做迁移学习。很多人不理解为什么要这个预训练权重简单说模型一开始就知道边缘、纹理、形状这些通用特征我们只需要在头盔/人头这类特定目标上做微调收敛速度和新数据量要求都会大幅下降。--cache把图片读入内存减少磁盘IO瓶颈能大幅提升训练速度。前提是你的内存够大不够的话建议不加。--device 0使用第一张GPU。4.3 训练阶段的监控与判断训练开始后控制台会逐轮输出loss、精度、召回率等指标。建议重点关注两个文件runs/train/exp*/results.png和runs/train/exp*/train_batch*.jpg。前者是训练曲线后者是每个epoch的增强后训练样本示例。训练曲线怎么判断好坏关键看三个指标的趋势box_loss边框回归损失、obj_loss目标置信度损失、cls_loss分类损失它们都应该整体向下走。如果曲线在中后期进入平台期不再下降是正常现象如果loss出现剧烈震荡甚至上升就要注意学习率是不是太大了或者batch size太小导致噪声过大。头盔检测属于两类分类问题模型相对简单我训练时到第80轮左右就已经收敛得比较好了。实际上做这个项目不需要把100轮全跑完每训练10轮左右在自己的验证图片上测试一次看效果差不多了就可以提前停止。4.4 训练时的几个坑训练过程中我遇到过几个比较典型的问题写出来帮大家省点时间第一个是显存不足CUDA out of memory。在windows本机跑一块6G显存的显卡batch16 img640基本是极限了。如果报OOM优先调小batch到8或者4不要一上来就换小模型。如果还想保留batch大小可以在train.py里找--workers参数调低一点减少同时加载图片的数量。第二个是训练到一半loss突然变NaN。这个通常是learning rate设置了太高或者训练后期梯度过大导致的。YOLOv5默认的lr会配合warmup机制自动调整但如果用的是自定义超参数文件最好先检查lr0、lrf这两个值是否合理。第三个是类别不均衡。如果数据集中helmet戴头盔的样本明显多于head未戴头盔的样本训练出来的模型会对head类别漏检严重。解决方法是数据层面做平衡把head类的图片多收集一些或者对head类图片做更多的增强来扩充样本。5. 实时检测与部署从模型到视频流5.1 用训练好的模型做视频与摄像头检测训练完成之后runs/train/exp*/weights/下会生成两个文件best.pt和last.pt。best.pt是验证集上精度最高的模型last.pt是最后一轮的模型。部署时用best.pt。YOLOv5自带的detect.py可以直接做推理支持图片、视频、摄像头以及RTSP流# 对本地视频做检测 python detect.py --weights runs/train/exp/weights/best.pt --source test.mp4 --conf 0.45 # 对摄像头做实时检测0代表第一个摄像头 python detect.py --weights runs/train/exp/weights/best.pt --source 0 --conf 0.45 # 对RTSP流做实时检测 python detect.py --weights runs/train/exp/weights/best.pt --source rtsp://user:password192.168.1.100:554/stream1 --conf 0.45--conf 0.45是置信度阈值。这个值的选择有讲究阈值设太高容易漏检特别是被部分遮挡的头盔设太低误检变多比如把背景里的圆形物体误判为头盔。我当时用0.45作为默认值在工地的场景下误检和漏检基本平衡。5.2 实时检测的代码实现与优化如果要集成到自己的业务系统里直接调用detect.py的方式不太灵活建议写一个简单的推理类把模型加载一次然后对视频流的每一帧做检测。这个推理类的核心逻辑用PyTorch实现大致如下import cv2 import torch from pathlib import Path class HelmetDetector: def __init__(self, weights_path, device0, conf0.45, imgsz640): # 使用YOLOv5仓库提供的模型加载接口简单可靠 self.model torch.hub.load(yolov5, custom, pathweights_path, sourcelocal) self.model.conf conf self.model.iou 0.45 self.model.classes [0, 1] # 只检测helmet和head self.device device def detect(self, frame): # frame是BGR格式的numpy数组 results self.model(frame) # results.pandas().xyxy[0] 返回边框坐标/置信度/类别 boxes results.pandas().xyxy[0] return boxes实测下来用torch.hub.load加载本地YOLOv5仓库在GPU上推理一帧640x640大概耗时10到15毫秒也就是每秒能处理70帧左右完全满足实时视频流的处理要求。如果目标是部署到服务器上处理多路摄像头可以考虑用多线程配合队列把视频流读取和模型推理解耦读取线程只负责抓帧推理线程专注做检测。这样可以将多路摄像头的处理能力提升至少两倍。5.3 部署到边缘设备的可选方案如果工地现场没有GPU服务器也可以在Jetson Nano、Jetson Orin这类边缘设备上部署。具体做法是先用GPU训练好模型然后做模型导出与转换把PyTorch的.pt权重导出为TensorRT的engine格式推理时用TensorRT引擎加载Jetson上能达到每秒30帧以上的推理速度。# YOLOv5支持直接导出TensorRT引擎 python export.py --weights best.pt --include engine --device 0不过这里需要提醒NVIDIA Jetson系列的JetPack版本对PyTorch和TensorRT的版本有严格限制部署前一定要确认好对应关系否则导出阶段会报一堆版本不匹配的错。6. 常见问题与排查技巧实录6.1 训练阶段高频问题速查表下面这些问题是按我实际经验里出现频率从高到低排的多数是环境或数据问题模型本身反而很少出问题。现象可能原因排查方向与解决训练时报CUDA out of memorybatch size过大、显存不足降低batch到4或8减小--img到480加--cache降低IO压力训练到中途loss变成NaN学习率过大、数据里有异常标签调低lr0检查标签是否有坐标超出图片边界比如坐标为2.1模型检测不到任何目标置信度阈值过高、测试集与训练集差异大调低--conf到0.25看看是不是类别名不匹配导致预测类被过滤训练时label文件一直报错标签txt和图片文件名对不上检查images和labels下的文件是否同名检查路径是否含中文推理时出现大量误检数据增强参数过于激进、训练不充分调低HSV扰动幅度增加训练epoch或增加负样本无头盔背景图6.2 用torch.hub.load加载模型时的坑有段时间我非常推荐用torch.hub.load加载模型因为代码简洁。但这个函数在第一次加载时会去GitHub下载YOLOv5的代码如果网络不稳定或者公司内网屏蔽了GitHub就会卡住不动。解决办法有两种一是先把YOLOv5仓库clone到本地然后用sourcelocal参数指定本地路径二是直接把yolov5目录放在项目路径下通过sys.path引用它的源码。import sys sys.path.insert(0, /path/to/yolov5) from models.experimental import attempt_load model attempt_load(runs/train/exp/weights/best.pt, map_locationcuda) model.eval()这种方式完全离线适合部署环境没有外网的场景。6.3 识别效果不稳定的调优顺序如果部署后识别效果不稳定我的建议是按以下顺序排查先看数据集测试图片是否和训练图片分布一致角度、光线、分辨率。这一步占80%的可能性。再看置信度阈值阈值调低一点从0.45降到0.3看是否漏检减少误检多了再往上调。看模型复杂度如果精度一直不够但推理时间有余量换YOLOv5m或者YOLOv5l再训一轮。最后看数据增强如果现场光线变化极大增加HSV增强强度让模型更关注形状和纹理特征。我最后实际部署时选了YOLOv5s的模型加0.45的置信度阈值在1080p的施工监控画面下头盔检测的mAP_0.5在验证集上到了91%左右单张图片推理耗时约15ms。这个水平够用而且模型显存占用不到3GB。7. 这个项目还能怎么扩展头盔检测跑通之后我发现整个技术栈可以很顺滑地扩展到其他安全防护场景。比如工地上的反光背心检测、安全绳检测、车辆区域入侵检测本质上都是目标检测只要更换数据集、重新训练、调整部署逻辑就行。数据层面有一个小技巧如果目标类别之间在视觉上比较相似比如头盔和脑袋在同一个人身上出现可以训练一个带上下文的模型比如把person也作为一个类别让检测器利用人体位置来辅助判定头顶区域是否佩戴头盔。虽然YOLOv5本身是独立目标检测但多一个person类别可以在后处理阶段做位置约束减少误检。如果想要更实时的处理能力还可以把YOLOv5的检测结果接到目标跟踪模块比如DeepSORT对同一个人在连续多帧内的佩戴状态做时序判断而不是单帧判断。这样能有效避免因为人员短暂转身、低头导致的误检在工地这种动态较强的场景下实用性会提升不少。最后再分享一个小细节训练好的模型在发布到生产环境前最好用现场历史监控视频做一个批量回放测试不要只凭验证集指标。我那次测得模型指标不错但回放一段有扬尘的施工视频时误检率明显上升——这种情况只在真实连续帧里才暴露得出来。跑完一遍回放测试你的模型才算真正可以上线。本文还有配套的精品资源点击获取
基于YOLOv5和PyTorch的工地头盔检测实时部署全流程解析
简介在计算机视觉领域目标检测是图像理解的核心任务之一其目标是在图像或视频中定位并分类特定对象广泛应用于安全监管与智能运维。YOLO系列作为单阶段检测器的代表凭借端到端的设计和出色的实时性能成为工业界部署的首选算法之一。PyTorch作为主流的深度学习框架为模型训练与迁移学习提供了高效的工程支持。将两者结合可以构建一套完整的实时检测系统。本文以工地安全帽佩戴检测为例从数据集准备、YOLOv5模型训练、PyTorch环境配置到摄像头视频流推理和边缘设备部署详细梳理了目标检测技术落地的完整链路帮助开发者理解从算法到工程的关键环节。 先说结论这个项目本身没有多玄乎但它是把深度学习、目标检测、模型部署串起来的一条完整链路认真跑通一遍你对整个CV流程的理解会上一个台阶。我去年帮一个工地做安全巡检方案时就把头盔检测作为第一块试验田。当时前后折腾了三周踩了不少坑从数据集标注、YOLOv5训练到摄像头实时推流每一步都有一些文档里不会写的东西。这篇就按照当时的实操路径把整个流程拆开讲清楚从环境搭建到部署上线尽量让一个刚接触深度学习的人也能照着跑通。1. 需求与技术选型为什么是YOLOv5加PyTorch做头盔检测1.1 先搞清楚头盔检测到底要解决什么问题工地、工厂、物流园区这些场景里安全帽、防砸头盔的佩戴监管是一件高频但容易疲劳的事情。靠人工盯监控盯着盯着就走神了靠传统图像处理颜色分割、形状匹配来做又扛不住光照变化、角度变化、遮挡这些现实干扰。头盔检测的核心需求可以拆成三点一是要能实时处理视频流二是要在复杂场景下保持较高准确率三是能够部署到普通GPU服务器甚至边缘设备上不能太吃算力。目标检测这个领域主流方案有两条路线两阶段检测器Faster R-CNN、Mask R-CNN这类精度高但速度慢适合离线分析单阶段检测器YOLO系列、SSD速度快、端到端适合实时场景。头盔检测属于施工安全类应用实时性要求排在第一位所以单阶段检测器是更合理的选择。YOLOv5在这个方向上是目前工程落地最成熟的选项训练流程完善、文档充足、社区体量大遇到问题几乎都能搜到现成解法。1.2 YOLOv5的版本与模型大小的选择YOLOv5官方仓库里提供了n/s/m/l/x五个尺寸的模型对应的模型深度和宽度逐级递增参数量也从几百万到几千万不等。做头盔检测这类目标不算太小、场景相对固定的任务我个人建议从YOLOv5s入手——它比n型精度高不少又没有m型那么吃显存训练速度和推理速度都比较均衡。选型上还要考虑一个问题头部目标在监控画面里通常只有几十到一百多像素的大小属于中小尺寸目标。YOLOv5的检测头分三层分别负责大中小目标的预测对头盔这种尺度范围的目标覆盖是比较友好的。如果以后想提升小目标检测能力可以在数据增强阶段增加小目标裁剪与复制也可以把输入图像分辨率从640提到1280但代价是训练和推理速度都会明显下降。如果你是在嵌入式设备上做部署比如Jetson Nano、RK3588这类建议直接用YOLOv5n或者YOLOv5s的FP16量化版在精度损失可接受的范围内推理速度能快一倍以上。2. 搭建PyTorch与YOLOv5环境最容易翻车的一步2.1 深度学习环境的总体规划头盔检测项目要跑起来底层需要三样东西Python环境、PyTorch深度学习框架、YOLOv5代码仓库。我强烈建议用Anaconda来管理Python环境不要直接往系统Python里装包因为YOLOv5对依赖库的版本有要求不同项目之间很容易出现版本冲突虚拟环境可以隔离掉这些麻烦。整个环境的层级关系是操作系统我用的Ubuntu 20.04→ NVIDIA驱动 → CUDA工具包 → cuDNN → PyTorch → YOLOv5依赖。如果只是做推理Post训练不一定非要自己编译CUDA但训练阶段没有GPU加速会慢到你怀疑人生。2.2 PyTorch与CUDA版本匹配的问题PyTorch的安装有一个很多人反复踩坑的点CUDA驱动版本和PyTorch默认编译的CUDA版本是两回事。你不需要自己去官网装一个完整的CUDA Toolkit如果你已经装了也没关系版本别太老就行PyTorch的pip安装包会把运行所需的CUDA运行库一并带过来只要NVIDIA驱动的版本能支撑对应的CUDA运行库版本就行。具体操作上给出一个稳定的组合2024年中这个组合非常稳# 创建虚拟环境 conda create -n helmet python3.10 -y conda activate helmet # 安装PyTorch 2.1.2对应CUDA 12.1 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121装完之后务必先验证一下GPU是否可用import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果输出是True和你的显卡型号环境这关就过了。如果打印False大概率是驱动装得不对或者版本太老用nvidia-smi查一下驱动版本低于470的话建议先升级驱动再回来装PyTorch。2.3 下载并安装YOLOv5YOLOv5的官方仓库在GitHub上可以直接用git拉取。git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt这里有个建议安装requirements.txt的时候用官方版本就行不用手动去升级或降级。我第一次跑项目时自作聪明把numpy、opencv都升到了最新版结果yolov5报了一堆兼容性错误后来重置环境重装才解决。YOLOv5对依赖的版本范围卡得很死官方给定的版本组合是最稳的。3. 头盔检测数据集准备、标注与增强3.1 数据从哪来做头盔检测数据集的获取有三条路一是用公开数据集比如SHWDSafety Helmet Wearing Dataset安全头盔佩戴数据集里面有7000多张已经标注好的图片类别分为helmet和head两类二是自己拍视频抽帧标注适合项目现场针对性强的情况三是两者混合先用公开数据集做预训练再采集现场数据做微调。我那次做的工地巡检项目,开始只用SHWD数据集训练模型在公共区域测着还行但一旦到了现场摄像头从高处往下俯拍的视角检测率掉得很厉害。后来加了两千多张现场摄像头抓拍的图片重新标注效果才恢复到可用的水平。这里想强调一个容易被忽略的点目标检测模型的训练集分布要贴合部署现场的视角和光照否则泛化能力很难保证。3.2 标注格式与目录结构YOLOv5使用的标注格式是每张图片对应一个同名的txt文本文件每一行表示一个目标内容包括类别ID和归一化后的目标框坐标。标注格式示例# classes: 0helmet 1head 0 0.512345 0.483210 0.213456 0.283456 1 0.621234 0.635432 0.145678 0.196543坐标的含义分别是类别ID、x_center、y_center、width、height都除以了图片宽高做了归一化。我习惯用LabelImg做标注选YOLO格式导出标注完自动生成同名txt文件。数据集的目录结构按YOLOv5要求组织:helmet_dataset/ ├── images/ │ ├── train/ │ └── val/ └── labels/ ├── train/ └── val/图片和标签文件必须同名train和val的目录对应存放。另外注意路径中不要出现中文否则yolov5读取数据时可能报编码错误。3.3 数据增强策略与经验YOLOv5自带了一套数据增强管线包括Mosaic把四张图拼成一张、随机仿射变换、HSV色彩空间扰动、水平翻转等。Mosaic增强在很大程度上提升了小目标的检测能力因为拼接后每个目标在最终图像中的尺寸更小相当于均衡了训练样本中的目标尺度分布。我实际操作中的体会是对于头盔检测场景Augment中的HSV扰动参数值得适当加大。施工现场的摄像头会随着时间变化出现各种色温变化从清晨的偏蓝、正午的偏白到傍晚的偏黄模型如果对色彩太敏感部署后就容易出现误检。把HSV-H、HSV-S、HSV-V三个参数各加一点模型对不同光照的鲁棒性会好很多。4. 模型训练配置、参数调优与训练监控4.1 准备数据配置文件YOLOv5训练前需要一个yaml文件来声明数据集路径、类别数和类别名。这个文件非常关键路径写错会让训练阶段直接报错。# helmet.yaml train: /home/yourname/helmet_dataset/images/train val: /home/yourname/helmet_dataset/images/val nc: 2 names: [helmet, head]这里的train和val路径建议都写成绝对路径。我之前用相对路径跑训练jupyter里能跑通换到终端里就报找不到图片排查了很久才发现是当前工作目录不一样导致的。4.2 训练命令与超参数选择YOLOv5的训练入口是train.py我用的训练命令如下python train.py \ --img 640 \ --batch 16 \ --epochs 100 \ --data helmet.yaml \ --weights yolov5s.pt \ --cache \ --device 0几个关键参数简单说一下--img 640输入图像大小默认640。头盔目标比较小可以考虑用1280但显存占用会增加许多batch大小不变的情况下显存变为约4倍先按640起步比较稳妥。--batch 16batch size越大训练越稳定但显存吃紧的话可以调低到8或4。--epochs 100首次训练建议100轮起步我用100轮训出来的模型基本已经收敛。--weights yolov5s.pt加载在COCO数据集上预训练好的权重做迁移学习。很多人不理解为什么要这个预训练权重简单说模型一开始就知道边缘、纹理、形状这些通用特征我们只需要在头盔/人头这类特定目标上做微调收敛速度和新数据量要求都会大幅下降。--cache把图片读入内存减少磁盘IO瓶颈能大幅提升训练速度。前提是你的内存够大不够的话建议不加。--device 0使用第一张GPU。4.3 训练阶段的监控与判断训练开始后控制台会逐轮输出loss、精度、召回率等指标。建议重点关注两个文件runs/train/exp*/results.png和runs/train/exp*/train_batch*.jpg。前者是训练曲线后者是每个epoch的增强后训练样本示例。训练曲线怎么判断好坏关键看三个指标的趋势box_loss边框回归损失、obj_loss目标置信度损失、cls_loss分类损失它们都应该整体向下走。如果曲线在中后期进入平台期不再下降是正常现象如果loss出现剧烈震荡甚至上升就要注意学习率是不是太大了或者batch size太小导致噪声过大。头盔检测属于两类分类问题模型相对简单我训练时到第80轮左右就已经收敛得比较好了。实际上做这个项目不需要把100轮全跑完每训练10轮左右在自己的验证图片上测试一次看效果差不多了就可以提前停止。4.4 训练时的几个坑训练过程中我遇到过几个比较典型的问题写出来帮大家省点时间第一个是显存不足CUDA out of memory。在windows本机跑一块6G显存的显卡batch16 img640基本是极限了。如果报OOM优先调小batch到8或者4不要一上来就换小模型。如果还想保留batch大小可以在train.py里找--workers参数调低一点减少同时加载图片的数量。第二个是训练到一半loss突然变NaN。这个通常是learning rate设置了太高或者训练后期梯度过大导致的。YOLOv5默认的lr会配合warmup机制自动调整但如果用的是自定义超参数文件最好先检查lr0、lrf这两个值是否合理。第三个是类别不均衡。如果数据集中helmet戴头盔的样本明显多于head未戴头盔的样本训练出来的模型会对head类别漏检严重。解决方法是数据层面做平衡把head类的图片多收集一些或者对head类图片做更多的增强来扩充样本。5. 实时检测与部署从模型到视频流5.1 用训练好的模型做视频与摄像头检测训练完成之后runs/train/exp*/weights/下会生成两个文件best.pt和last.pt。best.pt是验证集上精度最高的模型last.pt是最后一轮的模型。部署时用best.pt。YOLOv5自带的detect.py可以直接做推理支持图片、视频、摄像头以及RTSP流# 对本地视频做检测 python detect.py --weights runs/train/exp/weights/best.pt --source test.mp4 --conf 0.45 # 对摄像头做实时检测0代表第一个摄像头 python detect.py --weights runs/train/exp/weights/best.pt --source 0 --conf 0.45 # 对RTSP流做实时检测 python detect.py --weights runs/train/exp/weights/best.pt --source rtsp://user:password192.168.1.100:554/stream1 --conf 0.45--conf 0.45是置信度阈值。这个值的选择有讲究阈值设太高容易漏检特别是被部分遮挡的头盔设太低误检变多比如把背景里的圆形物体误判为头盔。我当时用0.45作为默认值在工地的场景下误检和漏检基本平衡。5.2 实时检测的代码实现与优化如果要集成到自己的业务系统里直接调用detect.py的方式不太灵活建议写一个简单的推理类把模型加载一次然后对视频流的每一帧做检测。这个推理类的核心逻辑用PyTorch实现大致如下import cv2 import torch from pathlib import Path class HelmetDetector: def __init__(self, weights_path, device0, conf0.45, imgsz640): # 使用YOLOv5仓库提供的模型加载接口简单可靠 self.model torch.hub.load(yolov5, custom, pathweights_path, sourcelocal) self.model.conf conf self.model.iou 0.45 self.model.classes [0, 1] # 只检测helmet和head self.device device def detect(self, frame): # frame是BGR格式的numpy数组 results self.model(frame) # results.pandas().xyxy[0] 返回边框坐标/置信度/类别 boxes results.pandas().xyxy[0] return boxes实测下来用torch.hub.load加载本地YOLOv5仓库在GPU上推理一帧640x640大概耗时10到15毫秒也就是每秒能处理70帧左右完全满足实时视频流的处理要求。如果目标是部署到服务器上处理多路摄像头可以考虑用多线程配合队列把视频流读取和模型推理解耦读取线程只负责抓帧推理线程专注做检测。这样可以将多路摄像头的处理能力提升至少两倍。5.3 部署到边缘设备的可选方案如果工地现场没有GPU服务器也可以在Jetson Nano、Jetson Orin这类边缘设备上部署。具体做法是先用GPU训练好模型然后做模型导出与转换把PyTorch的.pt权重导出为TensorRT的engine格式推理时用TensorRT引擎加载Jetson上能达到每秒30帧以上的推理速度。# YOLOv5支持直接导出TensorRT引擎 python export.py --weights best.pt --include engine --device 0不过这里需要提醒NVIDIA Jetson系列的JetPack版本对PyTorch和TensorRT的版本有严格限制部署前一定要确认好对应关系否则导出阶段会报一堆版本不匹配的错。6. 常见问题与排查技巧实录6.1 训练阶段高频问题速查表下面这些问题是按我实际经验里出现频率从高到低排的多数是环境或数据问题模型本身反而很少出问题。现象可能原因排查方向与解决训练时报CUDA out of memorybatch size过大、显存不足降低batch到4或8减小--img到480加--cache降低IO压力训练到中途loss变成NaN学习率过大、数据里有异常标签调低lr0检查标签是否有坐标超出图片边界比如坐标为2.1模型检测不到任何目标置信度阈值过高、测试集与训练集差异大调低--conf到0.25看看是不是类别名不匹配导致预测类被过滤训练时label文件一直报错标签txt和图片文件名对不上检查images和labels下的文件是否同名检查路径是否含中文推理时出现大量误检数据增强参数过于激进、训练不充分调低HSV扰动幅度增加训练epoch或增加负样本无头盔背景图6.2 用torch.hub.load加载模型时的坑有段时间我非常推荐用torch.hub.load加载模型因为代码简洁。但这个函数在第一次加载时会去GitHub下载YOLOv5的代码如果网络不稳定或者公司内网屏蔽了GitHub就会卡住不动。解决办法有两种一是先把YOLOv5仓库clone到本地然后用sourcelocal参数指定本地路径二是直接把yolov5目录放在项目路径下通过sys.path引用它的源码。import sys sys.path.insert(0, /path/to/yolov5) from models.experimental import attempt_load model attempt_load(runs/train/exp/weights/best.pt, map_locationcuda) model.eval()这种方式完全离线适合部署环境没有外网的场景。6.3 识别效果不稳定的调优顺序如果部署后识别效果不稳定我的建议是按以下顺序排查先看数据集测试图片是否和训练图片分布一致角度、光线、分辨率。这一步占80%的可能性。再看置信度阈值阈值调低一点从0.45降到0.3看是否漏检减少误检多了再往上调。看模型复杂度如果精度一直不够但推理时间有余量换YOLOv5m或者YOLOv5l再训一轮。最后看数据增强如果现场光线变化极大增加HSV增强强度让模型更关注形状和纹理特征。我最后实际部署时选了YOLOv5s的模型加0.45的置信度阈值在1080p的施工监控画面下头盔检测的mAP_0.5在验证集上到了91%左右单张图片推理耗时约15ms。这个水平够用而且模型显存占用不到3GB。7. 这个项目还能怎么扩展头盔检测跑通之后我发现整个技术栈可以很顺滑地扩展到其他安全防护场景。比如工地上的反光背心检测、安全绳检测、车辆区域入侵检测本质上都是目标检测只要更换数据集、重新训练、调整部署逻辑就行。数据层面有一个小技巧如果目标类别之间在视觉上比较相似比如头盔和脑袋在同一个人身上出现可以训练一个带上下文的模型比如把person也作为一个类别让检测器利用人体位置来辅助判定头顶区域是否佩戴头盔。虽然YOLOv5本身是独立目标检测但多一个person类别可以在后处理阶段做位置约束减少误检。如果想要更实时的处理能力还可以把YOLOv5的检测结果接到目标跟踪模块比如DeepSORT对同一个人在连续多帧内的佩戴状态做时序判断而不是单帧判断。这样能有效避免因为人员短暂转身、低头导致的误检在工地这种动态较强的场景下实用性会提升不少。最后再分享一个小细节训练好的模型在发布到生产环境前最好用现场历史监控视频做一个批量回放测试不要只凭验证集指标。我那次测得模型指标不错但回放一段有扬尘的施工视频时误检率明显上升——这种情况只在真实连续帧里才暴露得出来。跑完一遍回放测试你的模型才算真正可以上线。本文还有配套的精品资源点击获取