YOLOv8训练前准备2（数据集--含测试集）

上一篇文章基本确定了YOLOv8模型是可以使用的，那么下一步就开始进入炼丹环节了。在此之前还需要确定自己的数据集，要按照YOLOv8的格式来。（预训练模型）

注意：这篇文章不涉及数据集的制作，包含的是数据集的分类，数据集标签的转化（部分思路），所以阅读之前请先保证自己拿到了数据集 。

1：数据集格式转化

YOLOv8数据集的训练格式与YOLOv5是一样的，主要包括训练图片与图片标签（.txt格式）。常见的标签包括.txt .xml以及.json ，所以要先将非.txt的图片标签转化为.txt的标签。这篇文章主要说明xml格式转txt格式的方法，直接上代码。（注意：其他的转化思路是一样的，就是要找到合适的代码段）

需要调整的就是if __name__ == "__main__":函数里面的内容，调整为自己的就可以了。

#xml 转yolo
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 = ["0","1","2","3"]  # 这里需要改成自己的对应类别。
    # xml标签文件路径（绝对路径）
    xml_files1 = r'D:\\xx\\label_xlm'
    # 转化为txt标签文件存储路径
    save_txt_files1 = r'D:\\xx\\label_txt'
 
    convert_annotation(xml_files1, save_txt_files1, classes1)

转换完之后大概是这个样子，我们使用的是image与label_txt两个文件夹 。

2：数据集的分类 （训练集+验证集+测试集）

在拿到一个新的数据集时，一般包括图片和标签内容，如下图所示，images里面放的是所有的训练图片，labels里面放的是所有的图片标签。如果你的不是这样的，请调整一下，后面需要整体划分数据集，需要这样的格式。即上面的前两个文件夹。

调整成这样的文件夹之后，就可以直接上代码了。

# by CSDN 迪菲赫尔曼
import os
import random
import shutil

def copy_files(src_dir, dst_dir, filenames, extension):
    os.makedirs(dst_dir, exist_ok=True)
    missing_files = 0
    for filename in filenames:
        src_path = os.path.join(src_dir, filename + extension)
        dst_path = os.path.join(dst_dir, filename + extension)
        
        # Check if the file exists before copying
        if os.path.exists(src_path):
            shutil.copy(src_path, dst_path)
        else:
            print(f"Warning: File not found for {filename}")
            missing_files += 1

    return missing_files

#这里用来设置训练集，验证集，测试集的占比，这里是8：1：1，可以进行调整的。
def split_and_copy_dataset(image_dir, label_dir, output_dir, train_ratio=0.8, valid_ratio=0.1, test_ratio=0.1):
    # 获取所有图像文件的文件名（不包括文件扩展名）
    image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_dir)]

    # 随机打乱文件名列表
    random.shuffle(image_filenames)

    # 计算训练集、验证集和测试集的数量
    total_count = len(image_filenames)
    train_count = int(total_count * train_ratio)
    valid_count = int(total_count * valid_ratio)
    test_count = total_count - train_count - valid_count

    # 定义输出文件夹路径
    train_image_dir = os.path.join(output_dir, 'train', 'images')
    train_label_dir = os.path.join(output_dir, 'train', 'labels')
    valid_image_dir = os.path.join(output_dir, 'valid', 'images')
    valid_label_dir = os.path.join(output_dir, 'valid', 'labels')
    test_image_dir = os.path.join(output_dir, 'test', 'images')
    test_label_dir = os.path.join(output_dir, 'test', 'labels')

    # 复制图像和标签文件到对应的文件夹
    train_missing_files = copy_files(image_dir, train_image_dir, image_filenames[:train_count], '.jpg')
    train_missing_files += copy_files(label_dir, train_label_dir, image_filenames[:train_count], '.txt')

    valid_missing_files = copy_files(image_dir, valid_image_dir, image_filenames[train_count:train_count + valid_count], '.jpg')
    valid_missing_files += copy_files(label_dir, valid_label_dir, image_filenames[train_count:train_count + valid_count], '.txt')

    test_missing_files = copy_files(image_dir, test_image_dir, image_filenames[train_count + valid_count:], '.jpg')
    test_missing_files += copy_files(label_dir, test_label_dir, image_filenames[train_count + valid_count:], '.txt')

    # Print the count of each dataset
    print(f"Train dataset count: {train_count}, Missing files: {train_missing_files}")
    print(f"Validation dataset count: {valid_count}, Missing files: {valid_missing_files}")
    print(f"Test dataset count: {test_count}, Missing files: {test_missing_files}")

# 使用例子 第一行输入的是你的图片路径（绝对路径）；第二行输入的是你的标签路径（绝对路径）；第三行输入的是你期望输出的文件夹
image_dir = 'D:\\xx\\image'
label_dir = 'D:\\xx\\label_txt'
output_dir = 'D:\\xx\\my_data'

split_and_copy_dataset(image_dir, label_dir, output_dir)

运行完之后会生成如下的文件夹。

每个文件夹又包含两个文件夹，images和labels。检查一下是不是这样的，看一下图片数量与标签数量是不是一致的。如果不是，请重新跑一下代码（注意：重新跑代码之前，请重新建立输出的文件夹或者清空输出的文件夹，否则会叠加，从而生成必错的分类结果）。