数据集分离为训练集和验证集

Python代码

import os
import random
import shutil

def split_dataset_yolo(image_folder, label_folder, train_ratio):
    # 创建目标文件夹
    train_image_folder = os.path.join(image_folder, 'train')
    train_label_folder = os.path.join(label_folder, 'train')
    valid_image_folder = os.path.join(image_folder, 'val')
    valid_label_folder = os.path.join(label_folder, 'val')

    os.makedirs(train_image_folder, exist_ok=True)
    os.makedirs(train_label_folder, exist_ok=True)
    os.makedirs(valid_image_folder, exist_ok=True)
    os.makedirs(valid_label_folder, exist_ok=True)

    # 获取图片文件列表
    image_files = [f for f in os.listdir(image_folder) if f.endswith(('.jpg', '.png', '.jpeg'))]

    # 计算训练集和验证集的数量
    num_images = len(image_files)
    num_train = int(num_images * train_ratio)
    num_valid = num_images - num_train

    # 随机打乱图片文件列表
    random.shuffle(image_files)

    # 将图片文件和对应的标签文件分配到训练集和验证集
    for i, image_file in enumerate(image_files):
        label_file = image_file.replace('.' + image_file.split('.')[-1], '.txt')
        if i < num_train:
            shutil.copy2(os.path.join(image_folder, image_file), train_image_folder)
            shutil.copy2(os.path.join(label_folder, label_file), train_label_folder)
        else:
            shutil.copy2(os.path.join(image_folder, image_file), valid_image_folder)
            shutil.copy2(os.path.join(label_folder, label_file), valid_label_folder)

    print("Dataset split completed.")

# 指定图片文件夹和标签文件夹的路径
image_folder_path = ''
label_folder_path = ''

# 指定训练集所占比例（0.9表示训练集占总数据的90%）
train_ratio = 0.9

# 执行数据集划分
split_dataset_yolo(image_folder_path, label_folder_path, train_ratio)

效果图

图片

标签