前言
Amusi 发现一个很棒的开源项目,利用YOLOv5进行目标检测的"落地化"应用:安全帽佩戴检测。
该项目使用了YOLOv5s、YOLOv5m、YOLOv5l来训练安全帽佩戴检测数据集,代码和权重均已开源!安全帽佩戴检测数据集也是开源的(共含7581 张图像)!
项目教程也相当详细,推荐入门练手学习!而且有意思的是,该项目和数据集的两位作者均是中国人,点赞!
项目链接(文末附下载):
https://github.com/PeterH0323/Smart_Construction
数据集链接(文末附下载):
https://github.com/njvisionpower/Safety-Helmet-Wearing-Dataset
如果帮到您请给个 star ✨✨✨,您的 star 是我最大的鼓励!
Smart_Construction
该项目是使用 YOLOv5 来训练在智能工地安全领域中头盔目标检测的应用
指标
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官方教程 来写
环境准备
首先确保自己的环境:
Python >= 3.7
Pytorch == 1.5.x
训练自己的数据
提示:
关于增加数据集分类的方法,请看【5. 增加数据集的分类】
1.1 创建自己的数据集配置文件
因为我这里只是判断 【人没有带安全帽】、【人有带安全帽】、【人体】 3个类别 ,基于 data/coco128.yaml
文件,创建自己的数据集配置文件 custom_data.yaml
# 训练集和验证集的 labels 和 image 文件的位置
train: ./score/images/train
val: ./score/images/val
# number of classes
nc: 3
# class names
names: ['person', 'head', 'helmet']
1.2 创建每个图片对应的标签文件
使用标注工具类似于 Labelbox 、CVAT 、精灵标注助手 标注之后,需要生成每个图片对应的 .txt
文件,其规范如下:
- 每一行都是一个目标
- 类别序号是零索引开始的(从0开始)
- 每一行的坐标
class x_center y_center width height
格式 - 框坐标必须采用归一化的 xywh格式(从0到1)。如果您的框以像素为单位,则将
x_center
和width
除以图像宽度,将y_center
和height
除以图像高度。代码如下:
import numpy as np
def convert(size, box):
"""
将标注的 xml 文件生成的【左上角x,左上角y,右下角x,右下角y】标注转换为yolov5训练的坐标
:param size: 图片的尺寸: [w,h]
:param box: anchor box 的坐标 [左上角x,左上角y,右下角x,右下角y,]
:return: 转换后的 [x,y,w,h]
"""
x1 = int(box[0])
y1 = int(box[1])
x2 = int(box[2])
y2