0 前期教程
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使用较为简单，参考下图即可。
