YoloV5安全帽检测实战：从数据集到部署的完整指南

简介面向目标检测入门与边缘部署场景这套基于YoloV5s-V5.0的工地安全帽检测工程包整合模型代码、开源安全帽数据集与TensorRT加速方案适合毕业设计、课程设计、大作业或工程实训也适合在Jetson Nano等边缘设备上快速落地。压缩包共162个文件大小约3.38MB以Python训练/推理脚本、YAML配置、CMake工程文件为主附带TensorRT自定义插件生成源码以及sh、Dockerfile等部署辅助内容方便在Windows、Linux或Jetson平台还原环境。已有542人学习下载。内容覆盖从数据集准备、模型训练到TensorRT转换与推理部署的关键环节目录结构清晰零基础或进阶学习者可逐文件理解YOLOv5s-V5.0的实现细节开发者也能复用插件编译和加速经验减少边缘端实时检测的踩坑成本。 YoloV5做安全帽检测这个话题在工地上已经不算新鲜但每次有人问我怎么做我还是会建议从YoloV5-V5.0这个版本入手。原因很简单这套组合的生态太成熟了——开源数据集有现成的训练代码稳定部署方案遍地都是哪怕你是个完全没接触过目标检测的新手照着做也能在两三天内跑出自己的检测模型。这篇文章就把我从数据集选择、环境配置、训练调参到现场部署的完整过程梳理一遍把那些踩过的坑、试错后的优化方案一并写出来给你一条可以照抄的路线。1. 项目整体设计与技术选型思路1.1 为什么锁定YoloV5-V5.0而不是追新版本很多人会问现在YOLO都出到v8、v9了为什么还要用v5.0这个老版本我的回答是做实际项目稳定性远比版本号新旧重要。YoloV5-V5.0是2021年发布的版本代码结构清晰、文档完善、社区讨论沉淀多遇到任何报错基本都能搜到现成解决方案。新版本虽然在精度和速度上有提升但部分工具链、部署库的兼容性反而不如v5.0成熟。安全帽检测这类垂直场景模型的绝对精度并不是唯一指标可复现性、可维护性同样关键。v5.0的代码逻辑简单网络结构、训练流程、数据增强方式都容易理解和改写这对后续做模型裁剪、量化和边缘端部署非常友好。我实际测试对比过在SHWD数据集上v5.0训练出的模型经过合理调参后mAP0.5能达到95%以上这个精度在工地场景下已经足够实用没必要为了几个点的提升去折腾新的框架。1.2 技术路线拆解模型、数据和场景的三角关系整个项目的技术路线可以拆成三条线模型线、数据线、场景线。模型线解决“用什么样的网络结构去提取特征、定位目标”的问题数据线解决“用什么数据让模型学会区分戴帽子和没戴帽子”的问题场景线解决“模型的输入输出如何适配真实工地环境”的问题。这三条线里数据线是最容易被忽略但又最关键的。很多人拿到开源数据集就直接训练忽略了工地现场和数据集原始场景之间的差异。安全帽检测本质上是一个二分类目标检测问题目标小、背景杂、光照变化大模型能不能在真实场景中稳定工作很大程度上取决于训练数据有没有覆盖这些变化。所以在项目设计之初我就把“开源数据集为主、自采数据为辅”作为数据策略后面的效果也证明了这个决策是对的。1.3 开源数据集SHWD的使用价值与局限SHWDSafety Helmet Wearing Dataset是目前用的最广的开源安全帽检测数据集GitHub上可以直接下载包含7581张图片其中正样本约7000张标注了9044个佩戴安全帽的人头目标还有900多张负样本画面中没有人或没有安全帽类别分为0带安全帽和1未带安全帽。这个数据集有两个优点一是标注质量整体不错框选的基本都是人头区域而非全身适合做安全帽佩戴检测二是数据场景涵盖工地、道路、工厂等多种环境泛化能力有一定保障。但它的局限也很明显。数据集里的图片主要是平视或近景视角到了真实工地摄像头经常是俯拍的人小且密集如果直接用原始数据集训练模型在俯拍场景下漏检会很严重。另外数据集中夜晚、阴天、逆光等恶劣光照条件下的样本偏少这类场景也是后面容易出问题的地方。所以我的建议很明确开源数据集作为基座没问题但一定得在自己目标场景下做数据补充和测试验证。2. 环境配置与数据处理实操2.1 环境搭建的几个关键版本搭配YoloV5-V5.0的依赖不算复杂但版本搭配不对还是会折磨人。我最终稳定使用的组合是这样的Python 3.83.9和3.10也试过但3.8兼容性最稳PyTorch 1.10.0配套CUDA 11.3torchvision 0.11.0CUDNNOpenCV-Python 4.5以上安装时有几个注意点第一PyTorch的安装一定要去官网用对应CUDA版本的命令装不要用默认源否则装出来的CPU版本或者CUDA版本不匹配训练时根本调用不了GPU。第二YoloV5-V5.0的requirements.txt里依赖项比较多建议在虚拟环境里装避免污染系统环境。第三如果用的是NVIDIA显卡记得装完环境后运行一段测试代码确认CUDA可用不然训练跑起来才发现用的是CPU白白浪费时间。2.2 数据集的下载、整理与格式转换数据集方面SHWD仓库里提供的是VOC格式的XML标注而YoloV5训练需要的是YOLO格式的TXT标注。这里需要做一个格式转换把XML里的目标框坐标转换为归一化的中心点坐标格式。这一步可以用现成的转换脚本网上有很多现成代码核心逻辑就是从XML里提取每个目标的类别和边界框坐标再除以图片的宽和高。转换的时候有几个容易出错的地方一是类别标签对应关系别搞混SHWD里0对应带安全帽1对应未带安全帽如果你后续要加自己的类别标签编号一定要写在最前面并做好记录二是有些XML里可能存在坐标为负值或越界的目标框转换时要做好检查把异常框过滤掉否则训练时会报错或导致loss异常三是图片和标注文件的名字必须一一对应yolov5加载数据时靠的是同名机制。数据整理成这样的目录结构datasets/ safety_helmet/ images/ train/ val/ labels/ train/ val/划分比例我用的是9:1也就是把约90%的图片用于训练、10%用于验证这个比例在数据量几千张时比较合适能保证训练数据充足验证集又有足够的代表性。2.3 数据配置文件的编写要点YoloV5训练前需要一个数据配置文件指定训练集路径、验证集路径、类别数量和类别名称。我这里给出一个参考# safety_helmet.yaml train: ./datasets/safety_helmet/images/train val: ./datasets/safety_helmet/images/val nc: 2 names: [helmet, no-helmet]配置文件的路径建议使用相对路径并确保你在yolov5项目根目录下执行训练命令。很多初学者喜欢写绝对路径但在本机跑没问题换一台机器就会失效。另外names列表的顺序要和标注文件里的类别索引严格对应否则训练出来的模型类别会错乱这种错误排查起来很费劲。3. 核心训练实操与调参经验3.1 训练命令与关键参数选择环境配置好、数据整理好之后就可以开始训练了。YoloV5-V5.0的训练入口是train.py我常用的训练命令长这样python train.py --data safety_helmet.yaml --cfg yolov5s.yaml --weights yolov5s.pt --epochs 100 --batch-size 16 --img 640 --device 0逐个解释一下这些参数--data指定数据配置文件--cfg指定模型结构我用的是yolov5s兼顾速度和精度--weights指定预训练权重YoloV5官方提供了在COCO数据集上预训练好的权重使用预训练权重可以显著加快收敛速度并提升最终精度--epochs训练轮数在安全帽检测这个任务上100轮完全够用--batch-size批次大小取决于显存16在8G显存上是可以跑的--img输入图片尺寸640是精度和速度的平衡点。这里特别说一下预训练权重。不要从零开始训练整个网络那会耗费大量时间而且效果往往不如用预训练权重微调。YoloV5-V5.0的预训练权重可以在官方GitHub Release页面下载把下载好的yolov5s.pt放在yolov5项目根目录下即可。3.2 训练过程中的观测指标解读训练开始后终端会实时输出每个epoch的loss、精度、召回率等指标。很多人看不懂这些数字的起伏就干等着训练结束。实际上通过观察这些指标可以提前判断训练是否正常。第一个看box_loss和obj_loss这两个值应该随着训练稳步下降如果出现剧烈震荡或者不降反升说明学习率过大或数据有异常。第二个看precision和recall随着训练推进这两个值会逐渐上升然后趋于平稳。第三个看mAP0.5这是最直观的模型效果指标如果100轮跑完mAP还在90%以下就需要回查数据质量或模型结构是否合理。训练完成后结果会保存在runs/train/exp目录下里面有训练曲线图、混淆矩阵、检测效果示例图等。我一般先看混淆矩阵它能直观反映模型有没有“把戴帽子的识别成没戴帽子”这类错误——在安全帽检测场景中这个错误的严重性比误报更高因为漏掉一个未戴帽子的人等于埋下安全隐患。3.3 超参数调整的实践经验v5.0版本的超参数配置文件在data/hyp.scratch.yaml里面定义了学习率、动量、权重衰减、数据增强等参数。对于安全帽检测这类任务我调整最频繁的是这几个初始学习率lr0默认是0.01如果数据集较小或训练不稳定可以降到0.005马赛克增强mosaic默认是1.0表示100%概率使用mosaic增强这个增强方式对小目标检测很有帮助但如果发现训练后期过拟合可以适当降低到0.5颜色变换幅度hsv_h、hsv_s、hsv_v工地场景光照变化大适当增大颜色增强幅度可以让模型更鲁棒。训练轮数也值得单独说。虽然官方推荐300轮但安全帽检测数据集中等规模100~150轮已经能收敛。我用150轮训练时大约在第100轮左右mAP就趋于稳定后续50轮带来的提升微乎其微反而增加了过拟合风险。另外v5.0自带早停机制可通过设置--patience参数让模型在精度不再提升时自动停止训练省时省力。3.4 补充数据后的二次训练策略只用SHWD数据集训练出的模型在平视场景表现不错但到了真实工地尤其是俯拍场景就差很多。我的做法是在完成第一轮baseline训练后用自己拍摄或从监控视频中截取的工地画面做第二轮训练补充100~200张俯拍和夜间场景数据用小学习率继续微调。这里有个经验第二轮微调时学习率要设置得比第一轮小比如--lr0 0.001避免新数据对预训练权重产生过大的冲击同时新补充的数据要做到标注准确因为少量但是高质量的标注对模型的纠正作用远大于大量的粗糙标注。我在实际项目中补充了大约150张现场图片后俯拍场景的漏检率从之前的15%降到了3%左右这个提升效果非常显著。4. 常见问题与排查技巧实录4.1 训练时显存不足OOM怎么办训练过程中最常见的问题就是CUDA out of memory。遇到这个报错不用慌优先级从高到低有四个解决办法调小--batch-size从16降到8甚至4这是最直接的方案降低--img输入尺寸从640降到512或416显存占用会以二次方关系下降在模型配置文件里把yolov5s.yaml中的width_multiple调小比如从0.5改到0.25但这样会降低模型精度如果以上都不行就要考虑换更大显存的显卡或者使用梯度累积技术。实际经验是8G显存可以跑yolov5s640输入batch164G显存建议用416输入batch8这样的组合训练速度和显存占用相对均衡。4.2 loss为nan或训练发散的原因与处理训练过程中如果看到loss变成nan或者扶摇直上不下降十有八九是这几个原因第一学习率过大。解决方法是调小--lr0比如从0.01降到0.001第二数据中存在损坏的图片或标注越界的框这需要在数据预处理时严格检查第三batch-size过小导致梯度估计不稳定这时可以适当增大batch-size或降低学习率。另外我还遇到过一种情况在训练集上loss正常下降但验证集loss却在某个epoch后持续上升这是典型的过拟合信号。处理方法一是增加数据增强的强度和概率二是对模型加早停机制三是在第二轮训练时增加Dropout或增大权重衰减系数。4.3 俯拍小目标漏检严重怎么优化工地监控的典型画面是摄像头装在塔吊或围挡高处俯拍地面人头在画面里只占很小面积可能只有二三十个像素。SHWD数据集里的目标框通常较大导致模型在小目标上表现不佳。解决这个问题的思路有三个方向数据层面手动收集并标注俯拍场景的数据这是最有效的办法训练层面训练时把输入尺寸从640提高到960代价是显存占用增加和推理变慢推理层面使用TTA测试时增强或多尺度推理v5.0的detect.py里带了--augment参数开启后对小目标检测有正向帮助。我在一个实际项目中通过补充150张俯拍数据开启多尺度推理小目标检测能力提升非常明显俯拍画面下未戴安全帽的召回率从72%提高到了95%左右。4.4 推理部署时速度与精度的平衡训练好后接下来是部署。YoloV5-V5.0支持导出ONNX、TensorRT、OpenVINO等多种格式最常用的是ONNX。导出命令python export.py --weights runs/train/exp/weights/best.pt --include onnx --img 640导出ONNX后可以用ONNX Runtime做CPU推理也可以用TensorRT在NVIDIA显卡上做GPU推理后者速度提升非常显著。实际测量下来在同一张GTX 1660显卡上PyTorch推理一张640x640的图片耗时约15ms转成TensorRT后降到了约5ms对实时视频流检测是质变。部署时还需要注意一个细节真实工地往往是多路视频流并发推理时的batch策略和线程管理非常关键。最简单的做法是用多线程加队列方式处理多路视频每路视频独享一个模型实例通过限制输入帧率比如每秒处理5帧来降低计算压力实测下来24路视频流在单张服务器显卡上面能稳定运行。5. 从模型到落地安全帽检测项目的完整思考模型训练好了并不意味着项目结束了。安全帽检测的真正价值在于和工地安防平台联动实现告警推送、抓拍存档、数据分析等功能。我在部署阶段最大的体会是模型只是整个系统里的一环数据流的稳定性、告警规则的合理性、多路并发的资源调度这些工程问题往往比模型精度更影响最终效果。比如摄像头角度不同、光线不同、甚至安全帽颜色不同黄色、白色、蓝色都会影响模型表现这些都必须通过现场测试来验证和修正。此外模型的持续迭代也是必须考虑的。工地场景一直在变新开区域、不同的光照季节、不同施工阶段都会导致模型效果波动。我给现场项目配了一个简单的反馈闭环每天把模型漏检和误报的截图自动保存下来每周人工挑选一批有价值的坏例补充进训练集重新微调模型。这样运行一个月后模型在特定工地的表现会越来越贴合实际场景。另一个容易被忽略的点是告警策略。模型识别到未戴安全帽之后不能马上就触发声光报警否则误报会搞得现场工人很不耐烦。比较合理的做法是同一目标连续多帧比如3~5帧都被判定为未佩戴安全帽才触发告警这样能过滤掉大量瞬时误检测。加上这个帧级确认逻辑后系统的可用性提升了一个档次。最后说一句我的真实感受YoloV5-V5.0能够在安全帽检测这个方向成为事实标准不是因为它最先进而是因为它最稳妥、最透明。论文里的新模型可以带来几个点的精度提升但在真实工地上稳定性、可排查性和快速迭代能力才是维持系统的根本。这个项目做完之后我反而对“越简单越可靠”这句话体会更深。本文还有配套的精品资源点击获取