YOLOV5训练数据集的完整过程

1、下载YOLOV5

git clone https://github.com/ultralytics/yolov5.git

2、安装依赖包

cd yolov5
pip install -r requirements.txt

3、准备数据集（VOC）

1）在yolov5目录下创建face_dataset文件夹，将labelImage标注好的xml文件（face_dataset/Annotations）和原始图片（face_dataset/images）放到对应的目录下，目录结构如下：

Annotations文件夹下面为xml文件（标注工具labelImage生成的标注文件），内容如下：

images文件夹下面为VOC数据集格式中的JPEGImages，内容如下：

2）在face_dataset目录下创建split_train_val.py脚本，通过该脚本划分数据集，划分为：训练集、验证集和测试集，代码如下：

import os
import random

trainval_percent = 1.0
train_percent = 0.9
xmlfilepath = 'Annotations'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(num * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

# ImageSets目录不存在，就创建
if not os.path.exists('ImageSets/'):
    os.makedirs('ImageSets/')
# ImageSets/Main目录不存在，就创建
if not os.path.exists('ImageSets/Main/'):
    os.makedirs('ImageSets/Main/')

ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/Main/val.txt', 'w')

for i in list:
    name = '/opt/PycharmProjects/yolov5/face_dataset/images/' + total_xml[i][:-4] + '.jpg' + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

执行split_train_val.py脚本，会生成如下文件夹和文件（蓝色部分由代码生成的文件夹和文件）：

生成的test.txt、train.txt、trainval.txt、val.txt文件存放是images的绝对路径，以train.txt为例，内容如下：

test.txt内容为空，其他两个文件内容与train.xml相似，只是存放不同图片的文件路径而已

3）准备labels，在face_dataset目录下创建voc_label.py脚本，该脚本是将voc的数据集xml文件转化为yolo数据集的txt文件

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

image_sets = ['train', 'val', 'test']

classes = ["person"]


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

    isValid = True
    if x == 0 or y == 0 or w == 0 or h == 0:
        print(f'x :{x} y:{y} w:{w} h:{h}')
        isValid = False
    if box[0] == 0 or box[1] == 0 or box[2] == 0 or box[3] == 0:
        print(f'box[0] :{box[0]} box[1]:{box[1]} box[2]:{box[2]} box[3]:{box[3]}')
        isValid = False

    return (x, y, w, h), isValid


def convert_annotation(image_id):
    in_file = open('Annotations/%s.xml' % (image_id.split("/").pop().split(".")[0]))
    out_file = open('labels/%s.txt' % (image_id.split("/").pop().split(".")[0]), 'w')
    print(out_file)
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    # print(root['object'])
    # return
    objList = root.findall('object')
    # for obj in objList:
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        print('-------')
        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))
        bb, isValid = convert((w, h), b)
        if isValid == False:
            # out_file.close()
            print(image_id)
            return False
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
        print("文件内容：")
        print(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    out_file.flush()
    out_file.close()
    return True


if not os.path.exists('labels/'):
    os.makedirs('labels/')
for image_set in image_sets:
    # strip() 移除字符串的首尾字符，默认为空格
    # split() 字符串分割，默认为所有空字符，包含空格、换行、制表符
    image_ids = open('ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
    list_file = open('%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        isValid = convert_annotation(image_id)
        if isValid == True:
            list_file.write('%s\n' % (image_id))
        else:
            print(f'存在无效值:{image_id}')
    list_file.close()

执行voc_label.py脚本，会生成如下文件夹和文件（蓝色部分由代码生成的文件夹和文件）：

其中test.txt、train.txt、val.txt文件内容与2）中生成的test.txt、train.txt、val.txt文件内容是一样的，保存的是图片的完整路径，labels文件夹下的数据，就是yolo的txt文本，内容如下：

两行数据表示当前对应的图片包含两个目标类目，0表示类目下标（下标从0开始），0.21306818181818182和0.20666666666666667分别表示标注框的中心点的x和y的坐标（其中坐标已经归一化），0.30113636363636365和0.3244444444444444分别表示标注框的宽度和高度（同中心坐标一样也进行了归一化）

至此数据的准备部分已经完成了

4、配置文件

1）在yolov5/data文件夹下，创建person.yaml文件，文件内容如下：

train: /opt/PycharmProjects/yolov5/face_dataset/train.txt
val: /opt/PycharmProjects/yolov5/face_dataset/val.txt

# number of classes
nc: 1

# class names
names: ['person']

2）编辑模型配置文件

在yolov5/model文件夹下模型配置文件，包含了s、m、l、x4个版本，逐渐增大（随着架构的增大，训练时间也是逐渐增大），本文以s版本为例，修改yolov5s.yaml，修改后的文件内容如下（只修改蓝色部分的内容）：

5、训练模型

1）下载预训练模型，执行以下会自动下载模型文件

cd weights ./download_weights.sh

命令执行完后，在weights目录下会生成如下四个预训练模型文件（蓝色部分）

2）训练

yolov5的train.py脚本的参数如下：

参数说明：

epochs：指的就是训练过程中整个数据集将被迭代多少次,显卡不行你就调小点。

batch-size：一次看完多少张图片才进行权重更新，梯度下降的mini-batch,显卡不行你就调小点。

cfg：存储模型结构的配置文件

data：存储训练、测试数据的文件

img-size：输入图片宽高,显卡不行你就调小点。

rect：进行矩形训练

resume：恢复最近保存的模型开始训练

nosave：仅保存最终checkpoint

notest：仅测试最后的epoch

evolve：进化超参数

bucket：gsutil bucket

cache-images：缓存图像以加快训练速度

weights：权重文件路径

name： 重命名results.txt to results_name.txt

device：cuda device, i.e. 0 or 0,1,2,3 or cpu

adam：使用adam优化

multi-scale：多尺度训练，img-size +/- 50%

single-cls：单类别的训练集

根据自己的实际设备情况修改参数，然后执行脚本，我的执行参数如下：

python3 train.py --img-size 256 --batch-size 16 --epoch 10 --data data/person.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --device 'cpu'

执行完成之后，会生成runs文件夹，目录结构如下：

每次执行train.py脚本的时候都会在runs/trarin文件夹下生成一个exp文件夹，第一次文件夹的名字为exp，第二次文件夹的名字为exp1，第三次文件夹的名字为exp2，以此类推

6、模型测试

yolov5的test.py脚本的参数如下：

执行如下命令进行模型测试：

python3 test.py --data data/person.yaml --weights runs/train/exp/weights/best.pt --augment

执行完成之后，在runs文件夹下多出test文件夹，目录结构如下：

同训练一样，每次执行test.py脚本的时候都会在runs/test文件夹下生成一个exp文件夹，第一次文件夹的名字为exp，第二次文件夹的名字为exp1，第三次文件夹的名字为exp2，以此类推

7、模型推理

最后在没有标注的数据集上进行推理，yolov5目录下的detect.py脚本为推理脚本，脚本参数如下：

执行如下命令进行推理

python3 detect.py --weights runs/train/exp/weights/best.pt --img-size 256 --source inference/images/ --device cpu --save-txt

参数说明：

--weights：权重文件路径

--img-size：图片大小

--source：图片的路径

--device：cuda device, i.e. 0 or 0,1,2,3 or cpu

--save-txt：save results to *.txt

上行上面命令后，在runs文件夹下会多出detect文件夹，目录结构如下：

同训练和模型测试一样，每次执行上面命令，都会在runs/detect目录下生成exp文件，第一次执行的文件名为exp，第二次执行的文件名为exp1，第三次执行的文件名为exp3，以此类推

至此YOLOv5训练自己的数据集整个过程：制作数据集----模型训练----模型测试----模型推理阶段已全部完成。

最后附上整个项目的源码及数据集：https://gitee.com/hagler/yolov5.git