分类数据集拆分成train和val两个子集

import os
import random
import shutil

def split_dataset(src_dir, train_dir, test_dir, ratio):
    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    if not os.path.exists(test_dir):
        os.makedirs(test_dir)

    categories = os.listdir(src_dir)
    for category in categories:
        category_dir = os.path.join(src_dir, category)
        if os.path.isdir(category_dir):
            train_category_dir = os.path.join(train_dir, category)
            test_category_dir = os.path.join(test_dir, category)
            if not os.path.exists(train_category_dir):
                os.makedirs(train_category_dir)
            if not os.path.exists(test_category_dir):
                os.makedirs(test_category_dir)

            filenames = os.listdir(category_dir)
            random.shuffle(filenames)
            split_index = int(len(filenames) * ratio)
            train_filenames = filenames[:split_index]
            test_filenames = filenames[split_index:]

            for filename in train_filenames:
                src_file = os.path.join(category_dir, filename)
                dst_file = os.path.join(train_category_dir, filename)
                shutil.copy(src_file, dst_file)

            for filename in test_filenames:
                src_file = os.path.join(category_dir, filename)
                dst_file = os.path.join(test_category_dir, filename)
                shutil.copy(src_file, dst_file)

#将src_dir目录下的数据集按3比7的比例拆分到train_dir和test_dir目录下
split_dataset('src_dir', 'train_dir', 'test_dir', 0.7)

请确保将src_dir替换成你实际的数据集目录，train_dir为训练集保存的目录，test_dir为测试集保存的目录。0.7表示拆分比例为70%的数据用于训练集，30%的数据用于测试集。