YOLOv5自定义数据集训练

YOLOv5自定义数据集训练

简介

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

VOC数据集格式

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

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

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

自定义训练

下载源码

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

此时创建虚拟环境，并通过pip install -r requirements.txt安装依赖包，我这里测试过，最新的项目是兼容Pytorch 1.6的，1.6之前的Pytorch会有一些问题。

数据集处理

一般，符合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 os
import random
import argparse

parser = 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.0
train_percent = 0.8
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_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格式，每个图像对应一个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 ET
import os
from os import getcwd

sets = ['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 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h


def 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').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_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 = w
        if b4 > h:
            b4 = h
        b = (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.txt
val: dataset/val.txt
 
# number of classes
nc: 3
 
# class names
names: ['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       R      [email protected]    [email protected]:.95
all    400      970    0.376    0.441     0.35       0.235

Class  Images  Targets    P      R      [email protected]    [email protected]:.95
all     400     970     0.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地源码基础上修改完成，代码开源于我的Github，欢迎star或者fork。