YOLOv5训练自己的数据集

1.源代码获取

源代码的获取教程在我的另一篇博客中有提到，这里就不多赘述了，给个传送门

2.数据集获取

这里我提供了一个安全帽检测项目的数据集，大家可以自行下载，当然网上也有很多开源的数据集供大家下载尝试，这里给个安全帽数据集的传送门。

获取数据集后，将其解压到yolov5项目下的主目录中。

3.voc数据集转换为txt格式

由于voc数据集中的标签格式都是xml文件，不是项目所能读取的数据，因此我们需要先将其全部转成yolo能读懂的标签格式（txt格式）。
这里我们可以在我们项目的主目录下创建一个py文件，这里我命名为convert.py文件：

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


def convert(size, box):
    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]
    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]
    return (x, y, w, h)


def convert_annotation(xml_files_path, save_txt_files_path, classes):
    xml_files = os.listdir(xml_files_path)
    print(xml_files)
    for xml_name in xml_files:
        print(xml_name)
        xml_file = os.path.join(xml_files_path, xml_name)
        out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')
        out_txt_f = open(out_txt_path, 'w')
        tree = ET.parse(xml_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))
            # b=(xmin, xmax, ymin, ymax)
            print(w, h, b)
            bb = convert((w, h), b)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


if __name__ == "__main__":
    # 需要转换的类别，需要一一对应
    classes1 = ['hat', 'person']
    # 2、voc格式的xml标签文件路径
    xml_files1 = r'D:\detection\yolov5-5.0\yolov5-5.0\VOC2028\VOC2028\Annotations'
    # 3、转化为yolo格式的txt标签文件存储路径
    save_txt_files1 = r'D:\detection\yolov5-5.0\yolov5-5.0\dataset\annotations'

    convert_annotation(xml_files1, save_txt_files1, classes1)

这里main函数中的路径需要根据实际情况进行修改，看你数据的存放的位置，输入绝对路径即可。这里我将voc的所有标签数据放到了主路径下的一个dataset文件中的annotation文件夹中，方便后面进一步划分训练集和验证集和测试集。这里我将所有的图片数据集也转移到了dataset下面的image文件夹中：

4.制作训练集、验证集等

这里用到一个实现数据抽取功能的函数，我将其命名为divide.py文件：

import os, random, shutil


def moveimg(fileDir, tarDir):
    pathDir = os.listdir(fileDir)  # 取图片的原始路径
    filenumber = len(pathDir)
    rate = 0.72  # 自定义抽取图片的比例，比方说100张抽10张，那就是0.1
    picknumber = int(filenumber * rate)  # 按照rate比例从文件夹中取一定数量图片
    sample = random.sample(pathDir, picknumber)  # 随机选取picknumber数量的样本图片
    print(sample)
    for name in sample:
        shutil.move(fileDir + name, tarDir + "\\" + name)
    return


def movelabel(file_list, file_label_train, file_label_val):
    for i in file_list:
        if i.endswith('.jpg'):
            # filename = file_label_train + "\\" + i[:-4] + '.xml'  # 可以改成xml文件将’.txt‘改成'.xml'就可以了
            filename = file_label_train + "\\" + i[:-4] + '.txt'  # 可以改成xml文件将’.txt‘改成'.xml'就可以了
            if os.path.exists(filename):
                shutil.move(filename, file_label_val)
                print(i + "处理成功！")


if __name__ == '__main__':
    fileDir = r"D:\detection\yolov5-5.0\yolov5-5.0\dataset\images" + "\\"  # 源图片文件夹路径
    tarDir = r'D:\detection\yolov5-5.0\yolov5-5.0\VOCdevkit\images\train'  # 图片移动到新的文件夹路径
    moveimg(fileDir, tarDir)
    file_list = os.listdir(tarDir)
    file_label_train = r"D:\detection\yolov5-5.0\yolov5-5.0\dataset\annotations"  # 源图片标签路径
    file_label_val = r"D:\detection\yolov5-5.0\yolov5-5.0\VOCdevkit\labels\train"  # 移动标签的路径
    # 移动到新的文件路径
    movelabel(file_list, file_label_train, file_label_val)

main函数中的路径根据自己的需要设置。这里为了减少路径的更改配置，我都按照统一标准，在主路径下创建了一个VOCdevkit的文件夹：

在代码中，moveimg方法中，rate是从源数据集中随机获得数据的比例，可以自行设置来制作不同数量的数据集。

5.开始训练

这一部分，炮哥带你学的博客中写的十分清楚到位，跟着大佬的来做就可以了。

6.我遇到的bug以及解决方式

除了第五部分中，炮哥所讲解的训练前的一些设置，我自己本身还遇到了一些额外的问题，这里我将一一展示，并提供解决方法。
首先第一个bug就是：

AttributeError: Can't get attribute 'SPPF' on <module 'models.common' from 'D:\\detection\\yolov5-5.0\\yolov5-5.0\\models\\common.py'>

解决方式：在models/common.py中增加一个类：

import warnings

class SPPF(nn.Module):
    # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
    def __init__(self, c1, c2, k=5):  # equivalent to SPP(k=(5, 9, 13))
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * 4, c2, 1, 1)
        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
            y1 = self.m(x)
            y2 = self.m(y1)
            return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))

第二个bug就是：

RuntimeError: result type Float can't be cast to the desired output type __int64

解决方式：
第一步，先打开utils中的loss.py文件，使用cltr+f快捷键搜索：

for i in range(self.nl):

会在177行找到对应代码，将其改成：

        for i in range(self.nl):
            anchors, shape = self.anchors[i],p[i].shape
            gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]]  # xyxy gain

第二步，还是打开utils中的loss.py文件，使用cltr+f快捷键搜索：

indices.append

在第211行找到需要修改的代码，将其改为：

indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1))) 

到此为止，不出意外的话，应该是可以训练起来了，我们就能看到终端界面的网络层和训练进度：

7.对图像的检测和对视频的检测

这里同样可以参照炮哥带你学中的预测部分，不得不说炮哥讲的真的非常好！