YOLOv8训练自定义数据集

首先使用labelImg对数据进行标记

我的图像数据包含不同的来源，比如文件夹A、B、C，其中图像的命名方式都是image_i.jpg …
且图像命名没有重叠的情况（标记文件同理），在整理训练数据集的时候需要修改命名规则来保证文件能够放到一个文件夹下，可以参考如下的重命名的脚本：

import os
import shutil

# 设置原始文件夹路径
jpeg_images_dir = '/home/robot/Desktop/temp/JPEGImages'  # 修改为JPEGImages文件夹的实际路径
annotations_dir = '/home/robot/Desktop/temp/Annotations'  # 修改为Annotations文件夹的实际路径

# 设置新文件夹路径
new_images_dir = '/home/robot/Desktop/temp/JPEGImages'  # 修改为新JPEGImages文件夹的实际路径
new_annotations_dir = '/home/robot/Desktop/temp/Annotations'  # 修改为新Annotations文件夹的实际路径

# 设置起始编号
start_number = 0  # 或任何你希望开始的数字

# 创建新文件夹，如果它们还不存在的话
os.makedirs(new_images_dir, exist_ok=True)
os.makedirs(new_annotations_dir, exist_ok=True)

# 获取所有.bmp文件的文件名（无扩展名）
jpeg_files = [f for f in os.listdir(jpeg_images_dir) if f.endswith('.bmp')]

# 按文件名排序，并对每个文件进行编号和重命名
for index, old_file in enumerate(sorted(jpeg_files), start=start_number):
    # 获取没有扩展名的文件名
    base_filename = os.path.splitext(old_file)[0]
    
    # 路径设置
    old_image_path = os.path.join(jpeg_images_dir, old_file)
    new_image_name = f"frame_{index}.bmp"
    new_image_path = os.path.join(new_images_dir, new_image_name)
    
    old_annotation_file = f"{base_filename}.xml"
    old_annotation_path = os.path.join(annotations_dir, old_annotation_file)
    new_annotation_name = f"frame_{index}.xml"
    new_annotation_path = os.path.join(new_annotations_dir, new_annotation_name)
    
    # 仅当Annotations文件夹中存在对应的.xml文件时执行重命名
    if os.path.exists(old_annotation_path):
        # 移动并重命名JPEGImages中的.bmp文件
        shutil.move(old_image_path, new_image_path)
        # 移动并重命名Annotations中的.xml文件
        shutil.move(old_annotation_path, new_annotation_path)

# 执行完毕后，新文件夹中应该包含了以指定起始号码开始的重命名文件

开始制作yolov8的训练验证数据

到这步，文件夹下应该包含了Annotations 、ImageSets、JPEGImages，ImageSets文件夹中创建一个Main文件夹。
然后进行训练数据和验证数据的划分：splitTrainValTest.py

import os
import random

random.seed(0)

# 路径
xmlfilepath = r'./Annotations/'
saveBasePath = r'./ImageSets/Main/'

if not os.path.exists(saveBasePath):
    os.makedirs(saveBasePath)

trainval_percent = 0.95    #训练验证集占整个数据集的比重（划分训练验证集和测试集）
train_percent = 0.95 #训练集占整个训练验证集的比重（划分训练集和验证集）

temp_xml = os.listdir(xmlfilepath)
total_xml = []
for xml in temp_xml:
    if xml.endswith(".xml"):
        total_xml.append(xml)

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)

print("train and val size", tv)
print("train size", tr)
ftrainval = open(os.path.join(saveBasePath, 'trainval.txt'), 'w')
ftest = open(os.path.join(saveBasePath, 'test.txt'), 'w')
ftrain = open(os.path.join(saveBasePath, 'train.txt'), 'w')
fval = open(os.path.join(saveBasePath, '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()

紧接着，进行数据集生成：generateYoloDataset.py

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

datasets_path = "."
annotations_path = "{}/Annotations".format(datasets_path)
imagesets_path = "{}/ImageSets/Main".format(datasets_path)
jpegimages_path = "{}/JPEGImages".format(datasets_path)

sets = [('TrainVal', 'train'), ('TrainVal', 'val'), ('Test', 'test')]

classes = ["fluid", "dingzi"]

def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    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(mode, image_set, image_id):
    in_file = open("{path}/{image_id}.xml".format(path = annotations_path, image_id = image_id))
    out_file = open("{path}/labels/{mode}_{image_set}/{image_id}.txt".format(path = datasets_path, mode = mode, image_set = image_set, image_id = image_id), 'w', encoding = 'utf-8')
    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))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    in_file.close()
    out_file.close()
    
def copy_images(mode, iamge_set, image_id):
    in_file = "{path}/{image_id}.bmp".format(path = jpegimages_path, image_id = image_id)
    out_file = "{path}/images/{mode}_{iamge_set}/{image_id}.bmp".format(path = datasets_path, mode = mode, iamge_set = iamge_set, image_id = image_id)
    shutil.copy(in_file, out_file)

for mode, image_set in sets:
    label_dir = "{path}/labels/{mode}_{image_set}".format(path = datasets_path, mode = mode, image_set = image_set)
    if not os.path.exists(label_dir):
        os.makedirs(label_dir)
    
    image_dir = "{path}/images/{mode}_{image_set}".format(path = datasets_path, mode = mode, image_set = image_set)
    if not os.path.exists(image_dir):
        os.makedirs(image_dir)
    
    image_ids = open("{path}/{image_set}.txt".format(path = imagesets_path, image_set = image_set)).read().strip().split()
    list_file = open("{path}/{mode}_{image_set}".format(path = datasets_path, mode = mode, image_set = image_set), 'w')

    for iamge_id in image_ids:
        list_file.write("./images/{mode}_{image_set}/{id}.bmp\n".format(mode = mode, image_set = image_set, id = iamge_id))
        convert_annotation(mode, image_set, iamge_id)
        copy_images(mode, image_set, iamge_id)
    
    list_file.close()

至此，训练集生成操作完成。开始准备训练的环境，这部分我就跳过。

训练集最好放入一个自己创建的文件夹中：比如cup/ 然后将该文件夹放入自己创建好的工程文件夹下。

再创建一个data文件夹，文件夹下创建一个训练集的yaml文件：

path: /home/robot/spf/yoloSeries/yolov8/cup
train: TrainVal_train
val: TrainVal_val
test: Test_test

names:
  0: cup

然后在工程文件夹下放入比如 yolov8n.pt的预训练模型

训练

举个例子：进入自己的工程文件夹路径下，然后启动训练

cd /mnt/d/Code/OpenSource/YOLO/yolov8
yolo detect train model=yolov8n.pt data=./data/fruit_box.yaml epochs=100 imgsz=512 patience=0 batch=32 save=True workers=8

https://docs.ultralytics.com/zh/usage/cfg/#train-settings 该链接包含了CLI的不同选项的具体含义

验证数据集

yolo detect val model= ./runs/detect/train/weights/best.pt

使用训练验证集之外的照片进行测试

yolo predict model= ./runs/detect/train/weights/best.pt source= ./images/

pt文件转wts文件，序列化wts文件为engine文件，可部署在jetson平台

参考yolov8_tensorrt
把这个工程拉到本地，根据readme来进行配置