Yolo数据集按比例划分

最新推荐文章于 2024-08-01 11:01:10 发布

科研分母

最新推荐文章于 2024-08-01 11:01:10 发布

阅读量1.1k

点赞数 1

分类专栏：深度学习-Yolo 文章标签： YOLO 深度学习人工智能

本文链接：https://blog.csdn.net/LMH564123/article/details/130441795

版权

深度学习-Yolo 专栏收录该内容

5 篇文章 0 订阅

订阅专栏

该代码示例展示了如何使用Python将Yolo数据集按预设比例（训练集、验证集、测试集）划分，并确保原始图像子目录只包含图像文件，无其他格式。首先设置随机数种子，然后定义文件路径，获取所有图像和标签文件的名称，随机打乱顺序，计算各集合的数量，最后创建输出目录并将文件复制到相应集合的目录中。

摘要由CSDN通过智能技术生成

Yolo数据集按比例划分,其中一个必须修改的地方，一个需要注意的地方，还有就是要注意原始图像所在的子目录里只能存在图像格式的文件，不能有其他格式的文件

import os  # 用于处理文件路径、创建目录等操作
import random  # 用于生成随机数种子、打乱列表等操作
import shutil  # 用于生成随机数种子、打乱列表等操作

# 设置随机数种子
random.seed(123)

# 定义文件夹路径（需要按照自己文件目录修改）
image_dir = 'path/to/images'  # 原始图像所在的子目录
label_dir = 'path/to/lables'  # 原始标签所在的子目录
output_dir = 'yolo_dataset'  # 处理后的数据集输出目录

# 定义训练集、验证集和测试集比例（根据自己的需求修改，可改可不改）
train_ratio = 0.7  # 训练集比例
valid_ratio = 0.15  # 验证集比例
test_ratio = 0.15  # 测试集比例

# 获取所有图像文件和标签文件的文件名（不包括文件扩展名）
image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_dir)]  # 提取所有图像文件的文件名列表
label_filenames = [os.path.splitext(f)[0] for f in os.listdir(label_dir)]  # 提取所有标签文件的文件名列表

# 随机打乱文件名列表
random.shuffle(image_filenames)  # 打乱图像文件的文件名列表

# 计算训练集、验证集和测试集的数量
total_count = len(image_filenames)  # 总文件数
train_count = int(total_count * train_ratio)  # 训练集文件数
valid_count = int(total_count * valid_ratio)  # 验证集文件数
test_count = total_count - train_count - valid_count  # 测试集文件数

# 定义输出文件夹路径
train_image_dir = os.path.join(output_dir, 'train', 'images')  # 训练集图像输出目录
train_label_dir = os.path.join(output_dir, 'train', 'labels')  # 训练集标签输出目录
valid_image_dir = os.path.join(output_dir, 'valid', 'images')  # 验证集图像输出目录
valid_label_dir = os.path.join(output_dir, 'valid', 'labels')  # 验证集标签输出目录
test_image_dir = os.path.join(output_dir, 'test', 'images')  # 测试集图像输出目录
test_label_dir = os.path.join(output_dir, 'test', 'labels')  # 测试集标签输出目录

# 创建输出文件夹
os.makedirs(train_image_dir, exist_ok=True)  # 创建训练集图像输出目录
os.makedirs(train_label_dir, exist_ok=True)  # 创建训练集标签输出目录
os.makedirs(valid_image_dir, exist_ok=True)  # 创建验证集图像输出目录
os.makedirs(valid_label_dir, exist_ok=True)  # 创建验证集标签输出目录
os.makedirs(test_image_dir, exist_ok=True)  # 创建测试集图像输出目录
os.makedirs(test_label_dir, exist_ok=True)  # 创建测试集标签输出目录

# 将图像和标签文件划分到不同的数据集中
for i, filename in enumerate(image_filenames):
    # 如果文件数量小于训练数据集大小，则将文件复制到训练数据集目录中
    if i < train_count:
        output_image_dir = train_image_dir
        output_label_dir = train_label_dir
    # 如果文件数量小于训练数据集大小+验证数据集大小，则将文件复制到验证数据集目录中
    elif i < train_count + valid_count:
        output_image_dir = valid_image_dir
        output_label_dir = valid_label_dir
    # 否则，将文件复制到测试数据集目录中
    else:
        output_image_dir = test_image_dir
        output_label_dir = test_label_dir

    # 复制图像文件（注意是否为jpg txt保持一致）
    src_image_path = os.path.join(image_dir, filename + '.jpg')  # 获取图像文件的源路径
    dst_image_path = os.path.join(output_image_dir, filename + '.jpg')  # 获取图像文件的目标路径
    shutil.copy(src_image_path, dst_image_path)  # 复制图像文件到目标路径

    # 复制标签文件
    src_label_path = os.path.join(label_dir, filename + '.txt')  # 获取标签文件的源路径
    dst_label_path = os.path.join(output_label_dir, filename + '.txt')  # 获取标签文件的目标路径
    shutil.copy(src_label_path, dst_label_path)  # 复制标签文件到目标路径