【YOLO】yolov5训练自己的数据集

0 前期教程

• 【Python】朴实无华的yolov5环境配置（一）

1 前言

  上面前期教程中，大致介绍了yolov5开发环境的配置方法和yolov5项目的基本结构，下一步就是基于yolov5预训练模型来训练自己的数据集，这对于只是想要使用yolov5这个工具的人，还是想要深入研究yolov5类似的目标识别算法的人，都是绕不开的入门操作，本文将根据查找的资料以及自己的经验来简单介绍这个过程。

2 准备数据集

2.1 数据集来源

  由于数据标注这个过程过于繁琐，耗时长，所以个人觉得可以先拿网上公开数据集来尝试，学会之后再尝试自己标注数据然后训练，只是后者多了一个标注的步骤，其他都是完全一样的。
  这里推荐的数据集是PASCAL VOC2012数据集，它是一个世界级计算机视觉比赛专用的数据集，包含了生活中常见的20种目标，其数据类别分布如下图所示。

官网下载链接

  打开之后，如下图所示，直接点击第一个下载即可。

2.2 数据集结构介绍

  下载完得到的是一个压缩包，它文件夹结构如下所示

├───Annotations       //标注文件
├───ImageSets         //图像数据
│   ├───Action        //人物动作
│   ├───Layout        //人各个部位数据集
│   ├───Main          //主要数据集（含训练集和测试集）
│   └───Segmentation  //语义分割训练和测试集
├───JPEGImages        //所有图片
├───SegmentationClass //语义分割类别
└───SegmentationObject//语义分割物体

其中，Main文件夹下的数据集分布非常有规律：

每一种目标都对应三个txt文件，以train结尾的为训练集，以val结尾的为测试集，以trainval结尾的为训练集和测试集之和，这里以bottle该目标的数据集为例：

其中，前面一列为JPEGImages文件夹下对应的图片文件名，后一列表示目标信息：-1代表没有该目标；1表示有该目标；0表示该目标检测有困难。

2.3 标签格式的转换

  Annotations文件夹下存放着所有图片的标签信息，标签文件名和图片文件名是一一对应的，使用时根据文件名查找即可。
  VOC数据集中的标签是xml格式，但yolov5训练时要使用的是yolo格式，因此需要先进行转换。以其中一个标签为例：

对于目标识别任务来说，需要用到的主要是两个标签，分别是<size>标签和<object>标签。其中<size>标签主要是描述图片的大小信息；<object>标签描述图片中的目标信息，一个object标签对应一个目标，其中<name>为目标信息，<truncated>表示目标是否截断（1为截断）；<difficult>表示目标是否难以检测（1为难检测）；<bndbox>表示目标的左上角和右下角坐标。

  而yolo格式的标签如下所示

分别表示：[目标类别（一般用数字表示） x_center y_center width height]，因此，在训练该数据集之前，还需要进行标签格式的转换，主要使用到的就是xml这个库来对xml文件进行解析。

示例代码如下所示：

import xml.etree.ElementTree as ET
import os
import shutil
import tqdm

def convert(size, box):
    '''	@func: 将box的坐标转换为yolo需要的格式
        @para	size: 图片的尺寸, eg:[500, 200]
        @para	box: box的坐标, [xmin, xmax, ymin, ymax]
        @return: 转换后的yolo格式坐标[x_center, y_center, width, height]
    '''
    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 transfer(xmlfile:str, txtfile:str, classes:list[str]=['bottle']):
    '''	@func: 将xml文件转换为txt文件
        @para	xmlfile: xml文件的路径
        @para	txtfile: txt文件的路径
        @return: None
    '''
    in_file = open(xmlfile, encoding='utf-8')
    out_file = open(txtfile, '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'):
        cls = obj.find('name').text
        if cls not in classes: continue
        num_id = classes.index(cls) #找到该类别的序号
        xmlbox = obj.find('bndbox')
        # 坐标转换
        box = [float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
               float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text)]
        bb = convert((w, h), box)
        out_file.write(str(num_id) + ' ' + " ".join([str(a) for a in bb]) + '\n')


def construct(train_set, val_set, JPEG_Path, Ano_Path, dataset_path='./dataset', classes:list[str]=['bottle']):
    '''	@func: 将VOC数据集构建成yolo训练所需的数据集格式
        @para	train_set: 训练集的路径txt
        @para	val_set: 验证集的路径txt
        @para	JPEG_Path: VOC数据集图片文件的路径
        @para	Ano_Path: VOC数据集标注文件所在路径
        @para	dataset_path: 构造好的数据集所在路径,可以不存在, 默认为当前路径
        @para	classes: 类别列表
        @return: None
    '''
    os.makedirs(dataset_path, exist_ok=True)
    img_path = os.path.join(dataset_path, 'images'); os.makedirs(img_path, exist_ok=True)
    img_train_path = os.path.join(img_path, 'train'); os.makedirs(img_train_path, exist_ok=True)
    img_val_path = os.path.join(img_path, 'val'); os.makedirs(img_val_path, exist_ok=True)
    label_path = os.path.join(dataset_path, 'labels'); os.makedirs(label_path, exist_ok=True)
    label_train_path = os.path.join(label_path, 'train'); os.makedirs(label_train_path, exist_ok=True)
    label_val_path = os.path.join(label_path, 'val'); os.makedirs(label_val_path, exist_ok=True)

    with open(train_set, 'r') as f:
        print('Start converting train set...')
        for line in tqdm.tqdm(f.readlines()):
            num = line.strip().split()[1]
            filename = line.strip().split()[0]
            shutil.copy(os.path.join(JPEG_Path, filename + '.jpg'), img_train_path)
            if num == '-1':
                with open(os.path.join(label_train_path, filename + '.txt'), 'w') as f:
                    f.write("") # 写入空值
            else:
                transfer(os.path.join(Ano_Path, filename + '.xml'),
                         os.path.join(label_train_path, filename + '.txt'), classes)

    with open(val_set, 'r') as f:
        print('Start converting val set...')
        for line in tqdm.tqdm(f.readlines()):
            num = line.strip().split()[1]
            filename = line.strip().split()[0]
            shutil.copy(os.path.join(JPEG_Path, filename + '.jpg'), img_val_path)
            if num == '-1':
                with open(os.path.join(label_val_path, filename + '.txt'), 'w') as f:
                    f.write("")
            else:
                transfer(os.path.join(Ano_Path, filename + '.xml'),
                         os.path.join(label_val_path, filename + '.txt'), classes)

    print('Done!')

    # 写入yaml文件
    with open(os.path.join(dataset_path, 'dataset.yaml'), 'w') as f:
        f.write('path: {}\n'.format(dataset_path))
        f.write('train: images/train\n')
        f.write('val: images/val\n\n')
        f.write('nc: {}\n'.format(len(classes)))
        f.write('names: {}'.format(classes))

if __name__ ==  "__main__":
    dataset = r'C:\Users\Zeoy\Desktop\dataset' # 构造好的数据集所在路径
    train_set = r'C:\Users\Zeoy\Desktop\VOC2012\ImageSets\Main\bottle_train.txt'
    val_set = r'C:\Users\Zeoy\Desktop\VOC2012\ImageSets\Main\bottle_val.txt'
    JPEG_Path = r'C:\Users\Zeoy\Desktop\VOC2012\JPEGImages'
    Ano_Path = r'C:\Users\Zeoy\Desktop\VOC2012\Annotations'
    classes = ['bottle'] # 如果类别较多，可以用用numpy读取txt
    construct(train_set, val_set, JPEG_Path, Ano_Path, dataset, classes)

使用时，注意替换main部分的train_set, val_set, JPEG_Path, Ano_Path。

这个代码需要注意的是，在使用yolov5进行模型训练时，填入的参数并不是数据集的路径，而是yaml文件，而在yaml文件中，只有图片所在路径，而没有标注文件所在路径，如下图所示：

这是因为yolov5默认读取的文件夹路径是固定的，要求文件夹名字和下面保持一致：

├───images
│   ├───train
│   └───val
└───labels
    ├───train
    └───val

3 训练以及训练结果

3.1 训练

  上一步得到了数据集和yaml文件，下一步就是利用该数据集进行训练得到pt模型。这里使用的是最外层文件夹的train.py文件，使用方法参考文件开头的注释即可。

直接在命令行输入：

python train.py --data "C:\Users\Zeoy\Desktop\dataset\dataset.yaml" --weights yolov5s.pt --img 640 --batch-size -1

其中，--data参数为刚刚构建好的数据集所在的路径；--batch-size参数为一次性读取图片的数量，设置为-1程序将根据显卡的显存大小自动分配。

进入到如下界面，就可以耐心等待了，默认的epoch数为100，如果觉得不合适可以添加参数--epochs xxxx来设置。

3.2 测试

  训练完毕后，在run/train/exp?文件夹下有一个weights文件夹，里面即是训练得到的权值模型，以pt结尾，一般会有两个：best.pt 和 last.pt，一般选择前者。接下来就是利用detect.py文件读取模型并对测试图片进行处理。

  方法和运行train.py差不多，首先看一下文件开始的注释部分，就知道怎么使用了

可以看出，yolov5测试时的输入数据可以是图片，视频，包含图片或视频的文件夹，甚至是摄像头和网络视频流，使用非常方便。而使用的模型格式也可以是多种深度学习框架下的结果。

  复制一下pt文件的路径，然后在命令行再次输入：

python detect.py --weights 'C:\Users\Zeoy\Desktop\Code\Python\yolov5-master\runs\train\exp19\weights\best.pt' --source 'C:\Users\Zeoy\Desktop\img.png'

--weights参数即刚刚训练得到的pt文件路径，--source参数即需要测试的数据来源。

4 数据标注

  以上就是基于VOC数据集并利用yolov5训练并测试的全部过程了，前面也提到，使用开源数据集还是自己数据集的差别就是标注的这个过程，因此，再简单介绍一下数据集标注的方法。

  目前网上绝大多数标注图片数据的方法都是使用labelImg这个工具，它是python的一个第三方库，可以直接使用pip或conda安装：

pip install labelImg

  不过，尤为需要注意的是，这个库对于新版的python支持似乎不够好，在运行时会经常出现卡退的情况，需要修改该库的内部代码，我比较懒，这里是直接使用python3.8安装（恰好电脑上也有python3.8），运行是没有问题的。

  labelImg使用较为简单，参考下图即可。

文章来源地址https://www.toymoban.com/news/detail-501787.html

5 后续教程

• 【YOLO】目标识别模型的导出和opencv部署（三）

到了这里，关于【YOLO】yolov5训练自己的数据集的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！