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 osclass 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 + 1except:continueprint ('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 randomtrainval_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 getcwdsets = ['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 - 1y = (box[2] + box[3]) / 2.0 - 1w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn x, y, w, hdef 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').textif cls not in classes :continuecls_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 = bif b2 > w:b2 = wif b4 > h:b4 = hb = (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.txtnc: 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数据集的训练输出日志流程差不多，没有报错，这一步就完成了。