简介一份基于YOLOv5的安全帽识别完整项目含源码、训练好的模型与配套数据集面向计算机相关专业学生、老师及企业开发者可直接用于毕业设计、课程设计或安全帽检测场景的快速落地与二次开发。该项目的核心代码已通过实际测试运行作者标明为高分毕业设计评审达到95分说明完整性和可靠性有保障。资源包共13个文件整体约49.48MB涵盖Python脚本训练、导出、YAML模型配置、预训练权重pt、数据集压缩包、License、使用说明文档及notebook示例结构清晰便于按需取用。其中pt权重可加载即用yaml便于调整网络结构py脚本支持重新训练或自定义修改ipynb则适合逐步学习和演示。目前已有134人学习浏览适合从零入门YOLOv5的目标检测新手也适合需要快速完成课设、毕设或项目初稿的开发者下载后即可依据操作说明运行并在此基础上扩展其他检测功能。 安全帽识别是Yolov5在工业视觉里落地最成功的场景之一但和Faster R-CNN这类两阶段模型相比它在边缘设备上的推理速度优势实在太明显。尤其以yolov5s为骨架的轻量版本模型体量只有十几MB在Jetson Nano上就能达到每秒20帧以上的检测能力这正好满足工地实时监控的需求。这套资源把源码、训练好的权重、数据集和操作说明都打包好了适合直接拿来做毕业设计、课程设计或者作为工程初版的验证环境。接下来我从数据组织开始把整个链路拆开看一遍顺便标记出那些容易翻车的参数和路径问题。2. 项目结构和数据准备从hat.yaml看Yolov5的数据组织2.1 解压后我们拿到了什么压缩包解开以后目录结构和官方yolov5-master基本一致但额外带了好几个定制文件。按我的习惯拿到资源的第一件事不是跑训练而是先把文件作用理清楚否则后面改错配置文件损失的是训练时间。下面是我整理出来的关键文件清单后续所有操作都围绕这些文件展开。文件/目录作用使用场景yolov5-master/完整的Yolov5训练推理框架训练、验证、导出hat.yaml安全帽数据集的配置入口训练时告知数据路径和类别名yolov5s.yaml模型结构定义文件修改类别数后传给train.pySafety helmet.pt训练好的安全帽检测权重直接推理或者重新训练起点yolov5s.ptCOCO预训练权重迁移学习初始化YOLO-hat.py针对安全帽场景的推理脚本单张图片、文件夹、视频检测export.py官方模型导出工具导出ONNX、TensorRT等格式导出数据库格式.ipynb标签格式转换工具把VOC或别的格式转成Yolo标签171265889347208773632.zip原始数据集备份解压后用于重新训练这里面hat.yaml是核心配置它决定了训练阶段去哪里读图片、目标类别有哪些。出现训练完模型但检测结果完全离谱的情况十有八九是这里的类别顺序和标注文件不一致。Yolo是用类别ID索引来对应名字的names顺序错了模型学到的语义就错了。2.2 hat.yaml的字段和路径规则hat.yaml的内容按我的标准写法是这样的我直接加了注释说明每个字段的坑# hat.yaml train: ../datasets/helmet/images/train # 训练集图片目录 val: ../datasets/helmet/images/val # 验证集图片目录 nc: 2 # 类别数helmet和head names: [helmet, head] # 类别名列表索引0和1逻辑说明train和val路径是相对于yolov5根目录的相对路径。我把数据集放在yolov5的上级目录里的datasets/helmet下这样才能用../datasets/helmet/images/train定位。如果你的数据集放在别处需要改成对应的绝对路径或相对路径。nc必须和标注txt里的最大类别ID一致比如只有helmet和head两类那索引只有0和1如果多了一个person类就要把nc改成3names里也要对应加上。参数说明Windows下很容易忽略路径分隔符问题建议把所有反斜杠统一成/否则会报FileNotFoundError。另外Yolov5不会自动划分训练集和验证集必须手动在images和labels下分好train、val两个目录且图片和对应的txt文件名要完全一致否则训练时提示“found no labels”。2.3 数据集划分与标签格式转换项目中的导出数据库格式.ipynb专门处理标签转换。如果你手里是来自LabelImg的VOC xml或者Labelme的json都需要转成Yolo的txt格式。我一般用下面这段脚本来做转换纯Python不依赖额外库方便改成json格式。import os import xml.etree.ElementTree as ET from pathlib import Path classes [helmet, head] # 必须和hat.yaml的names顺序一致 def convert_box(size, box): dw 1.0 / size[0] dh 1.0 / size[1] x (box[0] box[2]) / 2.0 * dw y (box[1] box[3]) / 2.0 * dh w (box[2] - box[0]) * dw h (box[3] - box[1]) * dh return x, y, w, h def xml_to_yolo(xml_dir, out_dir): # 遍历xml目录将每个xml转成同名txt for xml_file in Path(xml_dir).glob(*.xml): tree ET.parse(xml_file) root tree.getroot() size (int(root.find(size/width).text), int(root.find(size/height).text)) out_path Path(out_dir) / (xml_file.stem .txt) with open(out_path, w) as f: for obj in root.iter(object): cls obj.find(name).text if cls not in classes: continue cls_id classes.index(cls) box [float(obj.find(bndbox/xmin).text), float(obj.find(bndbox/ymin).text), float(obj.find(bndbox/xmax).text), float(obj.find(bndbox/ymax).text)] x, y, w, h convert_box(size, box) f.write(f{cls_id} {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n)逻辑说明convert_box把VOC格式的绝对坐标转成中心点和宽高的相对值同时除以图片长宽归一化。Yolov5训练时会把输入图片缩放成640×640如果用绝对坐标缩放后的框位置就全错了。xml_to_yolo函数遍历所有xml找到合法的类别写出每一行的目标框。参数说明如果xml里有目标被标记成helmet但你在classes里拼写成了helm那这个目标会被直接跳过最终标签数量减少模型会漏检。建议转换完以后统计一下每张图的标签数量少得离谱的图要重点排查。接下来划分数据集可以用一条shell命令配合脚本也可以直接在Jupyter里用random.shuffle。我倾向于用Python完成随机抽样保证验证集分布和训练集一致。cd yolov5 mkdir -p datasets/helmet/images/train datasets/helmet/images/val mkdir -p datasets/helmet/labels/train datasets/helmet/labels/val python split_data.py逻辑说明前两行创建目录结构split_data.py负责把图片和对应txt文件按照8:2比例随机复制到train和val里。注意txt和图片要同进同出否则训练过程中会出现找不到标签的警告。参数说明如果你的原始标注有大量背景很相近的负样本建议验证集比例高一些比如9:1避免验证时mAP虚高。3. 训练自己的安全帽数据集yolov5超参数调整与模型训练3.1 环境配置和依赖安装Yolov5的框架依赖Pytorch常规安装步骤是先建虚拟环境再装依赖。项目根目录自带requirements.txt直接执行pip install -r requirements.txt pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118逻辑说明第一行安装的是Yolov5运行所需的基础库包括numpy、opencv、matplotlib、pyyaml等。第二行安装带CUDA 11.8支持的Pytorch。如果你用的是最新RTX 40系显卡cu118版本能够稳定支持如果是30系也可以降到cu113。如果你不打算用GPU训练跑这一步可以省略但训练速度会慢到不可接受。参数说明安装完以后用python -c import torch; print(torch.cuda.is_available())验证输出True才说明GPU可用。很多训练卡住的问题都是Pytorch装成了CPU版本训练时虽然不报错但loss一直不降。3.2 修改yolov5s.yaml的模型结构配置Yolov5s.yaml开头的nc默认是80对应COCO数据集。训练自己的类别前必须把它改成2。打开文件只需改动这一行# yolov5s.yaml nc: 2 # number of classes逻辑说明yolov5s.yaml定义了模型的深度因子、通道数和检测头结构。它不算完整网络定义而是让Yolov5自动生成Backbone和Head。更改nc后训练时就会输出2个类别的预测结果。如果你想要更高精度可以换yolov5l.yaml或yolov5x.yaml但显存占用和推理延迟都会成倍增加。参数说明有些项目会直接把整个yaml里的nc改成类别数而忽略anchors也需要对应调整。Yolov5会自动重新计算anchors但需要加--noautoanchor参数来禁止多数情况下保持默认即可。3.3 训练命令的核心参数拆解我用实际训练安全帽数据集的命令来拆解参数选择这个命令和官方文档保持一致也符合hat.yaml的配置结构。python train.py --img 640 --batch 16 --epochs 100 \ --data hat.yaml --cfg yolov5s.yaml \ --weights yolov5s.pt --name helmet_exp逻辑说明--img 640表示输入图片尺寸为640×640这是速度和精度的平衡点--batch 16在8G显存上比较稳妥如果你的显卡是12G以上可以调到32--epochs 100对于二分类任务来说能充分收敛再多就容易过拟合--data hat.yaml是刚才配置好的数据集入口--cfg yolov5s.yaml指定模型结构--weights yolov5s.pt表示加载COCO预训练权重进行迁移学习。用预训练权重比从头训练收敛速度快得多这也是Yolov5在小型数据集上的最佳实践。参数说明如果你显存只有6G把batch改为8同时加--cache-images来缓存图片到内存减少磁盘IO压力。如果看到loss数值震荡严重可以把--lr0改成0.005降低初始学习率。3.4 训练过程中的loss曲线与超参数调整训练结束后会在runs/train/helmet_exp/目录下生成weights/best.pt和last.pt。best.pt是验证集上mAP最高的权重也就是项目里那个Safety helmet.pt的来源。训练过程中我习惯用下面表格里的超参数做微调参数建议值影响batch8-32越大收敛越稳定但显存占用高lr00.01初始学习率过大会导致震荡mosaic1.0拼接4张图做增强过拟合时调大fl_gamma0.0focal loss参数的强弱正负样本极不平衡时调为1.5逻辑说明mosaic是Yolov5最核心的数据增强手段它把四张图片拼接成一张训练样本极大提高模型对遮挡和小目标的鲁棒性。如果你自己训练的模型在检测远处人员时经常漏检把mosaic从1.0提高到1.5会有效果但训练时间也会变长。fl_gamma是focal loss的调节因子当数据集中正样本戴头盔和负样本不戴头盔数量悬殊时提高它能让模型更多关注难分类样本。参数说明修改超参数有两种方式一是直接改data/hyps/hyp.scratch-low.yaml文件二是在命令行追加--hyp /path/to/yaml。我推荐用后者这样不同实验之间不会互相污染配置文件。训练过程中可以通过tensorboard --logdir runs/train实时观察loss曲线。如果发现验证集loss在第60轮开始反弹说明过拟合了需要减小模型复杂度或增加数据增强。4. 加载训练好的模型推理与ONNX导出4.1 用YOLO-hat.py做单张图片和视频推理项目自带的YOLO-hat.py是官方detect.py的精简版把标签和输出格式调成了安全帽场景。实际操作时一条命令就可以测试模型python YOLO-hat.py --source ./test.jpg --weights Safety helmet.pt --conf 0.25逻辑说明--source支持图片路径、目录、视频文件甚至摄像头RTSP流。脚本会先做letterbox缩放然后推理并画出边界框。绿色框是helmet类红色框是head类。--conf 0.25表示置信度阈值低于0.25的预测会被过滤掉如果你的现场误检较多可以调到0.5但会牺牲一部分召回率。参数说明如果你输入的视频帧率很高可以用--vid-stride 2让模型每隔一帧推理一次这样在CPU上也能做到接近实时的效果。输出结果保存在runs/detect/exp目录每次执行自动递增下标。4.2 用export.py导出ONNX格式项目中的BestONNX.zip说明作者已经把模型转成了ONNX格式。实地部署时ONNX是无处不在的标准交换格式因为很多推理引擎都支持它。导出命令如下python export.py --weights Safety helmet.pt --include onnx --simplify逻辑说明--include onnx指定导出ONNX--simplify调用onnx-simplifier对计算图做简化减少冗余算子。导出的Safety helmet.onnx会放在weights目录下。ONNX模型可以再接上OpenVINO、TensorRT或者ONNX Runtime运行适合将推理逻辑从训练框架中解耦。参数说明如果你后续要部署到NVIDIA嵌入式设备可以加--half导出半精度FP16权重体积缩小一半速度也有提升。如果遇到导出后模型输出维度对不上先确认Pytorch和onnx的版本兼容。4.3 ONNX模型输出维度和NMS后处理导出后的ONNX模型输出形状是[1, 25200, 7]25200等于三个尺度的锚框总数7分别代表[x, y, w, h, conf, class0, class1]。直接用onnxruntime推理时还需要自己写NMS。参考这个代码片段import onnxruntime as ort import numpy as np import cv2 sess ort.InferenceSession(Safety helmet.onnx) # 假设输入是640x640先做letterbox这里用简化示例 image cv2.imread(test.jpg) resized cv2.resize(image, (640, 640)) x np.transpose(resized.astype(np.float32) / 255.0, (2, 0, 1))[None] preds sess.run(None, {sess.get_inputs()[0].name: x})[0] # preds形状为(1, 25200, 7)这里对第一个batch做解析 boxes preds[0] # 过滤低置信度按类别分别NMS逻辑说明这段代码展示的是从ONNX推理到拿到原始预测框的过程。preds[0]是25200个候选框每个框有位置、置信度和类别得分。NMS通常用torchvision.ops.nms或onnxruntime扩展实现但最基本的做法是先过滤掉conf小于0.25的框再按IoU阈值去除重复框。这一部分在YOLO-hat.py里已经实现了你自己处理ONNX时要补上。参数说明不同Yolov5版本导出的ONNX输出格式可能不一样有的是一个张量有的是三个尺度的输出。建议导出后用onnxruntime打印输出节点的shape再做后续处理。5. 跑通整个流程的常见坑与实用技巧5.1 数据集路径错误与标签类别ID错位我在几十次实验里见过的最典型错误是训练后模型完全无法检测或者把所有的框都画错颜色。原因几乎都指向两个地方——路径和类别顺序。hat.yaml里的train和val路径要确保yolov5的工作目录能从相对路径找到你的图片否则训练会正常启动但数据量为0。标签转换时classes列表的顺序必须和names一致不然标注文件中写0模型就会认为它是names[0]里的helmet。5.2 显存不足时从yolov5s开始迭代训练安全帽模型时不需要一开始就上yolov5x。用yolov5s把流程跑通再把模型切换成yolov5l或yolov5m这能为你省下大把的调试时间。当显存溢出时第一时间把batch降到4并检查--workers参数是否过高推荐改成4。另外在训练命令中加入--cache-images可以大幅提升数据读取速度但需要至少4G可用内存否则会拖慢系统。5.3 用TensorRT加速边缘部署如果你的最终目标是进入真实生产我建议把ONNX再转成TensorRT engine。TensorRT在NVIDIA显卡上能获得3到5倍的推理速度提升尤其适合Jetson Nano这类功耗受限设备。转换方式很简单NVIDIA官方提供了trtexec工具一条命令即可完成trtexec --onnxSafety helmet.onnx --saveEnginesafety.engine --fp16。转换时需要显卡与目标部署环境一致比如在PC上转好再拷到Jetson上往往需要重新构建。这个优化留给那些想继续在项目上做扩展的同学答辩时提到这点会显得项目更完整。本文还有配套的精品资源点击获取