YOLOv5白皮书-第Y2周：训练自己的数据集

>- **🍨 本文为[🔗365天深度学习训练营]中的学习记录博客**
>- **🍖 原作者：[K同学啊]**

本人往期文章可查阅：深度学习总结

接着上一篇文章 YOLOv5白皮书-第Y1周：调用官方权重进行检测，配置完YOLOv5需要的环境后，我们试着用YOLOv5训练自己的数据。（在开始本项目之前，记得先跑一遍入门篇，确保环境是正常的）

🏡 我的环境：

语言环境：Python3.11

编译器：PyCharm

深度学习环境：Pytorch

- torch==2.0.0+cu118

- torchvision==0.18.1+cu118
显卡：NVIDIA GeForce GTX 1660

一、准备好自己的数据

我的目录结构是这样子的：

将会看到如下的目录结构：

Annotation文件夹为xml文件，我的文件如下：

images文件是.png格式，官方的为.jpg，不过问题不大，后面改一下代码即可(后面会讲解)

二、运行 split train val.py 文件

lmageSets文件夹下面有个Main子文件夹，其下面存放了 train.txt、val.txt、test.txt 和 trainval.txt 四个文件，它们是通过split_train_val.py 文件来生成的。
split_train_val.py 文件的位置如下：

split train val.py 文件的内容如下：

# 导入必要的库
import os,random,argparse

# 创建一个参数解析器
parser=argparse.ArgumentParser()

# 添加命令行参数，用于指定XML文件的路径，默认为‘Annotation'文件夹
parser.add_argument('--xml_path',default='Annotations',type=str,help='input xml label path')

#添加命令行参数，用于指定输出txt标签文件的路径，默认为'ImageSets/Main'文件
parser.add_argument('--txt_path',default='ImageSets/Main',type=str,help='output txt label path')

#解析命令行参数
opt=parser.parse_args()

#定义训练验证集和测试集的划分比例
trainval_percent=1.0  # 使用全部数据
train_percent=0.9  # 训练集占训练验证集的90%

# 设置XML文件夹的路径，根据命令行参数指定
xmlfilepath=opt.xml_path

# 设置输出txt标签文件的路径，根据命令行参数指定
txtsavepath=opt.txt_path

# 获取XML文件夹中的所有XML文件列表
total_xml=os.listdir(xmlfilepath)

# 如果输出txt标签文件的文件夹不存在，创建它
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

# 获取XML文件的总数
num=len(total_xml)

# 创建一个包含所有XML文件索引的列表
list_index=range(num)

# 计算训练验证集的数量
tv=int(num*trainval_percent)

# 计算训练集的数量
tr=int(tv*train_percent)

# 从所有XML文件索引中随机选择出训练验证集的索引
trainval=random.sample(list_index,tv)

# 从训练验证集的索引中随机选择出训练集的索引
train=random.sample(trainval,tr)

# 打开要写入的训练验证集、测试集、训练集、验证集的txt文件
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')

# 遍历所有XML文件的索引
for i in list_index:
    name=total_xml[i][:-4]+'\n' # 获取XML文件的名称（去掉后缀.xml）
    # 如果该索引在训练验证集中
    if i in trainval:
        file_trainval.write(name) # 写入训练验证集txt文件
        if i in train:  # 如果该索引在训练集中
            file_train.write(name)  # 写入训练集txt文件
        else:
            file_val.write(name)  # 写入验证集txt文件
    else:
        file_test.write(name)  # 否则写入测试集txt文件

# 关闭所有打开的文件
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

运行 split train_val.py文件后将得到 train.txt、val.txt、test.txt 和 trainval.txt 四个文件，结果如下：

注:如何修改数据集中训练集、验证集、测试集的比例?请参考:

三、生成 train.txt、test.txt、val.txt 文件

先看看我们要生成的文件位置：

现在我们需要的文件是voc_label.py，其位置如下：

voc_label.py文件的代码如下所示：

#-*- coding:utf-8 -*-

#导入必要的库
import xml.etree.ElementTree as ET
import os
from os import getcwd

#定义数据集的名称
sets =['train','val','test']

#请根据您的数据集修改这些类别名称
# 下面的两行代码，列表的里面的单词，除了首字母的大小写不一样，其他的是一样的，如果运行首字母是大写的代码，那么labels文件夹的txt文件的内容是空的，如果运行首字母是小写的代码，那么labels文件夹的txt文件的内容是有的，这很神奇，不知道原因是什么。
# classes =['Banana','Snake fruit','Dragon fruit','Pineapple']
classes = ["banana","snake fruit","dragon fruit","pineapple"]


#获取当前工作目录的绝对路径
abs_path = os.getcwd()
print(abs_path)

#定义一个函数，将边界框的坐标从绝对值转换为相对于图像大小的比例
def convert(size,box):
    dw = 1./(size[0])     #计算图像宽度的倒数
    dh = 1./(size[1])     #计算图像高度的倒数
    x = (box[0]+box[1])/2.0-1     #计算中心点的x坐标
    y = (box[2]+box[3])/2.0-1     # 计算中心点的y坐标
    w= box[1]- box[0]           # 计算边界框的宽度
    h = box[3] - box[2]         # 计算边界框的高度
    x = x * dw  # 缩放x坐标
    w = w * dw  # 缩放宽度
    y = y * dh  # 缩放y坐标
    h=h*dh      #缩放高度
    return x, y, w, h

# 定义一个函数，将标注文件从XML格式转换为YOLO格式
def convert_annotation(image_id):
    in_file = open('./Annotations/%s.xml' % (image_id),encoding='UTF-8')  # 打开XML标注文件
    out_file = open('./labels/%s.txt'%(image_id),'w')   # 打开要写入的YOL0格式标签文件
    tree =ET.parse(in_file)     #解析XML文件
    root = tree.getroot()

    filename = root.find('filename').text  # 获取图像文件名
    filenameFormat=filename.split(".")[1]  # 获取文件格式

    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)     #调用convert函数，将边界框坐标转换为Y0OLO格式
        out_file.write(str(cls_id)+" "+" ".join([str(a) for a in bb])+'\n')     #写入YOLO格式标签文件

    return filenameFormat


#获取当前工作目录
wd = getcwd()

#遍历每个数据集(train、val、test)
for image_set in sets:
    #如果labels目录不存在，创建它
    if not os.path.exists('./labels/'):
        os.makedirs('./labels/')
    #从数据集文件中获取图像ID列表
    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:
        filenameFormat=convert_annotation(image_id)
        list_file.write(abs_path +'/images/%s.%s\n'%(image_id,filenameFormat)) #注意你的图片格式，如果是.jpg记得修改
    list_file.close()

运行voc_label.py文件，你将会得到上面截图中 train.txt、test.txt、val.txt三个文件，文件内容如下：
train.txt的截图：

val.txt的截图：

四、创建 ab.yaml 文件

这个文件名是可以随意取的，是可以改变的，ab.yaml文件的位置如下：

ab.yaml文件的内容如下：

#path: E:\YOLO\yolov5-master    #dataset root dir
train: E:\YOLO\yolov5-master\paper_data\train.txt   #train images (relative to 'path') 118287 images
val: E:\YOLO\yolov5-master\paper_data\val.txt   #train images (relative to 'path') 5000 images
#上面的路径要使用绝对路径
#test: test-dev2017.txt

nc: 4    # number of classes，根据类别数目实际填写，本次训练的类别有4种

names: ["banana", "snake fruit", "dragon fruit", "pineapple"]    #改成自己的类别

五、开始用自己的数据集训练模型

在yolov5-master文件夹下打开cmd，查看自己的环境：

conda env list

进入我自己的pytorch环境中：

conda activate torch_env

输入命令：

python E:/YOLO/yolov5-master/train.py --img 900 --batch 2 --epoch 100 --data E:/YOLO/yolov5-master/data/ab.yaml --cfg E:/YOLO/yolov5-master/models/yolov5s.yaml --weights E:/YOLO/yolov5-master/yolov5s.pt

就可以直接训练自己的数据集啦，训练结果如下所示：

从图中可以看出：

在本项目的100轮次的训练中，总计费时0.925小时，大体55分钟。

YOLOv5s的网络结构有157层，参数量为7020913，梯度为0，GFLOPs是15.8。

还显示了类别的训练结果，包括P-R值、mAP50的值。

四种水果的模拟精度以及总体的模拟精度。

训练结果存储在 runs\train\exp8这个文件夹中。

六、遇到的问题

1. export GIT_PYTHON_REFRESH=quiet

初次运行时出现“export GIT_PYTHON_REFRESH=quiet”的报错。

这个环境变量是Git的一个配置选项，用于控制Git在执行某些操作时是否刷新Python环境。当设置为 quiet 时，Git在执行操作时不会打印有关刷新Python环境的日志信息。

可以进入train.py内进行修改，添加代码即可

os.environ["GIT_PYTHON_REFRESH"] = "quiet"

如图所示：

2. Arial.ttf 字体问题

训练开始时出现错误，提示下载Arial.ttf 字体失败，连接超时。直接从网上下载后按照提示的地址的复制粘贴进去即可。

链接：https://pan.baidu.com/s/1owG5tOmi2a2TGwBShrq4bw
提取码：q9md

七、训练结果

打开对应的文件夹可以看到训练结果大体情况：

抽取其中几张演示参考：

八、心得体会

初步在YOLOv5上运行了自己的数据集，如果单纯右键运行train.py文件则会自动下载并运行YOLOv5的数据集，如果想运行自己数据集则需要在train.py文件中进行修改。