YOLOv5解析 | 第二篇:用YOLOv5训练自己的数据集(超级详细完整版)

大家好,我是K同学啊!

我们接着上一篇文章 YOLOv5解析 | 第一篇:快速部署YOLOv5模型 配置完YOLOv5需要的环境后,今天我们试着用YOLOv5训练自己的数据。(在开始本教程前,记得先跑一遍入门篇,确保其他环境是正常的

有图有真相,先看看我昨天的运行结果

【YOLOv5 源码地址】

🚀 我的环境:

  • 语言环境:Python3.8
  • 编译器:PyCharm
  • 深度学习环境:
  • torch==1.10.0+cu113
  • torchvision==0.11.1+cu113
  • 显卡:GeForce RTX 3080
  • 文章目录

  • 一、准备好自己的数据
  • 二、运行 split_train_val.py 文件
  • 三、运行 voc_label.py 文件
  • 四、创建 ab.yaml 文件
  • 五、开始用自己的数据集训练模型
  • 一、准备好自己的数据

    我的目录结构是这样子的

  • 主目录
  • paper_data(自己创建一个文件夹,将数据放到这里)
  • Annotations(放置我们的.xml文件)
  • images(放置图片文件)
  • ImageSets
  • Main(会在该文件夹内自动生成 train.txt、val.txt、test.txt 和 trainval.txt 四个文件,存放训练集、验证集、测试集图片的名字)
  • 你将会看如下的目录结构:

    Annotations文件夹为xml文件,我的文件如下:


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

    二、运行 split_train_val.py 文件

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

    split_train_val.py 文件的位置如下:

    split_train_val.py 的内容如下:

    # coding:utf-8
    
    import os
    import random
    import argparse
    
    parser = argparse.ArgumentParser()
    #xml文件的地址,根据自己的数据进行修改 xml一般存放在Annotations下
    parser.add_argument('--xml_path', default='Annotations', type=str, help='input xml label path')
    #数据集的划分,地址选择自己数据下的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
    xmlfilepath = opt.xml_path
    txtsavepath = opt.txt_path
    total_xml = os.listdir(xmlfilepath)
    if not os.path.exists(txtsavepath):
        os.makedirs(txtsavepath)
    
    num = len(total_xml)
    list_index = range(num)
    tv = int(num * trainval_percent)
    tr = int(tv * train_percent)
    trainval = random.sample(list_index, tv)
    train = random.sample(trainval, tr)
    
    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')
    
    for i in list_index:
        name = total_xml[i][:-4] + '\n'
        if i in trainval:
            file_trainval.write(name)
            if i in train:
                file_train.write(name)
            else:
                file_val.write(name)
        else:
            file_test.write(name)
    
    file_trainval.close()
    file_train.close()
    file_val.close()
    file_test.close()
    

    运行 split_train_val.py 文件后你将得到 train.txt、val.txt、test.txt 和 trainval.txt 四个文件,结果如下:

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

    三、运行 voc_label.py 文件

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

    开始办事,现在我们需要的是 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']
    classes = ["unripe citrus"]   # 改成自己的类别
    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
        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 convert_annotation(image_id):
        in_file = open('./Annotations/%s.xml' % (image_id), encoding='UTF-8')
        out_file = open('./labels/%s.txt' % (image_id), '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'):
            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)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    
    wd = getcwd()
    for image_set in sets:
        if not os.path.exists('./labels/'):
            os.makedirs('./labels/')
        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:
            list_file.write(abs_path + '/images/%s.png\n' % (image_id)) # 注意你的图片格式,如果是.jpg记得修改
            convert_annotation(image_id)
        list_file.close()
    

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

    四、创建 ab.yaml 文件

    这个文件名是我随意取的,这个可以做出改变的哈

    ab.yaml 文件的位置如下:

    我的 ab.yaml 文件内容如下:

    #path: ../datasets/coco  # dataset root dir
    train: ./paper_data/train.txt  # train images (relative to 'path') 118287 images
    val: ./paper_data/val.txt  # train images (relative to 'path') 5000 images
    #test: test-dev2017.txt
    
    nc: 1  # number of classes
    
    names: ['unripe citrus'] # 改成自己的类别
    

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

    输入命令:

    python train.py --img 900 --batch 2 --epoch 100 --data data/ab.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --device '0'
    

    就可以直接训练我们自己的数据集啦,我最后的运行结果如下:

    如果你是采用pycharm右键训练的,记得将batch-size下调,否则可能会报yolov5训练时报错 OSError: [WinError 1455] 页面文件太小,无法完成操作。 Error loading的错误

    来源:K同学啊

    物联沃分享整理
    物联沃-IOTWORD物联网 » YOLOv5解析 | 第二篇:用YOLOv5训练自己的数据集(超级详细完整版)

    发表评论