yolov5数据集划分脚本

你可以使用以下脚本来划分YOLOv5数据集：

import os
import random
import shutil

def split_dataset(data_dir, train_ratio=0.8, val_ratio=.1, test_ratio=0.1):
    assert train_ratio + val_ratio + test_ratio == 1, "The sum of train_ratio, val_ratio, and test_ratio must be 1."
    
    file_list = os.listdir(data_dir)
    random.shuffle(file_list)
    
    train_size = int(len(file_list) * train_ratio)
    val_size = int(len(file_list) * val_ratio)
    
    train_files = file_list[:train_size]
    val_files = file_list[train_size:train_size+val_size]
    test_files = file_list[train_size+val_size:]
    
    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)
    
    for file in train_files:
        shutil.move(os.path.join(data_dir, file), os.path.join(train_dir, file))
        
    for file in val_files:
        shutil.move(os.path.join(data_dir, file), os.path.join(val_dir, file))
        
    for file in test_files:
        shutil.move(os.path.join(data_dir, file), os.path.join(test_dir, file))

你可以调用`split_dataset`函数来划分数据集。`data_dir`是包含所有图像和标签的文件夹路径。`train_ratio`，`val_ratio`和`test_ratio`是用来指定训练集、验证集和测试集的比例。调用该函数后，它将会在`data_dir`下创建`train`，`val`和`test`三个文件夹，并将数据文件移动到相应的文件夹中。

请注意，这个脚本假设你的数据集中图像和标签的文件名是一一对应的，并且它们都存储在同一个文件夹中。