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
PyQt5==5.15.3
PyQtChart==5.15.3
PyQt5-tools
GPUtil
或者使用我的环境(推荐)
pip install -r requirements.txt
官方权重
我已上传到一份到百度云:下载链接 , 密码: 44qm
训练自己的数据
提示:
关于增加数据集分类的方法,请看【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 创建每个图片对应的标签文件
你可以使用 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 除以图像高度。代码
如下:
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 = int(box[3])
dw = np.float32(1. / int(size[0]))
dh = np.float32(1. / int(size[1]))
w = x2 - x1
h = y2 - y1
x = x1 + (w / 2)
y = y1 + (h / 2)
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return [x, y, w, h]
生成的 .txt 文件放置的名字是图片的名字,放置在 label 文件夹中,例如:
./score/images/train/00001.jpg # image
./score/labels/train/00001.txt # label
生成的 .txt 例子:
1 0.1830000086920336 0.1396396430209279 0.13400000636465847
0.15915916301310062