自动划分数据集脚本

自动将数据集划分为测试集、验证集、测试集（1 - train - val）

文件夹格式：

source_dir/
|
|-- class1/
|   |-- 1.png
|   |-- 2.png
|   |-- ...
|   |
|-- class2/
|   |-- 1.png
|   |-- 2.png
|   |-- ...
|
|-- class3/
|   |-- 1.png
|   |-- 2.png
|   |-- ...
|-- ...

脚本：

import os       # 操作文件和文件夹路径
import random   # 用于打乱文件顺序
import shutil   # 用于复制文件

def split_dataset(source_dir, dest_dir, train_ratio=0.8, val_ratio=0.1):
    """
    按指定比例划分图像数据集为训练集、验证集和测试集。
    
    参数：
        source_dir: 原始数据集路径，要求每个类别是一个子文件夹
        dest_dir: 输出路径，将创建 train/val/test 子目录
        train_ratio: 训练集所占比例（默认0.8）
        val_ratio: 验证集所占比例（默认0.1），测试集 = 1 - train - val
    """
    # 获取所有类别名称（即 source_dir 下的子文件夹）
    class_names = os.listdir(source_dir)

    for class_name in class_names:
        # 获取每个类别下的所有文件路径（例如图片）
        files = os.listdir(os.path.join(source_dir, class_name))
        
        # 打乱文件顺序，确保划分是随机的
        random.shuffle(files)

        # 计算训练集和验证集的分界点
        train_split = int(len(files) * train_ratio)
        val_split = int(len(files) * (train_ratio + val_ratio))

        # 划分文件列表
        subsets = {
            'train': files[:train_split],               # 前 train_ratio 的文件
            'val': files[train_split:val_split],        # 中间 val_ratio 的文件
            'test': files[val_split:]                   # 剩下的作为测试集
        }

        # 复制文件到对应的目标路径
        for subset, subset_files in subsets.items():
            # 构建目标子文件夹路径，如：dest_dir/train/class_name/
            subset_dir = os.path.join(dest_dir, subset, class_name)
            
            # 创建子文件夹（如果不存在）
            os.makedirs(subset_dir, exist_ok=True)
            
            for f in subset_files:
                # 构建原始文件路径和目标路径
                src_path = os.path.join(source_dir, class_name, f)
                dst_path = os.path.join(subset_dir, f)
                
                # 复制文件
                shutil.copy(src_path, dst_path)

# 示例调用：
# split_dataset('data/original', 'data/split')
# 假设 data/original 中是按类别分文件夹的结构，调用后将在 data/split 下创建 train/val/test 三个子文件夹，每个内部保留类别结构。