YOLOv5实现佩戴安全帽检测和识别
1. 前言
安全帽是作业场所作业时头部防护所用的头部防护用品,它对使用者的头部在受坠落物或小型飞溅物体等其他因素引起的伤害起到防护作用。近年来,因不佩戴安全帽、不规范佩戴安全帽等原因导致的安全生产事故屡禁不止,事故发生背后的影响是巨大的,不仅为家人带来巨大的伤痛,也为企业的利益带来巨大的损失。而如何使员工规范佩戴安全帽、保障员工和企业双方利益,成为了一直以来各方坚持不懈想要实现的目标。因此,研究佩戴安全帽的监测算法,具有重大的意义和广泛的应用价值。
本篇博客,将手把手教你搭建一个基于YOLOv5的佩戴安全帽目标检测识别方法。目前,基于YOLOv5s的目标检测的佩戴安全帽识别方法的平均精度平均值mAP_0.5=0.93,mAP_0.5:0.95=0.63,基本满足业务的性能需求。
另外,为了能部署在手机Android平台上,本人对YOLOv5s进行了模型轻量化,开发了一个轻量级的版本yolov5s05_416和yolov5s05_320,在普通Android手机上可以达到实时的检测和识别效果,CPU(4线程)约30ms左右,GPU约25ms左右 ,基本满足业务的性能需求。
先展示一下Python版本佩戴安全帽检测和识别Demo效果(head表示(头部)未佩戴安全帽,helmet表示佩戴了安全帽):
2. 佩戴安全帽检测和识别的方法
Helmet-Asian亚洲人佩戴安全帽数据集:总共有19000+图片,VOC的XML数据格式,可直接用于目标检测模型训练。
Helmet-Europe欧洲人佩戴安全帽数据集:总共有3000+图片,VOC的XML数据格式,可直接用于目标检测模型训练。
(2)自定义数据集
如果需要增/删类别数据进行训练,或者需要自定数据集进行训练,可参考如下步骤:
采集图片,建议不少于200张图片
使用Labelme等标注工具,对目标进行拉框标注:labelme工具:GitHub - wkentaro/labelme: Image Polygonal Annotation with Python (polygon, rectangle, circle, line, point and image-level flag annotation).
将标注格式转换为VOC数据格式,参考工具:labelme/labelme2voc.py at main · wkentaro/labelme · GitHub
生成训练集train.txt和验证集val.txt文件列表
修改engine/configs/voc_local.yaml的train和val的数据路径
重新开始训练

4. 基于YOLOv5的佩戴安全帽模型训练
(1)YOLOv5安装
训练Pipeline采用YOLOv5: https://github.com/ultralytics/yolov5 , 原始代码训练需要转换为YOLO的格式,不支持VOC的数据格式。为了适配VOC数据,本人新增了LoadVOCImagesAndLabels用于解析VOC数据集进行训练。另外,为了方便测试,还增加demo.py文件,可支持对图片,视频和摄像头的测试。
Python依赖环境,使用pip安装即可:
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.2
Pillow
PyYAML>=5.3.1
scipy>=1.4.1
torch>=1.7.0
torchvision>=0.8.1
tqdm>=4.41.0
tensorboard>=2.4.1
seaborn>=0.11.0
pandas
thop # FLOPs computation
pybaseutils
(2)准备Train和Test数据
下载Helmet-Asian和Helmet-Europe佩戴安全帽数据集,总共约2W+的图片
Helmet-Asian Helmet-Europe
(3)配置数据文件
修改训练和测试数据的路径:engine/configs/voc_local.yaml
注意数据路径分隔符使用【/