yolov5划分数据集

对于YOLOv5，划分数据集的一种常见方法是按照一定比例将数据集分为训练集、验证集和测试集。以下是一个简单的示例代码，可以帮助你实现数据集的划分：

import os
import random
import shutil

def split_dataset(dataset_dir, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1):
    # 创建保存划分后数据集的目录
    train_dir = os.path.join(dataset_dir, 'train')
    val_dir = os.path.join(dataset_dir, 'val')
    test_dir = os.path.join(dataset_dir, 'test')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 获取所有图像文件的路径
    image_files = []
    for root, dirs, files in os.walk(dataset_dir):
        for file in files:
            if file.endswith('.jpg') or file.endswith('.jpeg') or file.endswith('.png'):
                image_files.append(os.path.join(root, file))
    
    # 随机打乱图像文件列表
    random.shuffle(image_files)

    # 计算划分后的样本数量
    total_samples = len(image_files)
    train_samples = int(total_samples * train_ratio)
    val_samples = int(total_samples * val_ratio)
    test_samples = total_samples - train_samples - val_samples

    # 划分数据集并移动文件到相应目录
    for i, file in enumerate(image_files):
        if i < train_samples:
            shutil.copy(file, train_dir)
        elif i < train_samples + val_samples:
            shutil.copy(file, val_dir)
        else:
            shutil.copy(file, test_dir)

    print(f"Successfully split the dataset into {train_samples} train samples, {val_samples} validation samples, and {test_samples} test samples.")

# 调用划分数据集函数
dataset_dir = '/path/to/dataset'  # 数据集所在目录
split_dataset(dataset_dir, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1)

在上述示例代码中，你需要将`dataset_dir`替换为你自己的数据集所在目录。可以根据需要调整训练集、验证集和测试集的比例。划分后的数据集将保存在相应的