Yolov8_使用自定义数据集训练模型1

前面几篇文章介绍了如何搭建Yolov8环境、使用默认的模型训练和推理图片及视频的效果、并使用GPU版本的torch加速推理、导出.engine格式的模型进一步利用GPU加速，本篇介绍如何自定义数据集，这样就可以训练出识别特定物体的模型。

《Yolov8_使用自定义数据集训练模型1》——主要是怎么创建自定义数据集，测试demo；

《Yolov8_使用自定义数据集训练模型2》——搜集更多的图片去标注、训练，重点关注训练后的实际效果；

1、创建自定义数据集

1.1、创建自定义数据集——总体流程

收集图片：收集一批带有目标物体的图片【images文件夹下.png图片】
标注目标物体：使用标注工具对图片中的目标物体进行标注【xml_labels文件夹下.xml文件】
划分数据集：将整个数据集按一定比例分为训练集、验证集、测试集【使用split_dataset.py脚本生成split_dataset_txt文件夹中的.txt文件，.txt文件内容是不含后缀.xml的文件名】
生成Yolo标注文件及各数据集使用的图片路径：使用xml_to_txt.py脚本将xml标注文件转成Yolo需要的.txt标注文件【labels文件夹下.txt标注文件】，同时脚本生成训练集、验证集、测试集所使用图片的绝对路径【当前目录下test.txt、train.txt、val.txt】

最终的效果就是下面这个文件夹：

farmland.yaml是进行yolo训练时配置文件，不属于创建数据集，后面再说明；

图片是截图来的命名比较乱，batch_rename.py用于批量重命名，下面会附代码；

1.2、收集图片

收集一批带有目标物体的图片，图片的多少和质量关乎训练出模型的效果，这里只找了几张图片是为了跑一下自定义数据集的流程。

图片是截图来的，自动保存的文件名与内容无关，相信你也不想一个一个rename，下面是对图片批量重命名的batch_rename.py代码：

import os

class BatchRename():
    def __init__(self):
        self.path = './images'  
        
    def rename(self):
        filelist = os.listdir(self.path)    
        total_num = len(filelist)          
        i = 1 
        for item in filelist:
            if 1:  
                src = os.path.join(os.path.abspath(self.path), item)  
                dst = os.path.join(os.path.abspath(self.path), 'farmland' + format(str(i), '0>4s') + '.png')    
                try:
                    os.rename(src, dst)   
                    print('converting %s to %s ...' % (src, dst))
                    i = i + 1
                except:
                    continue
        print ('total rename %d files.' % (total_num))
 
if __name__ == '__main__':
    demo = BatchRename()
    demo.rename()

1.3、标注目标物体

1.3.1、标注工具_没有使用labelimg

大部分人使用的标注工具是labelimg，但是安装labelimg需要安装pyqt5等依赖，pyqt5等不支持python3.10，不至于为了这个标注工具去修改现在Linux的Python环境。虽然labelimg也支持Windows，但是看教程又需要Anaconda环境，没必要这么麻烦，所以不想使用labelimg。

1.3.2、标注工具_使用Colabeler

发现Colabeler的界面还算好看，支持计算机视觉、NLP、语音三大领域的标注，功能强大且免费，所以试试看。

Colabeler官网：Colabeler - Best annotation tool for AI dataset labeling

Windows安装Colabeler，标注后的.xml文件传给Linux

标注的目的是得到.xml文件，所以完全可以在Windows安装该软件并标注，然后使用SSH传给Linux，使用该方式是因为图片本身也要传给Linux，这样等于是把搜集图片和标注这两步合一起在自己的Windows上先做好。

Windows安装Colabeler没必要说了，直接下一步下一步就能安装成功，然后打开界面左上角创建项目，然后选择Localization、填入项目名称、图片路径、分类名称（多个分类用逗号隔开，这里只写了一个farmland，就是想根据图片判断是否是农田）。

使用“Rectangle”工具标注出目标物体，右侧LabelList选择目标所属类别，一张图片所有的目标物体标注完成后单击下方对号确认，使用“Next”进入下一张图片标注，所有图片标注完成单击“Export”导出XML文件。

1.4、划分数据集

目的是将整个数据集按一定比例分为训练集、验证集、测试集。

使用split_dataset.py脚本随机划分数据集，生成split_dataset_txt文件夹中的.txt文件，.txt文件内容是不含后缀.xml的文件名，下面是split_dataset.py以及各文件内容：

import os
import random
 
 
trainval_percent = 0.9
train_percent = 0.9
xmlfilepath = './xml_labels'
txtsavepath = './split_dataset_txt'
total_xml = os.listdir(xmlfilepath)
 
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
 
ftrainval = open('./split_dataset_txt/trainval.txt', 'w')
ftest = open('./split_dataset_txt/test.txt', 'w')
ftrain = open('./split_dataset_txt/train.txt', 'w')
fval = open('./split_dataset_txt/val.txt', 'w')
 
for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)
 
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

1.5、生成Yolo标注文件及各数据集使用的图片路径

使用xml_to_txt.py脚本将xml标注文件转成Yolo需要的.txt标注文件【labels文件夹下.txt标注文件】，同时脚本生成训练集、验证集、测试集所使用图片的绝对路径【当前目录下test.txt、train.txt、val.txt】，下面是xml_to_txt.py以及各文件内容：

如果使用的标注工具不同，解析xml过程可能会报错，此时任意可以打开一个.xml文件根据实际结构修改。

 
import xml.etree.ElementTree as ET
import os
from os import getcwd
 
sets = ['train', 'val', 'test']
classes = ['farmland']
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('./xml_labels/%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)
    
    objects = root.find('outputs').find('object')
      
    for obj in objects.iter('item'):      
        cls = obj.find('name').text
        if cls not in classes :
            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('./split_dataset_txt/%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))
        convert_annotation(image_id)
    list_file.close()

2、使用数据集进行yolo训练

2.1、编写.yaml配置文件

给出训练集、验证集、测试集的路径，训练的目标总数，具体的序号和目标名称列表。

train: /home/lgzn/datasets/farmland_dataset/train.txt
val: /home/lgzn/datasets/farmland_dataset/val.txt
test: /home/lgzn/datasets/farmland_dataset/test.txt

nc: 1
names:
  0: farmland

2.2、使用自定义数据集训练

修改.yaml配置文件的路径，测试能否使用刚才制作的这个数据集进行训练。

yolo train data='/home/lgzn/datasets/farmland_dataset/farmland.yaml' model=yolov8n.pt epochs=2 lr0=0.01

如果和之前文章中使用coco128数据集的训练输出日志流程差不多，没有报错，这一步就完成了。