按7:2:1比例随机划分yolo数据集

import os
import shutil
import random# 设置路径
base_dir = 'D:\ChromeDownload\***'
images_dir = os.path.join(base_dir, 'images')
labels_dir = os.path.join(base_dir, 'labels')# 输出路径
output_dir = 'D:\ChromeDownload\***'
splits = ['train', 'val', 'test']# 创建输出目录
for split in splits:os.makedirs(os.path.join(output_dir, split, 'images'), exist_ok=True)os.makedirs(os.path.join(output_dir, split, 'labels'), exist_ok=True)# 获取所有图片名（去掉扩展名）
image_files = sorted([f for f in os.listdir(images_dir) if f.endswith('.jpg') or f.endswith('.png')])
base_names = [os.path.splitext(f)[0] for f in image_files]# 打乱并划分
random.shuffle(base_names)
total = len(base_names)
train_end = int(0.7 * total)
val_end = int(0.9 * total)train_files = base_names[:train_end] #0-7
val_files = base_names[train_end:val_end] #7-9
test_files = base_names[val_end:] #9~endprint(f'Total: {total}, Train: {len(train_files)}, Val: {len(val_files)}, Test: {len(test_files)}')# 拷贝函数
def copy_files(file_list, split):for name in file_list:img_src = os.path.join(images_dir, name + '.jpg')  # 假设是.jpgif not os.path.exists(img_src):  # 或者是.pngimg_src = os.path.join(images_dir, name + '.png')label_src = os.path.join(labels_dir, name + '.txt')img_dst = os.path.join(output_dir, split, 'images', os.path.basename(img_src))label_dst = os.path.join(output_dir, split, 'labels', os.path.basename(label_src))shutil.copy(img_src, img_dst)if os.path.exists(label_src):shutil.copy(label_src, label_dst)# 执行复制
copy_files(train_files, 'train')
copy_files(val_files, 'val')
copy_files(test_files, 'test')