Python实现自动分离coco128数据集，按比例分为训练验证测试三部分

import os
import random
import shutil

# 源数据集文件夹路径
source_dir = 'E:\\Folder\\dataset\\coco128'

# 目标数据集文件夹保存路径
target_dir = 'E:\\Folder\\dataset\\coco128_split'

# 创建目标数据集文件夹
os.makedirs(target_dir, exist_ok=True)

# 创建train、valid和test子文件夹
subsets = ['train', 'valid', 'test']
for subset in subsets:
	subset_dir = os.path.join(target_dir, subset)
	os.makedirs(subset_dir, exist_ok=True)
	os.makedirs(os.path.join(subset_dir, 'images'), exist_ok=True)
	os.makedirs(os.path.join(subset_dir, 'labels'), exist_ok=True)

# 列出源数据集中的所有文件
all_files = os.listdir(os.path.join(source_dir, 'images\\train2017'))
random.shuffle(all_files)  # 随机打乱文件列表

# 计算划分数据的大小
total_files = len(all_files)
train_size = int(0.7 * total_files)
valid_size = int(0.2 * total_files)

# 分离数据并复制到目标文件夹
for i, file in enumerate(all_files):
	src_image = os.path.join(source_dir, 'images\\train2017', file)
	src_label = os.path.join(source_dir, 'labels\\train2017', file.replace('.jpg', '.txt'))

	if not os.path.exists(src_image) or not os.path.exists(src_label):
		# 如果文件不存在，跳过当前循环
		continue

	if i < train_size:
		dst_subset = 'train'
	elif i < train_size + valid_size:
		dst_subset = 'valid'
	else:
		dst_subset = 'test'

	dst_image = os.path.join(target_dir, dst_subset, 'images', file)
	dst_label = os.path.join(target_dir, dst_subset, 'labels', file.replace('.jpg', '.txt'))

	shutil.copy(src_image, dst_image)
	shutil.copy(src_label, dst_label)

print("数据集分离完成。")