YOLOv5训练自定义数据集

之前对YOLO的一些基础知识以及数据集的内容进行了简单介绍。下面通过一个实际的项目来加深对YOLO的应用。

在这里要实现一个检测人物是否带头盔的需求,特别是一些安全生产的项目中。对应的数据集选择是开源的harijawahar/Helmet_Detection。

该数据集采用COCO格式进行标注,因此其中的坐标值分别为左上角x,y值及边界框的宽高。可以使用如下的代码将其数据集转换为YOLO的格式:

import json
import shutil
from pathlib import Path

image_dict = {}
image_arr = []

def mkdir_dir(path):
	p = Path(path)
	if not p.exists():
		p.mkdir()
output_directory = "helmet"
mkdir_dir(output_directory)
mkdir_dir("helmet/train")
mkdir_dir("helmet/train/images")
mkdir_dir("helmet/train/labels")
mkdir_dir("helmet/val")
mkdir_dir("helmet/val/images")
mkdir_dir("helmet/val/labels")

with open("train/_annotations.coco.json",encoding="utf-8") as f:
	data = f.read()
	json_data = json.loads(data)
	images = json_data.get("images", [])
	annotations = json_data.get("annotations", [])
for d in images:
	_id = d.get("id")
	filename = d.get("file_name")
	height = d.get("height")
	width = d.get("width")
	image_arr.append((filename,width,height))
	shutil.copy2("train/{}".format(filename),"helmet/train/images/{}.jpg".format(_id))
for d in annotations:
	image_id = d.get("image_id")
	category_id = d.get("category_id")
	bbox = d.get("bbox")
	x1,y1,width2,height2 = bbox
	x2 = x1 + width2
	y2 = y1 + height2
	x1 = int(x1)
	x2 = int(x2)
	y1 = int(y1)
	y2 = int(y2)
	with open("helmet/train/labels/{}.txt".format(image_id),"a+",encoding="utf-8") as f:
		f.write("{} {} {} {} {}\n".format(category_id,(x1+width2//2)/width,(y1+height2//2)/height,width2/width,height2/height))

我们将对应的数据集下载解压到YOLOv5项目目录下,首先创建相应的训练及验证目录,每个目录下有2个子目录images和labels分别存储图片及对应的标签文件。

根据数据集每个目录下的_annotations.coco.json的内容得到文件名称、分类类型、边界框的值。其中类型主要有3类,分别是不确定、带头盔、没带头盔。将左上角坐标加上边界框一般的宽高后归一化得到YOLO的格式。

之后在YOLOv5源代码data目录下新增1个helmet.yaml的配置文件,其内容如下:

train: ./helmet/helmet/train/
val: ./helmet/helmet/val/

nc: 3

names: ["undefined","helmet","no helmet"]

下面就可以开始进行模型训练后,其对应的目录如下:

python train.py --data data/helmet.yaml --cfg models/yolov5m.yaml --weights weights/yolov5m.pt --epochs 100 --batch-size 8

这里模型使用的是YOLOv5m,训练轮数为100次,每个批次大小为8。经过几个小时的训练后得到最终的模型文件。

之后就可以使用其检测代码了:

$ python detect.py --weights /path/to/best.pt --source helmet/test

我们对测试集目录下的所有图片进行检测。其最终的效果如下:

可以看到实际该数据集对图片进行强制拉伸为640x640的大小,从而导致图片会出现一些变形。更多情况还是推荐使用图像强化的方式会更好些。关于数据增强可以考虑使用albumentations。

如果喜欢这篇文章或对您有帮助，可以：[☕] 请我喝杯咖啡 | [💓] 小额赞助

码力全开 / YOLOv5训练自定义数据集