免费获取学习方案
ARTICLE DETAIL

资讯详情

深耕编程基础知识与建站技术分享的一线实战洞察。

YOLOv5电动自行车头盔识别:从数据集标注到TensorRT部署全流程

YOLOv5电动自行车头盔识别:从数据集标注到TensorRT部署全流程 简介这份资源面向计算机视觉方向的本科生与研究生提供一套可直接运行的电动自行车头盔佩戴识别检测方案适用于毕业设计、期末大作业或课程实践。项目基于YOLOv5算法完成训练与推理全流程代码注释详尽新手也能快速理解模型结构与检测逻辑部署后即可用于图片或视频中的头盔佩戴状态判别。压缩包共187个文件约134.13MB包含55个Python源码文件、45张png效果图、22个yaml配置、7个pt权重文件以及前端页面所需的js、css、html与Dockerfile等覆盖训练、测试、可视化展示各环节。目前已有385人学习下载。资源内含完整模型权重与数据配置可省去从零搭建的繁琐过程便于读者直接复现检测效果、撰写论文实验章节并在此基础上做二次开发或功能扩展。1. 电动自行车头盔识别到底难在哪从 YOLOv5 的落地边界说起电动自行车头盔佩戴识别检测说白了就是给路口、园区、小区门口装一双电子眼自动判断骑行人头上有没有戴头盔。这件事听起来像是一个标准的目标检测任务但真正做过的人都知道它比识别猫狗难得多。原因在于头盔目标小、遮挡严重、逆光与夜间场景多、骑行人姿态千变万化而且戴了和没戴之间的视觉差异有时候只有几十个像素。YOLOv5 之所以成为这类毕设和工程项目的首选是因为它在精度和速度之间取得了很好的平衡640 分辨率下单卡推理可以轻松跑到几十甚至上百 FPS模型文件小、部署链路成熟从训练到 ONNX 再到 TensorRT 都有现成路径。这个方向适合谁如果你是做计算机视觉毕设的学生或者需要在园区、工地、社区做轻量级安全合规检测的工程师这套方案是性价比很高的起点。它不需要你从零设计网络结构也不需要海量算力一张消费级显卡就能完成训练和验证。但要注意能跑通 demo 和能上线用是两回事中间隔着数据集质量、标注一致性、后处理逻辑和场景适配这几道坎。接下来我会把从数据准备到模型导出、再到实际部署的完整路径拆开讲重点放在那些真正会让人翻车的地方。2. 数据集准备与标注头盔识别的地基怎么打2.1 类别设计与标注规范头盔识别最常见的类别划分有两种二分类helmet / no_helmet和三分类helmet / no_helmet / person。二分类适合只关心头上有没有头盔的场景标注框直接框住头部区域三分类则多一个 person 类用于先定位骑行人再判断头部状态。我一般推荐二分类起步因为标注成本低、类别间混淆少而且实际部署时你只需要关心头部区域。标注工具用 LabelImg 或 X-AnyLabeling 都可以导出 YOLO 格式的 txt 文件。关键规范有几条第一标注框要紧贴头盔或头部轮廓不要框到肩膀或车身第二对于模糊、严重遮挡的目标如果人眼都难以判断直接跳过不标不要强行给一个模棱两可的标签第三夜间和逆光样本必须单独检查一遍这类样本的标注一致性最容易出问题。一个容易被忽略的点是负样本的采集。很多人只收集戴头盔和没戴头盔的骑行人图片但实际场景中还有大量行人、停放的车辆、广告牌上的人像等干扰目标。如果训练集里没有这些负样本模型上线后会把广告牌上的头盔图案也识别成目标。建议负样本占比控制在 10% 到 20% 之间。2.2 数据增强策略与 YOLOv5 超参数配置YOLOv5 自带的数据增强已经相当丰富包括 mosaic、mixup、随机缩放、随机裁剪、HSV 色彩空间扰动等。对于头盔识别任务我建议在默认基础上做几处调整。mosaic 增强默认开启概率为 1.0它把四张图拼成一张对小目标检测非常友好但如果你发现训练后期 loss 震荡严重可以把 mosaic 关闭概率调高让模型在最后几十个 epoch 用原始分布收敛。超参数文件在data/hyps/hyp.scratch-low.yaml中关键参数如下# 学习率与动量 lr0: 0.01 # 初始学习率头盔任务建议 0.01 起步 lrf: 0.01 # 最终学习率系数余弦退火到 lr0 * lrf momentum: 0.937 # SGD 动量 weight_decay: 0.0005 # 损失权重 box: 0.05 # 框回归损失权重 cls: 0.5 # 分类损失权重 cls_pw: 1.0 # 分类正样本权重 obj: 1.0 # 目标置信度损失权重 # 数据增强 hsv_h: 0.015 # 色调扰动 hsv_s: 0.7 # 饱和度扰动 hsv_v: 0.4 # 亮度扰动 degrees: 0.0 # 旋转角度头盔任务不建议大角度旋转 translate: 0.1 # 平移 scale: 0.5 # 缩放 shear: 0.0 # 剪切 perspective: 0.0 # 透视变换 flipud: 0.0 # 上下翻转头盔场景不建议开启 fliplr: 0.5 # 左右翻转 mosaic: 1.0 # mosaic 增强概率 mixup: 0.0 # mixup 增强小数据集可开 0.1这里有几个参数值得展开说。degrees我设为 0因为骑行人头部在画面中基本是正立或轻微倾斜的大角度旋转会引入不真实的样本分布。flipud上下翻转也关掉现实中不会出现倒着骑车的场景。scale设为 0.5 是为了让模型适应不同距离的目标路口摄像头和园区门口摄像头的安装距离差异很大缩放增强能提升泛化能力。数据集目录结构按 YOLOv5 的标准来dataset/ ├── images/ │ ├── train/ │ ├── val/ │ └── test/ ├── labels/ │ ├── train/ │ ├── val/ │ └── test/ └── helmet.yamlhelmet.yaml内容path: ./dataset train: images/train val: images/val test: images/test nc: 2 names: [helmet, no_helmet]划分比例建议 7:2:1如果数据量少于 3000 张可以调整为 8:1:1验证集至少保留 200 张以上才能稳定评估。标注完成后一定要写一个脚本检查标签文件确认没有越界坐标、没有空标签文件、类别 ID 没有超出 nc 范围。这些低级错误在训练时不会报错但会直接拉低 mAP排查起来非常费时间。3. 用 YOLOv5 训练头盔检测模型从命令行到收敛判断3.1 环境搭建与最小训练命令环境搭建这一步我踩过最多的坑是 PyTorch 版本和 CUDA 版本不匹配。YOLOv5 对 PyTorch 版本有一定要求太新的版本可能遇到 API 变动太旧的版本又不支持某些算子。我一般用 PyTorch 1.13 到 2.0 之间的版本CUDA 11.7 或 11.8配合对应的 cuDNN。安装命令# 创建虚拟环境 conda create -n helmet python3.9 -y conda activate helmet # 安装 PyTorch以 CUDA 11.8 为例 pip install torch2.0.1 torchvision0.15.2 --index-url https://download.pytorch.org/whl/cu118 # 克隆 YOLOv5 并安装依赖 git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt验证环境是否正常python -c import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))如果输出True和显卡型号说明环境没问题。如果输出False检查 CUDA 驱动版本是否匹配或者重新安装对应版本的 PyTorch。最小训练命令python train.py \ --img 640 \ --batch 16 \ --epochs 100 \ --data dataset/helmet.yaml \ --weights yolov5s.pt \ --cfg models/yolov5s.yaml \ --hyp data/hyps/hyp.scratch-low.yaml \ --name helmet_exp \ --project runs/train参数说明--img 640是输入分辨率头盔目标较小如果显存允许可以提到 1280但训练时间会显著增加--batch 16根据显存调整8G 显存用 yolov5s 跑 640 分辨率大概能开到 16 到 24--weights yolov5s.pt使用预训练权重从头训练在小数据集上几乎不可能收敛--cfg指定模型结构yolov5s 是最小的如果精度不够可以换 yolov5m 或 yolov5l。3.2 训练过程监控与收敛判断训练启动后终端会输出每个 epoch 的损失值和 mAP。但只看终端输出不够YOLOv5 会自动生成 TensorBoard 日志在runs/train/helmet_exp/目录下。启动 TensorBoardtensorboard --logdir runs/train重点看几条曲线train/box_loss、train/obj_loss、train/cls_loss应该整体下降并趋于平稳metrics/mAP_0.5和metrics/mAP_0.5:0.95应该上升并趋于饱和。如果 box_loss 下降但 mAP 不涨说明模型在过拟合需要增加数据增强或减少模型复杂度。如果 obj_loss 震荡严重可能是学习率太大把lr0降到 0.005 试试。训练过程中还有一个容易被忽略的指标val/obj_loss。如果训练集的 obj_loss 持续下降但验证集的 obj_loss 开始上升这是过拟合的明确信号。我一般会在验证集 obj_loss 连续 10 个 epoch 不下降时提前停止训练而不是死等 100 个 epoch 跑完。关于训练轮数头盔识别任务在 3000 到 5000 张数据集上yolov5s 通常 80 到 120 个 epoch 就能收敛。如果超过 150 个 epoch 还在涨说明数据量不够或者标注质量有问题加轮数不如加数据。3.3 模型评估与推理验证训练完成后用val.py在测试集上评估python val.py \ --data dataset/helmet.yaml \ --weights runs/train/helmet_exp/weights/best.pt \ --img 640 \ --task test \ --conf-thres 0.25 \ --iou-thres 0.45--conf-thres 0.25是置信度阈值低于这个值的检测框会被过滤--iou-thres 0.45是 NMS 的 IoU 阈值控制重叠框的合并程度。这两个参数对最终指标影响很大建议在验证集上做一轮网格搜索找到最适合你场景的组合。推理单张图片python detect.py \ --weights runs/train/helmet_exp/weights/best.pt \ --source test_images/ \ --img 640 \ --conf-thres 0.25 \ --save-txt \ --save-conf--save-txt会把检测结果保存为 YOLO 格式的 txt 文件--save-conf会在 txt 中附带置信度。这两个选项在后续做误检分析时非常有用可以快速定位哪些图片的检测结果异常。评估指标方面头盔识别任务我一般关注三个数mAP0.5 达到 0.85 以上算及格0.90 以上算良好no_helmet 类的召回率要特别关注因为漏检一个没戴头盔的人比误检一个戴头盔的人后果更严重推理速度在目标硬件上要满足实时性要求通常单帧处理时间不超过 50ms。4. 避坑与排查头盔识别项目里最容易翻车的五件事4.1 现象模型把头盔识别成 no_helmet置信度还很高原因训练集中戴头盔和没戴头盔的样本在视觉上高度相似尤其是深色头盔和黑色头发在低分辨率下几乎无法区分。标注时如果边界模糊模型学到的特征就是混乱的。解决第一检查标注一致性把那些疑似戴了但不确定的样本重新过一遍统一标准第二在数据增强中增加亮度扰动和对比度扰动让模型学会关注形状而非颜色第三如果两类样本数量差异大用--weights加载预训练权重后先冻结主干训练几个 epoch再解冻全量微调。4.2 现象训练 loss 正常下降但验证集 mAP 始终在 0.5 左右上不去原因最常见的原因是标注格式错误。YOLO 格式要求坐标是归一化后的中心点 x、中心点 y、宽度、高度且都在 0 到 1 之间。如果标注工具导出的是 VOC 格式的绝对坐标直接拿来训练就会出现这个问题。解决写一个检查脚本遍历所有标签文件确认坐标范围在 0 到 1 之间且宽度和高度大于 0。另外检查类别 ID 是否从 0 开始YOLOv5 要求类别 ID 从 0 到 nc-1。import os def check_labels(label_dir): issues [] for fname in os.listdir(label_dir): if not fname.endswith(.txt): continue path os.path.join(label_dir, fname) with open(path, r) as f: lines f.readlines() if len(lines) 0: issues.append(f{fname}: 空标签文件) continue for i, line in enumerate(lines): parts line.strip().split() if len(parts) ! 5: issues.append(f{fname} 第{i1}行: 字段数不对) continue cls_id int(parts[0]) coords [float(x) for x in parts[1:]] if cls_id 0 or cls_id 1: issues.append(f{fname} 第{i1}行: 类别ID越界 {cls_id}) if any(c 0 or c 1 for c in coords): issues.append(f{fname} 第{i1}行: 坐标越界 {coords}) if coords[2] 0 or coords[3] 0: issues.append(f{fname} 第{i1}行: 宽高非正 {coords}) return issues issues check_labels(dataset/labels/train) for issue in issues[:20]: print(issue) print(f共发现 {len(issues)} 个问题)4.3 现象白天检测正常夜间几乎全部漏检原因训练集中夜间样本太少模型没有学到夜间场景下的头盔特征。另外夜间红外摄像头拍出的画面是灰度的如果训练集全是彩色图模型对灰度图的泛化能力会很差。解决第一补充夜间样本至少占总数据的 20% 到 30%第二在数据增强中增加灰度化处理让模型适应灰度输入第三如果硬件支持在推理前做直方图均衡化提升夜间画面的对比度。4.4 现象推理速度远低于预期达不到实时要求原因可能是输入分辨率设得太大或者模型选得太大或者没有做推理优化。yolov5s 在 640 分辨率下用 TensorRT 加速后单卡可以跑到几百 FPS但如果用 PyTorch 原生推理且不做任何优化可能只有几十 FPS。解决第一导出 ONNX 或 TensorRT 引擎用export.py脚本python export.py \ --weights runs/train/helmet_exp/weights/best.pt \ --include onnx engine \ --img 640 \ --batch 1 \ --device 0第二如果部署在边缘设备上考虑用 yolov5n 或 yolov5s并开启 FP16 量化第三检查数据预处理和后处理是否成为瓶颈尤其是 NMS 部分可以用 GPU 版本的 NMS 替代 CPU 版本。4.5 现象模型在测试集上表现很好但实际部署后误检率很高原因测试集和实际场景的数据分布不一致。测试集可能来自同一个摄像头、同一时间段而实际部署后摄像头角度、光照、天气都在变化。解决第一测试集要覆盖不同时间段、不同天气、不同摄像头角度的样本第二部署后持续收集误检样本定期做增量训练第三在推理后处理中加入业务规则过滤比如根据检测框的宽高比、位置、与人体框的相对关系来排除明显不合理的检测结果。5. 模型导出与部署从 PyTorch 权重到可用的推理服务5.1 ONNX 导出与推理验证ONNX 是跨平台部署的通用格式导出命令python export.py \ --weights runs/train/helmet_exp/weights/best.pt \ --include onnx \ --img 640 \ --batch 1 \ --opset 12 \ --simplify--opset 12是 ONNX 算子集版本12 兼容性较好--simplify会调用 onnx-simplifier 简化计算图减少冗余算子。导出完成后用 onnxruntime 验证import onnxruntime as ort import numpy as np import cv2 # 加载 ONNX 模型 session ort.InferenceSession(best.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider]) # 预处理 img cv2.imread(test.jpg) img cv2.resize(img, (640, 640)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img img.transpose(2, 0, 1).astype(np.float32) / 255.0 img np.expand_dims(img, axis0) # 推理 inputs {session.get_inputs()[0].name: img} outputs session.run(None, inputs) print(outputs[0].shape) # 应该是 (1, 25200, 7)25200 是候选框数量7 是 4 坐标 1 置信度 2 类别ONNX 输出的后处理需要自己实现包括置信度过滤、NMS、坐标还原。YOLOv5 的官方仓库里有utils/general.py中的non_max_suppression函数可以直接复用。5.2 TensorRT 加速与边缘设备部署如果部署在 NVIDIA 显卡或 Jetson 设备上TensorRT 能带来数倍的推理加速。导出命令python export.py \ --weights runs/train/helmet_exp/weights/best.pt \ --include engine \ --img 640 \ --batch 1 \ --device 0 \ --half--half开启 FP16 精度在支持 FP16 的显卡上可以进一步提速精度损失通常在 1% 以内。TensorRT 引擎是硬件相关的换显卡需要重新导出。对于树莓派、RK3568 这类边缘设备TensorRT 不可用需要走 ONNX Runtime 或厂商提供的推理框架。RK3568 有 RKNN 工具链可以把 ONNX 转成 RKNN 格式利用 NPU 加速。转换流程大致是ONNX → RKNN 量化 → RKNN 推理。量化时需要用一批校准图片建议从训练集中随机抽取 200 到 500 张覆盖不同光照和场景。5.3 推理服务封装与性能调优实际部署时通常需要把模型封装成一个 HTTP 或 gRPC 服务。用 FastAPI 封装一个最小推理服务from fastapi import FastAPI, UploadFile import numpy as np import cv2 import onnxruntime as ort app FastAPI() session ort.InferenceSession(best.onnx, providers[CUDAExecutionProvider]) def preprocess(image_bytes): img cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR) img cv2.resize(img, (640, 640)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img img.transpose(2, 0, 1).astype(np.float32) / 255.0 return np.expand_dims(img, axis0) def postprocess(output, conf_thres0.25, iou_thres0.45): # 这里复用 YOLOv5 的 non_max_suppression 逻辑 # 返回检测框列表 pass app.post(/detect) async def detect(file: UploadFile): image_bytes await file.read() input_tensor preprocess(image_bytes) outputs session.run(None, {session.get_inputs()[0].name: input_tensor}) detections postprocess(outputs[0]) return {detections: detections}性能调优方面几个关键点第一预处理用 GPU 加速OpenCV 的cuda模块或者 DALI 库都可以第二批处理推理如果并发量高把多张图拼成一个 batch 一起推理吞吐量能提升数倍第三后处理的 NMS 用 GPU 实现避免 CPU 成为瓶颈第四如果服务部署在多卡机器上用多进程或多线程并行推理每个进程绑定一张卡。关于量化INT8 量化能进一步提速但对小目标检测的精度影响较大。头盔目标本身就小INT8 量化后 mAP 可能掉 5 到 10 个点。如果精度要求高建议用 FP16 而不是 INT8。如果必须用 INT8一定要用足够的校准数据并且在量化后重新评估 mAP确认精度损失在可接受范围内。6. 把 mAP 再拉高几个点几个我反复验证过的技巧训练完第一版模型后大多数人会卡在一个瓶颈上mAP 到 0.85 左右就上不去了。这时候加数据、加轮数效果都不明显需要换思路。我试过几种方法下面这几个是投入产出比最高的。第一个技巧是用 TTATest Time Augmentation做推理增强。YOLOv5 的detect.py和val.py都支持--augment参数开启后会对每张图做多尺度、多翻转推理然后合并结果。代价是推理时间增加 2 到 3 倍但 mAP 通常能涨 1 到 3 个点。如果服务对延迟不敏感这个技巧几乎零成本。python val.py \ --data dataset/helmet.yaml \ --weights runs/train/helmet_exp/weights/best.pt \ --img 640 \ --augment \ --task test第二个技巧是调整 NMS 的 IoU 阈值和置信度阈值。默认的 0.45 和 0.25 不一定适合头盔场景。头盔目标密集时IoU 阈值太高会导致漏检太低会导致误检。我一般会在验证集上跑一组网格搜索import subprocess import itertools conf_list [0.15, 0.2, 0.25, 0.3, 0.35] iou_list [0.3, 0.4, 0.45, 0.5, 0.6] best_map 0 best_params None for conf, iou in itertools.product(conf_list, iou_list): result subprocess.run([ python, val.py, --data, dataset/helmet.yaml, --weights, runs/train/helmet_exp/weights/best.pt, --img, 640, --conf-thres, str(conf), --iou-thres, str(iou), --task, test ], capture_outputTrue, textTrue) # 从输出中解析 mAP for line in result.stdout.split(\n): if mAP0.5 in line and all in line: map_val float(line.split()[-1]) if map_val best_map: best_map map_val best_params (conf, iou) break print(f最佳参数: conf{best_params[0]}, iou{best_params[1]}, mAP{best_map:.4f})第三个技巧是对 no_helmet 类做重采样。如果 no_helmet 的召回率偏低可以在训练时给这个类更高的损失权重。YOLOv5 的hyp.yaml里有cls_pw参数默认是 1.0可以调到 1.5 或 2.0让模型更关注分类错误。另外如果两类样本数量不平衡可以在数据集层面做过采样把 no_helmet 样本复制到与 helmet 样本数量相当。第四个技巧是用更大的输入分辨率做微调。先用 640 训练到收敛再用 1280 微调 10 到 20 个 epoch。大分辨率对小目标检测提升明显但训练时间和显存占用都会增加。如果显存不够可以把 batch 降到 4 或 8同时开启梯度累积。最后一个习惯每次训练完我都会把误检和漏检的样本单独挑出来看一遍。YOLOv5 的detect.py支持--save-txt和--save-conf把检测结果和原图对照着看能发现很多指标上看不出来的问题。比如模型可能把某个特定颜色的头盔全部漏掉或者把某个角度的骑行人全部误判。这些问题靠调参解决不了只能靠补充针对性数据。我一般会把这些 bad case 整理成一个单独的文件夹下一轮训练时优先补充这些场景的样本。希望这些经验能帮你在头盔识别这个方向上少走点弯路。本文还有配套的精品资源点击获取
返回列表