YOLOv5自定义数据集训练

简介

本文介绍如何在自己的VOC格式数据集上训练YOLO5目标检测模型。

VOC数据集格式

首先，先来了解一下Pascal VOC数据集的格式，该数据集油5个部分组成，文件组织结构如下，目前主要的是VOC和VOC.

- VOC- JPEGImages- 1.jpg- 2.jpg- ...- Annotations- 1.xml- 2.xml- ...- ImageSets- Main- train.txt- val.txt- test.txt- trainval.txt- ...- SegmentationClass- SegmentationObject

第一个文件夹JPEGImages为所有的图像，也就是说，训练集、验证集和测试集需要自己划分；Annotations为JPEGImages文件夹中每个图片对应的标注，xml格式文件，文件名与对应图像相同；ImageSets主要的子文件夹为Main，其中有四个文本文件，为训练集、验证集、测试集和训练验证集的图片文件名；SegmentationClass和SegmentationObject文件夹存放分割的结果图，前者为语义分割，后者为实例分割。

上述xml标注文件，格式如下。对其具体标注解释。

<annotation><folder>down</folder> # 图片所处文件夹<filename>1.jpg</filename> # 图片文件名及后缀<path>./savePicture/train_29635.jpg</path> # 存放路径<source> #图源信息<database>Unknown</database> </source><size> # 图片尺寸和通道<width>640</width><height>480</height><depth>3</depth></size><segmented>0</segmented> #是否有分割label，0无1有# 图像中包含的所有目标，一个目标一个object标签<object><name>car</name> # 目标类别<pose>Unspecified</pose> # 目标的姿态<truncated>0</truncated> # 目标是否被部分遮挡（>15%）<difficult>0</difficult> # 是否为难以辨识的目标， 需要结合背景才能判断出类别的物体<bndbox> # 目标边界框信息<xmin>2</xmin><ymin>156</ymin><xmax>111</xmax><ymax>259</ymax></bndbox></object><object><name>multi_signs</name><editType /><pose>Unspecified</pose><truncated>0</truncated><difficult>0</difficult><bndbox><xmin>81</xmin><ymin>98</ymin><xmax>154</xmax><ymax>243</ymax></bndbox></object></annotation>

也就是说，遇到这种文件格式的数据（主要特点为图像全放在一个文件夹，标注格式如上等），将其作为VOC格式的数据集，将自己的数据集重构为VOC格式以便开源项目的处理。

自定义训练

下载源码

通过git clone git@:ultralytics/yolov5.git将YOLOv5源码下载到本地，本文后面的内容也可以参考官方的自定义数据集训练教程，不同于我的教程，该教程全面包含了VOC格式和COCO格式数据集的处理方法。

此时创建虚拟环境，并通过pip install -r requirements.txt安装依赖包，我这里测试过，最新的项目是兼容Pytorch 1.6的，1.6之前的Pytorch会有一些问题（这是基于本文发布的时候，随着yolo5的更新对版本的要求会更高，后续只要按照官方教程配置环境即可）。

数据集处理

一般，符合VOC格式的数据集至少包含图像和标注两个文件夹，结构如下。我这里假定测试集是独立的，该数据集实际上为训练集，只需要划分出训练集和验证集即可。这里建议将文件夹重命名如下，否则后续可能出现数据集加载失败的情况。

- 根目录- images- Annotations

下面，编写脚本划分数据集，split_train_val.py脚本内容如下（参考Github上的开源脚本），只需要执行python split_train_val.py --xml_path dataset_root/Annotations/ --txt_path dataset_root/anno_txt/就得到了划分结果的文件列表，如训练集对应的train.txt如下图，里面与训练图片所有的文件名。

import osimport randomimport argparseparser = argparse.ArgumentParser()parser.add_argument('--xml_path', type=str, help='input xml label path')parser.add_argument('--txt_path', type=str, help='output txt label path')opt = parser.parse_args()trainval_percent = 1.0train_percent = 0.8xmlfilepath = opt.xml_pathtxtsavepath = opt.txt_pathtotal_xml = os.listdir(xmlfilepath)if not os.path.exists(txtsavepath):os.makedirs(txtsavepath)num = len(total_xml)list_index = range(num)tv = int(num * trainval_percent)tr = int(tv * train_percent)trainval = random.sample(list_index, tv)train = random.sample(trainval, tr)file_trainval = open(txtsavepath + '/trainval.txt', 'w')file_test = open(txtsavepath + '/test.txt', 'w')file_train = open(txtsavepath + '/train.txt', 'w')file_val = open(txtsavepath + '/val.txt', 'w')for i in list_index:name = total_xml[i][:-4] + '\n'if i in trainval:file_trainval.write(name)if i in train:file_train.write(name)else:file_val.write(name)else:file_test.write(name)file_trainval.close()file_train.close()file_val.close()file_test.close()

接下来，我们要做的就是每个xml标注提取bbox信息为txt格式（这种数据集格式成为yolo_txt格式），每个图像对应一个txt文件，文件每一行为一个目标的信息，包括类别 xmin xmax ymin ymax。使用的脚本voc_label.py内容如下（注意，类别要替换为当前数据集的类别列表），在数据集根目录（此时包含Annotations、anno_txt以及images三个文件夹的目录）下执行该脚本，如python ../../utils/voc_label.py。

# -*- coding: utf-8 -*-import xml.etree.ElementTree as ETimport osfrom os import getcwdsets = ['train', 'val', 'test']classes = ['window_shielding', 'multi_signs', 'non_traffic_signs']abs_path = os.getcwd()def convert(size, box):dw = 1. / (size[0])dh = 1. / (size[1])x = (box[0] + box[1]) / 2.0 - 1y = (box[2] + box[3]) / 2.0 - 1w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn x, y, w, hdef convert_annotation(image_id):in_file = open('Annotations/%s.xml' % (image_id))out_file = open('labels/%s.txt' % (image_id), 'w')tree = ET.parse(in_file)root = tree.getroot()size = root.find('size')w = int(size.find('width').text)h = int(size.find('height').text)for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),float(xmlbox.find('ymax').text))b1, b2, b3, b4 = b# 标注越界修正if b2 > w:b2 = wif b4 > h:b4 = hb = (b1, b2, b3, b4)bb = convert((w, h), b)out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')wd = getcwd()for image_set in sets:if not os.path.exists('labels/'):os.makedirs('labels/')image_ids = open('anno_txt/%s.txt' % (image_set)).read().strip().split()list_file = open('%s.txt' % (image_set), 'w')for image_id in image_ids:list_file.write(abs_path + '/images/%s.jpg\n' % (image_id))convert_annotation(image_id)list_file.close()

这时候，我们的目标检测数据集就构建完成了，其内容如下，其中labels中为不同图像的标注文件，train.txt等几个根目录下的txt文件为划分后图像所在位置的绝对路径，如train.txt就含有所有训练集图像的绝对路径。

配置文件

下面需要两个配置文件用于模型的训练，一个用于数据集的配置，一个用于模型的配置。

首先是数据集的配置，在根目录下的data目录下新建一个yaml文件，内容如下，首先是训练集和验证集的划分文件，这个文件在上面一节最后生成得到了，然后是目标的类别数目和具体类别列表，这个列表务必和上一节最后voc_label.py中的一致。

train: dataset/train/train.txtval: dataset/train/val.txt# number of classesnc: 3# class namesnames: ['window_shielding', 'multi_signs', 'non_traffic_sign']

然后，编辑模型的配置文件，此时需要先在项目根目录下的weights目录下执行其中的download_weights.sh这个shell脚本来下载四种模型的权重。然后，选择一个模型，编辑项目根目录下models目录中选择的模型的配置文件，将第一个参数nc改为自己的数据集类别数即可，例如我使用yolov5x模型，则修改yolov5x.yaml文件。这里weights的下载可能因为网络而难以进行，我也将其上传到了百度网盘，地址给出，提取码为vjlx。

模型训练

此时，可以使用下面的命令进行模型的训练，训练日志默认保存在./runs/下，包括模型参数、Tensorboard记录等。此时TensorBoard以已经默认打开，浏览器访问效果如下图（由于数据量很小，很快过拟合）。

python train.py --img 640 --batch 8 --epoch 300 --data ./data/ads.yaml --cfg ./models/yolov5x.yaml --weights weights/yolov5x.pt --device '0'

模型测试

接着，就是在有标注的测试集或者验证集上进行模型效果的评估，在目标检测中最常使用的指标为mAP。通过下面的命令进行模型测试，由于这是个比赛，测试集没有标注，这里使用验证集作为测试用数据，下述命令只需要指定数据集配置文件和训练结果模型即可。

python test.py --data ./data/ads.yaml --weights ./runs/exp0/weights/best.pt --augment

不进行测试时数据增强和进行测试时数据增强（TTA）在验证集上的表现分别如下。

Class Images Targets P RmAP@.5 mAP@.5:.95all 400970 0.376 0.4410.35 0.235

Class Images Targets PRmAP@.5 mAP@.5:.95all4009700.272 0.532 0.366 0.24

模型推理

最后，模型在没有标注的数据上进行推理，使用下面的命令（该命令中save-txt选项用于生成结果的txt标注文件，不指定则只会生成结果图像）。其中，weights使用最满意的实验即可，source则提供一个包含所有测试图片的文件夹即可。

python detect.py --weights runs/exp0/weights/best.pt --source ./dataset/test/ --device 0 --save-txt

这样，对每个测试图片会在默认的inference/output文件夹中生成一个同名的txt文件，按照我的需求修改了detect.py文件后，每个txt会生成一行一个目标的信息，信息包括类别序号 置信度 xcenter ycenter w h，后面四个为bbox位置，均未归一化。如下图。

我这里因为是一个比赛，再将这个txt处理为了json文件。不论是这里的处理代码还是上面对detect.py修改的代码，都可以在文末给出的Github仓库找到。

补充说明

本文介绍了如何使用YOLOv5在自己的数据集上进行训练，按部就班地进行了讲解。该项目在YOLOv5地源码基础上修改完成。