如何分配训练集、验证集、测试集的标签和图片（YOLO系列）

preprocessing_second.py

import random
import os
import shutil

train_images_path = "C:/Users/DELL/Desktop/ultralytics/ultralytics/cfg/datasets/bvn/images/train"
train_labels_path = "C:/Users/DELL/Desktop/ultralytics/ultralytics/cfg/datasets/bvn/labels/train"
valid_images_path = "C:/Users/DELL/Desktop/ultralytics/ultralytics/cfg/datasets/bvn/images/val"
valid_labels_path = "C:/Users/DELL/Desktop/ultralytics/ultralytics/cfg/datasets/bvn/labels/val"
test_images_path = "C:/Users/DELL/Desktop/ultralytics/ultralytics/cfg/datasets/bvn/images/test"
test_labels_path = "C:/Users/DELL/Desktop/ultralytics/ultralytics/cfg/datasets/bvn/labels/test"

def empty_directory(path):
    # 确保路径存在并且是一个目录
    if os.path.exists(path) and os.path.isdir(path):
        # 遍历并删除目录中的所有文件和子目录
        for filename in os.listdir(path):
            file_path = os.path.join(path, filename)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print(f"Error while handling {file_path}: {e}")
    else:
        print(f"The path {path} does not exist or is not a directory.")

    # 使用绝对路径或相对路径来指定要清空的文件夹


empty_directory(train_images_path)
empty_directory(train_labels_path)
empty_directory(valid_images_path)
empty_directory(valid_labels_path)
empty_directory(test_images_path)
empty_directory(test_labels_path)
print("文件夹已全部清空")

ano_paths = []
i = 0

for dirname, _, filenames in os.walk(
        'C:/Users/DELL/Desktop/my_file/RDD2022_all_countries/China_Drone/train/labels'):
    for filename in filenames:
        if filename[-4:] == '.txt' and i < 1000:
            ano_paths += [(os.path.join(dirname, filename))]
            i += 1

n = len(ano_paths)
print(n)
N = list(range(n))
random.shuffle(N)

train_ratio = 0.6
valid_ratio = 0.2
test_ratio = 0.2

train_size = int(train_ratio * n)
valid_size = int(valid_ratio * n)

train_i = N[:train_size]
valid_i = N[train_size:train_size + valid_size]
test_i = N[train_size + valid_size:]

for i in train_i:
    ano_path = os.path.join(ano_paths[i])
    img_name = os.path.basename(ano_path)[:-4] + '.jpg'  # 获取不带扩展名的文件名
    img_path = os.path.join('C:/Users/DELL/Desktop/my_file/RDD2022_all_countries/China_Drone/train/images', img_name)

    try:
        shutil.copy(ano_path, train_labels_path)  # 使用 shutil 来复制文件
        shutil.copy(img_path, train_images_path)
    except Exception as e:  # 捕获所有异常
        print(f"Error occurred: {e}")
        continue

for i in valid_i:
    ano_path = os.path.join(ano_paths[i])
    img_name = os.path.basename(ano_path)[:-4] + '.jpg'
    img_path = os.path.join('C:/Users/DELL/Desktop/my_file/RDD2022_all_countries/China_Drone/train/images', img_name)

    try:
        shutil.copy(ano_path, valid_labels_path)
        shutil.copy(img_path, valid_images_path)
    except Exception as e:
        print(f"Error occurred: {e}")
        continue

for i in test_i:
    ano_path = os.path.join(ano_paths[i])
    img_name = os.path.basename(ano_path)[:-4] + '.jpg'
    img_path = os.path.join('C:/Users/DELL/Desktop/my_file/RDD2022_all_countries/China_Drone/train/images', img_name)

    try:
        shutil.copy(ano_path, test_labels_path)
        shutil.copy(img_path, test_images_path)
    except Exception as e:
        print(f"Error occurred: {e}")
        continue

运行结果