yolo
安全帽检测
该项目是使用 YOLOv5 v2.x
的程序来训练在智能工地安全领域中头盔目标检测的应用, 先来一波演示!
指标
yolov5s 为基础训练,epoch = 50
分类 | P | R | mAP0.5 |
---|---|---|---|
总体 | 0.884 | 0.899 | 0.888 |
人体 | 0.846 | 0.893 | 0.877 |
头 | 0.889 | 0.883 | 0.871 |
安全帽 | 0.917 | 0.921 | 0.917 |
对应的权重文件:百度云,提取码: b981
yolov5m 为基础训练,epoch = 100
分类 | P | R | mAP0.5 |
---|---|---|---|
总体 | 0.886 | 0.915 | 0.901 |
人体 | 0.844 | 0.906 | 0.887 |
头 | 0.9 | 0.911 | 0.9 |
安全帽 | 0.913 | 0.929 | 0.916 |
对应的权重文件:百度云,提取码: psst
yolov5l 为基础训练,epoch = 100
分类 | P | R | mAP0.5 |
---|---|---|---|
总体 | 0.892 | 0.919 | 0.906 |
人体 | 0.856 | 0.914 | 0.897 |
头 | 0.893 | 0.913 | 0.901 |
安全帽 | 0.927 | 0.929 | 0.919 |
对应的权重文件:百度云,提取码: a66e
1.YOLO v5训练自己数据集教程
使用的数据集:Safety-Helmet-Wearing-Dataset ,感谢这位大神的开源数据集!
本文结合 YOLOv5官方教程 来写
环境准备
首先确保自己的环境:
1 | Python >= 3.7 |
官方权重
我已上传到一份到百度云:下载链接 , 密码: 44qm
训练自己的数据
提示:
关于增加数据集分类的方法,请看【5. 增加数据集的分类】
1.1 创建自己的数据集配置文件
因为我这里只是判断 【人没有带安全帽】、【人有带安全帽】、【人体】 3个类别 ,基于 data/coco128.yaml
文件,创建自己的数据集配置文件 custom_data.yaml
1 |
|
1.2 创建每个图片对应的标签文件
你可以使用 data/gen_data/gen_head_helmet.py
来将 VOC
的数据集转换成 YOLOv5
训练需要用到的格式。
使用标注工具类似于 Labelbox 、CVAT 、精灵标注助手 标注之后,需要生成每个图片对应的 .txt
文件,其规范如下:
- 每一行都是一个目标
- 类别序号是零索引开始的(从0开始)
- 每一行的坐标
class x_center y_center width height
格式 - 框坐标必须采用归一化的 xywh格式(从0到1)。如果您的框以像素为单位,则将
x_center
和width
除以图像宽度,将y_center
和height
除以图像高度。代码如下:
1 | import numpy as np |
生成的 .txt
文件放置的名字是图片的名字,放置在 label 文件夹中,例如:
1 | ./score/images/train/00001.jpg # image |
生成的 .txt
例子:
1 | 1 0.1830000086920336 0.1396396430209279 0.13400000636465847 0.15915916301310062 |
1.3 文件放置规范
文件树如下
1.4 聚类得出先验框(Yolov5 内部已做适配,可选)
使用代码 ./data/gen_anchors/calculate_anchors.py
,修改数据集的路径
1 | FILE_ROOT = r"xxx" # 根路径 |
跑完会生成一个文件 anchors.txt
,里面有得出的建议先验框:
1 | Best Accuracy = 79.72% |
1.5 选择一个您需要的模型
在文件夹 ./models
下选择一个你需要的模型然后复制一份出来,将文件开头的 nc =
修改为数据集的分类数,下面是借鉴 ./models/yolov5s.yaml
来修改的
1 | # parameters |
1.6 开始训练
这里选择了 yolov5s
模型进行训练,权重也是基于 yolov5s.pt
来训练
1 | python train.py --img 640 --batch 16 --epochs 10 --data ./data/custom_data.yaml --cfg ./models/custom_yolov5.yaml --weights ./weights/yolov5s.pt |
其中,yolov5s.pt
需要自行下载放在本工程的根目录即可,下载地址 官方权重
1.7 看训练之后的结果
训练之后,权重会保存在 ./runs
文件夹里面的每个 exp
文件里面的 weights/best.py
,里面还可以看到训练的效果
2. 侦测
侦测图片会保存在 ./inferenct/output/
文件夹下
运行命令:
1 | python detect.py --source 0 # webcam |
例如使用我的 s
权重检测图片,可以运行以下命令,侦测图片会保存在 ./inferenct/output/
文件夹下
1 | python detect.py --source 图片路径 --weights ./weights/helmet_head_person_s.pt |
3. 检测危险区域内是否有人
3.1 危险区域标注方式
我这里使用的是 精灵标注助手 标注,生成了对应图片的 json 文件
3.2 执行侦测
侦测图片会保存在 ./inferenct/output/
文件夹下
运行命令:
1 | python area_detect.py --source ./area_dangerous --weights ./weights/helmet_head_person_s.pt |