【yolov5】将标注好的数据集进行划分（附完整可运行python代码）

在YOLOv5中，我们通常需要将数据集划分为训练集、验证集和测试集。以下是一个简单的Python脚本，用于将标注好的数据集进行划分。

import os
import random
import shutil
 
def split_dataset(data_dir, train_ratio, val_ratio, test_ratio):
    assert sum([train_ratio, val_ratio, test_ratio]) == 1  # 确保划分比例之和为1
    train_dir = os.path.join(data_dir, 'train')
    val_dir = os.path.join(data_dir, 'val')
    test_dir = os.path.join(data_dir, 'test')
    
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)
    
    images = os.listdir(os.path.join(data_dir, 'images'))
    annotations = os.listdir(os.path.join(data_dir, 'labels'))
    
    for image, annotation in zip(images, annotations):
        random.seed(image)
        if random.random() < train_ratio:
            shutil.copy(os.path.join(data_dir, 'images', image), train_dir)
            shutil.copy(os.path.join(data_dir, 'labels', annotation), train_dir)
        elif random.random() < train_ratio + val_ratio:
            shutil.copy(os.path.join(data_dir, 'images', image), val_dir)
            shutil.copy(os.path.join(data_dir, 'labels', annotation), val_dir)
        else:
            shutil.copy(os.path.join(data_dir, 'images', image), test_dir)
            shutil.copy(os.path.join(data_dir, 'labels', annotation), test_dir)
 
if __name__ == '__main__':
    data_dir = 'path/to/your/dataset'  # 指定数据集目录
    split_dataset(data_dir, 0.8, 0.1, 0.1)  # 设置训练集:验证集:测试集的比例

请确保将 data_dir 变量设置为您的数据集的路径，并根据需要调整 train_ratio、val_ratio 和 test_ratio 以反映您希望的数据集划分比例。这个脚本会将图像和对应的标注文件按照指定的比例复制到相应的 train, val, 和 test 目录中。