← 返回首页
自定义数据集划分
发表时间:2025-04-13 11:25:55
自定义数据集划分

在训练YOLO(You Only Look Once)目标检测模型时,自定义数据集的划分是一个重要步骤。数据集的划分主要包括训练集、验证集和测试集的划分。这样做有助于模型在训练过程中进行有效的学习和评估。

1.获得原始数据集

在yolov8项目文件夹下创建存储原始数据集的目录,比如:origin_dataset(可以是其它名字),在这个目录下创建labels和images这两个文件夹。

2.自定义数据集划分(代码实现)

在yolov8项目文件夹下,创建split.py文件,用于存储划分数据的代码。

这里划分比例:训练集:验证集:测试集 = 7 : 1 :2。

split.py

import os
import random
import shutil
from pathlib import Path

#这里划分比例:训练集:验证集:测试集 = 7 : 1 :2。
def split_dataset(dataset_dir, output_dir, train_ratio=0.7, test_ratio=0.2, val_ratio=0.1):
    """
    自定义数据集划分函数,按照 YOLO 数据集格式划分数据集
    :param dataset_dir: 数据集的根目录(包含 images 和 labels 子文件夹)
    :param output_dir: 划分后的数据集保存目录
    :param train_ratio: 训练集占比
    :param test_ratio: 测试集占比
    :param val_ratio: 验证集占比
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 创建子目录结构
    train_images_dir = os.path.join(output_dir, "train", "images")
    train_labels_dir = os.path.join(output_dir, "train", "labels")
    val_images_dir = os.path.join(output_dir, "valid", "images")
    val_labels_dir = os.path.join(output_dir, "valid", "labels")
    test_images_dir = os.path.join(output_dir, "test", "images")
    test_labels_dir = os.path.join(output_dir, "test", "labels")

    os.makedirs(train_images_dir, exist_ok=True)
    os.makedirs(train_labels_dir, exist_ok=True)
    os.makedirs(val_images_dir, exist_ok=True)
    os.makedirs(val_labels_dir, exist_ok=True)
    os.makedirs(test_images_dir, exist_ok=True)
    os.makedirs(test_labels_dir, exist_ok=True)

    # 获取所有图像文件路径
    image_files = []
    for ext in ['.jpg', '.png', '.jpeg']:  # 支持的图像格式
        image_files.extend(list(Path(dataset_dir).rglob(f'images/*{ext}')))

    # 打乱文件列表
    random.shuffle(image_files)

    # 计算划分数量
    total_files = len(image_files)
    train_size = int(total_files * train_ratio)
    test_size = int(total_files * test_ratio)
    val_size = int(total_files * val_ratio)

    # 划分数据集
    train_files = image_files[:train_size]
    val_files = random.sample(train_files, val_size)  # 从训练集中随机抽取验证集
    test_files = image_files[train_size:train_size + test_size]

    # 复制文件到对应目录
    def copy_files(file_list, images_dir, labels_dir):
        for file in file_list:
            #print("file=>",file)
            # 复制图像文件
            shutil.copy(file, images_dir)
            # 获取对应的标签文件路径
            image_name = os.path.basename(file)
            image_name_without_ext = os.path.splitext(image_name)[0]
            label_file = os.path.join(dataset_dir, "labels", f"{image_name_without_ext}.txt")
            if os.path.exists(label_file):
                shutil.copy(label_file, labels_dir)
            else:
                print(f"警告:标签文件 {label_file} 不存在,跳过复制。")

    # 复制训练集文件
    #print("train_images_dir:",train_images_dir,"train_labels_dir:",train_labels_dir)
    copy_files(train_files, train_images_dir, train_labels_dir)

    print(f"训练集: {len(train_files)} 张图片")

    # 复制验证集文件
    copy_files(val_files, val_images_dir, val_labels_dir)
    print(f"验证集: {len(val_files)} 张图片")

    # 复制测试集文件
    copy_files(test_files, test_images_dir, test_labels_dir)
    print(f"测试集: {len(test_files)} 张图片")

    print(f"数据集划分完成!")

# 示例调用
dataset_dir = "./origin_dataset"  # 数据集的根目录(包含 images 和 labels 子文件夹)
output_dir = "./output/dataset"  # 划分后的数据集保存目录
split_dataset(dataset_dir, output_dir)

运行效果:

训练集: 2488 张图片
验证集: 355 张图片
测试集: 711 张图片
数据集划分完成!

运行后,划分后的数据集自动保存在了项目的output目录下,如下图所示: