关于图像分割任务中按照比例将数据集随机划分成训练集和测试集

2024-01-07 19:26:57

1. 前言

之前写了分类和检测任务划分数据集的脚本,三大任务实现了俩,基于强迫症,也实现一下图像分割的划分脚本

分类划分数据:关于图像分类任务中划分数据集,并且生成分类类别的josn字典文件

检测划分数据:关于目标检测中按照比例将数据集随机划分成训练集和测试集

之前写了目标检测的可视化脚本,明天看看要不要把分类和分割的可视化也实现一下,之前的博文里实现过,但都是大项目里的一个环节,看看能不能单独实现一下

?

这里的脚本和未划分数据的目录在同一路径

2. 完整代码

如下:

import random
import os
import shutil
from tqdm import tqdm


# 划分数据函数
def split_data(root, test_rate):
    images_path = [os.path.join(root, i) for i in os.listdir(root)]  # 获取所有图片路径
    test_split_path = random.sample(images_path, k=int(len(images_path) * test_rate))  # 随机采样测试集

    # 遍历所有图片
    for image_path in tqdm(images_path, desc='loading...'):
        # 获取相应的标注文件,这里需要根据目录、后缀更改
        mask_path = image_path.replace('images', 'masks')

        # 划分数据
        if image_path in test_split_path:  # 在测试集
            shutil.copy(image_path, './data/test/images')
            shutil.copy(mask_path, './data/test/masks')
        else:
            shutil.copy(image_path, './data/train/images')
            shutil.copy(mask_path, './data/train/masks')


if __name__ == '__main__':
    rawDataSet = './rawDataSet/images'  # 原始数据的图片路径

    if os.path.exists('./data'):  # 如果之前有,那么删除
        shutil.rmtree('./data')

    os.makedirs('./data/train/images')  # 训练集图片
    os.makedirs('./data/train/masks')   # 训练集标签
    os.makedirs('./data/test/images')  # 测试集图片
    os.makedirs('./data/test/masks')   # 测试集标签

    # 划分数据
    split_data(root=rawDataSet, test_rate=0.2)

运行过程:这里测试五个数据

运行结果:

?

3. 注意的点

根据训练图片找到对应mask模板的时候,这里要根据目录更换

文章来源:https://blog.csdn.net/qq_44886601/article/details/135442324
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。