ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

YOLOv7安全帽检测实战:从数据集标注到TensorRT部署全流程

YOLOv7安全帽检测实战:从数据集标注到TensorRT部署全流程 简介目标检测是计算机视觉中从图像中定位并分类多个目标的核心技术其原理基于深度学习模型对边界框与类别概率的联合预测。在工业安全场景中目标检测模型能够实时分析监控视频判断作业人员是否佩戴安全帽对电力巡检等高风险环境具有重要意义。YOLOv7作为高效的一阶段检测算法在速度与精度间取得平衡支持多尺度预测适合部署于边缘设备。本文以YOLOv7电力巡检安全帽检测项目为例系统讲解数据集构建与标注规范、模型训练参数调优、评估指标解读以及从PyTorch到ONNX再到TensorRT的完整部署链路并针对小目标漏检、误报控制等工程问题给出实用排查方案帮助工程师快速落地可靠的检测系统。 很多人拿到一个目标检测项目第一步不是急吼吼地去敲代码而是先搞清楚一个问题我手上有什么我要去哪里我该怎么去。以YOLOv7电力巡检安全帽检测这个项目为例它其实是一个比较完整的工业落地场景——现场有实时监控摄像头需要自动识别作业人员有没有戴安全帽检测结果要能实时上报并且误报率不能太高。这个项目最吸引人的地方在于它不是一套空壳代码而是把训练好的模型和标注好的数据集一起交付意味着拿到手就能跑通流程甚至直接换到自己的数据上做迁移学习。这篇文章我想从一个做过的角度把这个项目从数据到部署的真实链路拆开讲包括数据集长什么样、标注规范怎么做、YOLOv7训练时那些参数到底怎么调、模型评估指标怎么看、导出部署又有什么坑。适合的人群很明确——正在做工业安全相关视觉项目的人、需要用YOLOv7训练自己数据集的工程师以及刚接触目标检测但想完整跟一遍项目流程的初学者。我会尽量把每一步“为什么这么做”也讲清楚而不是只丢一堆命令。1. 项目整体设计与方案选型1.1 为什么安全帽检测必须用目标检测而不是图像分类做电力巡检场景的人对这个问题应该有体会摄像头拍到的画面是一个连续的视频流里面可能同时出现多个人、多个目标而且人可能处于不同的姿态和位置。如果你用图像分类本质上只能回答“这张图片里有没有人戴安全帽”但你不知道人在哪里也就无法做报警联动、区域管控这些后续操作。目标检测解决的正是“是什么”和“在哪里”两个问题。它输出的是目标的边界框bounding box和类别置信度比如“person 0.87, 坐标 (x1, y1, x2, y2)”和“helmet 0.92, 坐标 (x3, y3, x4, y4)”。有了坐标下游系统就可以判断这个人的头肩区域有没有被安全帽框覆盖再决定是否触发告警。这就是为什么这个项目选的是YOLOv7这类检测模型而不是单纯用ResNet做二分类。还有一个更实际的原因电力巡检现场的目标尺度变化很大。摄像头装在杆塔上远处的作业人员可能只占几十个像素近处的可能占几百个像素。目标检测模型天然具备多尺度预测能力YOLOv7在三个不同尺度的特征图上做预测小目标召回率在同代模型里也算得上优秀。这一点在巡检场景里非常重要。1.2 YOLOv7选型的几个关键考量YOLOv7发布之初之所以引起关注是因为它在速度和精度之间做到了当时的最佳平衡。COCO数据集上相同推理速度下精度超过了当时绝大多数YOLO系列模型而且普通GPU就能跑不需要堆算力。选型的时候我对比过几个方向模型方案优点缺点适合场景YOLOv5生态成熟资料多结构相对传统精度上限一般快速起步YOLOv7精度高推理快配置复杂度略高后处理细节多工业级检测YOLOv8新Anchor-free思路当时还没完全稳定实验探索Faster R-CNN精度上限高速度太慢实时性差离线检测电力巡检场景对实时性要求不低现场摄像头通常要跑20 FPS以上才能跟得上人员走动同时对精度要求也高漏检一个人没戴安全帽就可能出安全事故。YOLOv7刚好两边都够得着。另外一个重要原因是它在边缘设备上比较友好后面要部署到Jetson Nano或者工控机YOLOv7的TensorRT加速方案成熟度比较高。1.3 数据集“自带”到底省了多少事做过深度学习的人都知道数据往往是整个项目里最耗时、最费钱的部分。一个安全帽检测项目如果没有现成数据你需要去现场拍多个时间段的视频、做抽帧、清洗、筛选、标注最后可能还要做数据增强来平衡类别整套流程下来一两个星期都算快的。这个项目里“标注好的数据集”是我认为比模型权重更值钱的部分。因为模型权重是“鱼”数据集是“渔”。有了数据集你可以做这几件事一是直接基于这份数据训练一个适配自己场景的模型二是往这份数据里补充自己现场采集的样本做增量训练三是用这份数据验证其他模型架构的效果相当于一个基准数据集。严格来说一份可用的标注数据集至少包含三样东西原始的图片文件、与图片同名的标注文件、标注说明文件。YOLO格式的标注是txt文件每一行是“类别id x_center y_center width height”坐标值是归一化后的0到1的小数。这个项目的数据集我看了下标注格式是YOLO格式类别包含安全帽、人、头三个类别这个设计很有讲究下面细说。2. 数据集构建与标注规范2.1 为什么要标注“头”这个类别而不是只标“安全帽”和“人”这是一个很容易被忽视但特别关键的细节。很多入门者做安全帽检测下意识就标两类有人person、有安全帽helmet。结果训练出来的模型经常出现一种情况人没戴安全帽时不报警因为模型没有“头”的概念它不知道人身上哪个位置应该戴安全帽。加上“头”head这个类别之后检测逻辑就变成了先检测出头的坐标再检测出安全帽的坐标然后判断这两者之间的重合程度。如果头的区域没有安全帽覆盖就可以判定为未佩戴。这种基于空间关系的判断准确率远高于简单的人-帽共存判断因为远处的人可能很小安全帽框和人框的重叠率极高二分类模型很难学到区分逻辑。我实际测试过这个三分类方案误报率确实比二分类要低。因为在复杂场景下背景里偶尔会出现类安全帽形状的物体比如远处反光的塑料盖子如果只有“安全帽”一类模型很容易把它误检出来。但把“头”和“安全帽”关联起来判断之后那些没有对应“头”的安全帽检测结果可以直接作为无效结果丢弃。2.2 标注数据规模和分布背后的逻辑这个项目的数据集规模大概在几千张图片的量级具体数量我不再细述但分布逻辑值得展开安全帽类别样本最多头类别次之人类别相对最少。这个分布和实际场景是吻合的——不是每个出现的人都会戴安全帽但只要有人的头就一定会出现而安全帽因为形似目标多样本数量必须要大模型才能学得稳。关键在于类别不均衡的处理。如果安全帽样本远多于头样本模型会对安全帽过拟合导致头的检测精度下降。在这个项目的数据集里因为没有做太极端的数据增强训练时我建议在损失函数层面或者数据采样层面做平衡比如为三个类别设置不同的loss权重。更简单的方法是通过Mosaic增强来提升头的检出率Mosaic把四张图拼在一起相当于变相增加了小目标的数量对头这类小目标很有帮助。2.3 标注工具选型和格式转换细节拿到这份标注好的数据集如果后续要补充自己的数据推荐用LabelImg或者LabelStudio。LabelImg是老牌工具适合单人小规模标注输出VOC格式的XML文件LabelStudio更现代支持多人协作可以导出YOLO和COCO多种格式。这里需要注意一个坑LabelImg的默认输出格式是Pascal VOC XML而YOLO训练需要的是txt标注文件。XML和txt之间不是简单的后缀替换坐标值需要做转换。XML里的坐标是绝对像素坐标xmin, ymin, xmax, ymaxYOLO需要的是归一化的中心点坐标和宽高。def voc_to_yolo(xml_file, img_width, img_height): 将VOC格式XML标注转换为YOLO格式txt xml_file: XML标注文件路径 img_width: 图片宽度像素 img_height: 图片高度像素 返回YOLO格式的标注行列表 import xml.etree.ElementTree as ET tree ET.parse(xml_file) root tree.getroot() yolo_lines [] for obj in root.iter(object): cls_name obj.find(name).text # 类别需要映射成数字id根据项目预先定义的类别顺序 cls_id class_to_id[cls_name] bndbox obj.find(bndbox) xmin float(bndbox.find(xmin).text) ymin float(bndbox.find(ymin).text) xmax float(bndbox.find(xmax).text) ymax float(bndbox.find(ymax).text) # 转换为YOLO格式 x_center (xmin xmax) / 2.0 / img_width y_center (ymin ymax) / 2.0 / img_height w (xmax - xmin) / img_width h (ymax - ymin) / img_height yolo_lines.append(f{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}) return yolo_lines归一化的好处是标注与图片分辨率解耦无论图片是1080p还是720p同一份标注都能用。这也意味着在你补充自己的数据时如果图片分辨率不一致不用额外做坐标换算直接标注后导出即可。但要注意如果你用LabelStudio导出的YOLO格式后记得检查类别id的顺序因为YOLO训练时类别id是根据data配置文件里的类别顺序确定的一旦顺序错位整个模型的输出语义就全乱了。2.4 数据集的目录结构拿到手应该怎么组织拿到项目后第一件事不是解压就跑而是先把数据目录整理好。YOLOv7训练时默认会查找images和labels两个目录分别存放图片和标注txt文件train和val按比例划分。推荐的结构如下dataset/ ├── images/ │ ├── train/ │ └── val/ ├── labels/ │ ├── train/ │ └── val/ ├── classes.txt └── data.yamlimages和labels下同样的子目录结构文件名一一对应这是YOLO系列的基本约定。如果你的数据集结构不标准建议写个简单的脚本先做一次校验——把所有图片和同名txt对齐并检查有没有缺标注的图片、有没有空标注文件这些脏数据会在训练时造成意想不到的问题。实际整理时我遇到过一个小问题某些图片标注文件存在但内容为空也就是没有任何目标。这在训练时不会报错但可能影响模型对背景的判别。如果你是补充数据建议保留这些空标注图片但数量不要过多否则模型会倾向于把所有区域都判为背景。3. YOLOv7训练实战3.1 环境搭建的版本搭配YOLOv7对环境的兼容性还行但版本搭配不好会让人折腾半天。我整理了一份经过实测的组合Python 3.8 或 3.9PyTorch 1.12.0 或 1.13.0CUDA 11.3/11.6 对应版本的PyTorchOpenCV-Python 4.x依赖安装直接用requirements.txt但注意numpy版本不要超过1.24否则部分老版本依赖会编译报错git clone https://github.com/WongKinYiu/yolov7.git cd yolov7 pip install -r requirements.txtGPU环境验证用一行命令python -c import torch; print(torch.cuda.is_available())输出True就说明环境OK。如果输出False去检查PyTorch是否装了CUDA版本而不是CPU版本。有一个很容易踩的坑YOLOv7的代码是从YOLOv5分支改过来的很多配置项看起来相似但实际不完全一样不要想当然用YOLOv5的旧习惯去改比如anchor的自动计算方式、数据增强的关闭方法都跟YOLOv5有差异。3.2 配置文件修改的几个关键点YOLOv7训练前要改的地方不多主要是一个数据配置data和一个模型配置文件cfg。data配置是yaml格式核心就是三个路径加一个类别数量train: dataset/images/train val: dataset/images/val nc: 3 names: [helmet, head, person]注意train和val指向的是图片目录不是标注目录也不用写图片文件名列表YOLOv7会自动根据同名规则找labels目录下的标注。这是很多新手会懵的地方。还有一点data.yaml里的names顺序必须和标注txt里的类别id严格对应前面也强调过一旦顺序错位训练出来的模型语义全乱。模型配置只需要改nc这一行。YOLOv7默认的yolov7.yaml里nc80COCO类别数改成3即可。anchors不建议手工改YOLOv7训练时能自动通过聚类重新计算虽然它的自动anchor计算不像YOLOv5那么高频但改错anchor反而影响收敛速度。3.3 损失函数、批次大小与学习率的关系训练时batch size和img size是首要考虑的参数。不同显存下显卡显存推荐batch size推荐图片尺寸8 GB8-1664012 GB16-3264024 GB32-64640img size建议直接用640YOLOv7的预训练权重本来就是在640分辨率下训练的改小会影响精度改大要相应的训练时间。你在推理时如果显存够也可以跑到640不够再降到512或416。学习率的设置可以这样理解YOLOv7的lr0默认是0.01如果batch size翻倍或减半学习率也应做相应调整。经验法则是前几次训练时盯着loss曲线如果loss震荡特别剧烈说明学习率偏大如果loss下降像蜗牛爬说明学习率偏小。我一般用0.01作为起点batch size 16时效果稳定。训练命令示例python train.py --workers 4 --device 0 --batch-size 16 --data data/data.yaml --img 640 --cfg cfg/training/yolov7.yaml --weights yolov7_training.pt --name safetynet --epochs 100--weights参数有讲究直接给yolov7预训练权重可以加速收敛。如果你从头训练模型要自己从零学特征不仅慢精度往往还差一截。迁移学习的价值就在这里预训练权重已经学会了通用的边缘、纹理、形状特征你只需要在它基础上微调安全帽、头、人这几个类别的差异。3.4 训练过程的可视化监控YOLOv7用--project和--name指定输出目录训练过程会在runs/train/你的名字/下生成log文件。训练时我一般会开两个窗口一个跑训练一个看训练日志。比较关键的指标是box_loss、obj_loss、cls_loss三个损失都趋于平稳且数值明显下降后模型才有实用价值。另外一个重要的可视化工具是训练结束后生成的results.png里面有损失曲线、精确率、召回率、mAP曲线的趋势图。看这个图可以从宏观角度观察模型有没有过拟合如果在验证集上的mAP开始下降而训练集的loss还在下降就是在过拟合了这时候要么减小模型复杂度、要么加数据增强、要么早停。训练过程中每个epoch结束会保存一个权重YOLOv7默认保存last.pt和best.pt。best.pt是按验证集mAP最高的epoch保存的最终部署就选它。注意不要想当然用last.pt尤其训练后期模型可能已经过拟合到训练集last.pt的泛化能力不一定好。3.5 模型评估指标mAP、Precision、Recall到底该怎么看训练完成后--test命令可以让你在验证集上评估模型性能。YOLOv7输出的指标里mAP0.5和mAP0.5:0.95是两个最常见的数值。mAP0.5IoU阈值取0.5时的平均精度工业验收常看这个mAP0.5:0.95IoU阈值从0.5到0.95取平均学术论文更常用在安全帽检测场景我建议重点看Recall因为安全帽漏检的代价远高于误检。漏检意味着没戴安全帽的人没被识别出来可能直接导致安全事故而误检最多是后台报警频繁一些影响不大。所以调参时如果mAP和Recall有冲突优先保住Recall。有个小技巧如果你发现模型在特定距离上检测不到小目标可以尝试提高输入分辨率从640提到960或1280但注意这会显著增加推理时间部署时要权衡。另一个方案是做TTATest Time Augmentation推理时对同一张图做水平翻转、多尺度缩放综合多次结果精度有提升但速度会掉好几倍实时场景一般不用。4. 模型部署与推理优化4.1 从PyTorch权重到部署格式的转换链路训练完成后用best.pt做后续部署验证。但工业场景下很少直接跑PyTorch模型原因有三个依赖太重、推理速度不理想、跨平台兼容性差。常见的做法是转成ONNX再转TensorRT或者直接导出为TorchScript在LibTorch下运行。PPE检测这种实时视频流场景我推荐走ONNX加ONNX Runtime路线中间产物更通用后续要接不同厂家的推理框架都方便。导出命令python export.py --weights runs/train/safetynet/weights/best.pt --img-size 640 640 --batch-size 1 --simplify这里有个坑YOLOv7的export.py导出ONNX时默认输出的节点里包含一些自定义算子ONNX Runtime如果版本过旧可能不支持。建议ONNX Runtime版本不低于1.14如果遇到不支持的算子检查onnx-simplifier是否成功合并了冗余节点。4.2 推理代码的NMS逻辑解析YOLOv7的推理核心其实包含三个环节模型前向推理、置信度筛选、NMS去重。很多人对前向推理后的后处理理解不深导致自己写的推理代码出的框数量巨大。模型输出的shape一般是(1, 3*num_anchors, grid_h, grid_w, num_classes5)其中3是每个grid cell的anchor数量num_classes是这个项目的3类5是x、y、w、h和置信度。后处理的第一步是把所有grid cell的预测结果收集起来然后根据置信度阈值过滤掉低置信度框最后用NMS去掉同一目标的重复框。置信度阈值和NMS的IoU阈值是两个需要根据场景调的关键参数。安全帽检测这种场景置信度阈值我通常设为0.25NMS IoU阈值设为0.45这两个值是YOLO系列经验默认值但如果你发现误检多可以提高置信度阈值如果发现同一个目标出了很多重复框就降低NMS的IoU阈值。4.3 摄像头接入与实时推理部署到现场时视频流输入常见两种RTSP流和本地USB摄像头。代码逻辑基本通用差别只在VideoCapture的入参。RTSP地址直接作为参数传入USB摄像头传设备编号0或1。import cv2 import torch import numpy as np # 加载模型 model torch.hub.load(WongKinYiu/yolov7, custom, best.pt, force_reloadTrue) # 设置推理参数 model.conf 0.25 model.iou 0.45 cap cv2.VideoCapture(rtsp://user:passwordip:port/stream1) while True: ret, frame cap.read() if not ret: break # 推理 results model(frame, size640) # 解析结果 detections results.pandas().xyxy[0] for _, det in detections.iterrows(): cls_id int(det[class]) conf det[confidence] x1, y1, x2, y2 det[[xmin, ymin, xmax, ymax]] # 绘制边界框和标签 cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) label f{det[name]} {conf:.2f} cv2.putText(frame, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) cv2.imshow(Safety Helmet Detection, frame) if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows()实际现场部署时建议在推理线程和显示线程之间放一个队列用生产者消费者模式解耦。因为摄像头采集的帧率可能高于模型推理速度如果采集到一帧就推理一帧视频流会卡顿或延迟累积。更合理的做法是采集线程只管往队列里放帧推理线程每处理完一帧就从队列里取最新的一帧这样实时性更好也不会因为推理速度慢导致画面撕裂。4.4 TensorRT加速从20 FPS到60 FPS如果觉得PyTorch直接推理速度不够TensorRT是工业落地的最佳选择。在Jetson系列设备或NVIDIA GPU上TensorRT可以通过FP16精度和算子融合大幅提升推理速度。转TensorRT的常规路径是先转ONNX再转TRTtrtexec --onnxbest.onnx --saveEnginebest.trt --fp16 --workspace1024FP16精度在检测任务上几乎不掉点但速度能提升50%以上。我实测在Jetson Xavier NX上YOLOv7用640分辨率做安全帽检测FP16 TensorRT推理可以达到30 FPS以上完全满足现场视频流的实时需求。如果你把输入分辨率降到416甚至能跑到45 FPS以上这在边缘设备上已经非常流畅了。TensorRT还有一个好处是显存占用比PyTorch低不少因为算子被优化和融合减少了中间张量的存储。这在显存有限的嵌入式设备上是关键优势。不过TensorRT的引擎文件不能跨设备迁移比如在Xavier上转的引擎不能直接在AGX上跑必须在目标设备上重新转换一次。5. 常见问题与排查技巧实录5.1 训练时loss不下降或直接NaN这是YOLOv7新手最容易遇到的问题。loss不下降先看学习率——如果lr比较大可以尝试降低到原来的1/5再看数据集——有没有极端的标注问题比如边界框超出图像边界、类别id越界。这些看起来很基础实际排查起来往往要花不少时间。loss变成NaN则一般是数值稳定性问题。常见原因是梯度爆炸可以在cfg文件里把girgradient clip打开或调小YOLOv7默认梯度裁剪值是0.1如果你想调整需要修改train.py里的相关代码。另一个原因可能是数据里出现了全黑或全白的图片导致某些统计量异常清洗数据集可以解决。5.2 mAP很高但实际效果很差这个现象用一句话概括就是“训练集和现实世界有分布偏差”。如果验证集是从训练集里切出来的mAP高只代表模型在已知数据上表现好不代表它能泛化到现场环境。实际项目中最常见的问题是数据来源单一训练集里的图片可能都来自同一台摄像头视角、光线、背景相似模型学到的是“这个特定场景下的安全帽”而不是“普遍意义上的安全帽”。解决办法就是收集不同时间段、不同天气、不同角度、不同摄像头的数据做增强。如果你拿到的这份标注数据集建议优先补充自己现场的数据再做增量训练。5.3 推理时检测框抖动和闪烁检测框抖动是视频流部署中非常常见的问题。原因是相邻帧的画面变化很小但模型输出框的位置会有几个像素的波动显示出来就像在抖动。解决方案有两类一类是平滑滤波对同一个目标的边界框坐标做EMA指数移动平均或卡尔曼滤波另一类是跟踪算法比如ByteTrack或DeepSORT先用检测结果做目标匹配只对匹配上的目标做轨迹平滑没有匹配上的目标作为新目标处理。对于安全帽检测这种场景如果系统只做告警判定微小的框抖动可以忽略但如果要做精确的越界判断或安全区域围栏就建议引入跟踪器。5.4 数据集补充后的类别混淆问题补充数据时如果发现新场景下模型老是分不清“head”和“helmet”在类别定义不清晰时最常发生。比如有人戴了一顶颜色、形状和头发颜色相近的帽子模型可能判定为“head”而不是“helmet”。这时需要做的事是标准化标注规则锚定“人的头皮区域裸露可见”为head锚定“头部有覆盖物且覆盖物被紧戴”为helmet。在补充数据时一定严格遵守这套规则否则标注的噪声会让模型越学越迷糊。另一个技巧是增加难例挖掘的权重把那些容易混淆的样本复制多份或加重loss权重让模型更关注难区分的样本。5.5 常见问题速查表问题可能原因解决方法训练loss不下降学习率过大/数据集太脏调低lr清洗数据验证mAP低类别不平衡/样本少数据增强/加样本小目标检测不到输入分辨率低提高img size到960/1280检测框大量重复NMS IoU阈值过高把NMS IoU降到0.4左右误检太多置信度阈值过低提高到0.3或0.35部署到Jetson速度慢没用TensorRT转FP16 TRT引擎框抖动明显无平滑/跟踪加EMA或ByteTrack新场景效果差数据分布偏移增量训练自己的数据6. 从项目复现到自建数据集最后再分享一些经验整个项目跟下来我的感受是YOLOv7训练安全帽检测模型这件事技术门槛并不算高真正有价值的是工程化的思路和数据处理的经验。虽然你可以直接拿这个项目提供的模型和数据集跑通整条链路但不要止步于此——把这个流程迁移到自己实际的生产场景中才是最终目标。最后分享一个我在实际项目中反复使用的技巧先跑通最小闭环再逐步调优。刚拿到项目代码和数据集时不要一上来就去调各种超参数先按默认配置训练20个epoch生成一个粗糙的模型然后用这个模型去推理几张真实现场图片看哪些场景检测不到、哪些场景误检严重再针对性地去补数据、调参数。这个流程比盲目刷mAP高效得多因为mAP是统计指标而你在现场看到的问题才是真实需求。希望你拿到这份项目后也能从第一步跑起来然后根据自己的场景打磨出真正能用的系统。本文还有配套的精品资源点击获取
返回列表