如何在python中将图像数据集拆分为测试/训练/验证集？

import os
import random
import shutil

dataset_dir = 'path/to/dataset'  # 数据集目录
train_ratio = 0.7  # 训练集比例
val_ratio = 0.15  # 验证集比例
test_ratio = 0.15  # 测试集比例

image_files = os.listdir(dataset_dir)
random.shuffle(image_files)  # 随机打乱文件列表顺序

total_images = len(image_files)
train_size = int(total_images * train_ratio)
val_size = int(total_images * val_ratio)
test_size = total_images - train_size - val_size

train_dir = 'path/to/train'  # 训练集目录
val_dir = 'path/to/validation'  # 验证集目录
test_dir = 'path/to/test'  # 测试集目录

os.makedirs(train_dir, exist_ok=True)
os.makedirs(val_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)

train_files = image_files[:train_size]
val_files = image_files[train_size:train_size+val_size]
test_files = image_files[train_size+val_size:]

for file in train_files:
    shutil.copy(os.path.join(dataset_dir, file), train_dir)

for file in val_files:
    shutil.copy(os.path.join(dataset_dir, file), val_dir)

for file in test_files:
    shutil.copy(os.path.join(dataset_dir, file), test_dir)