240720_图像语义分割中数据预处理-划分训练集与验证集

240720_图像语义分割中数据预处理-划分训练集与验证集

在划分的过程中，我们的基本理念为：

分别读取原图文件夹以及标签文件夹中的文件到列表中，然后计算出文件数量，使用设定好的划分比例乘这个数量，得到对应的训练集验证集数量。再使用随机采样确定训练集和验证集的文件列表，去掉后缀后写入txt文本文件中，后续训练的时候读取这些文本文件。

# data_split.py

import os
import random
import numpy as np
from PIL import Image
from tqdm import tqdm


# 定义训练集和验证集比例，此时定义的训练集train：验证集val=9：1
# 此处没有定义测试集，但是预留了定义测试集的代码，需要可自行修改
# trainval+test=1,train<=trainval
trainval_rate=1
train_rate=0.9
test_rate=0
# 指向数据集所在的路径，此处只给一个总文件夹，后面进行拼接
data_path='Datasets'

if __name__ == '__main__':
    # 获取路径
    jpg_path=os.path.join(data_path, 'RoadCrack0720/ltimages')
    mask_path=os.path.join(data_path,'RoadCrack0720/ltmasks')
    save_path=os.path.join(data_path,'RoadCrack0720/Segmentation')
    # 把所有标签图名字放在列表中，以供后续写入
    temp_png=os.listdir(mask_path)
    # 这里是为了避免混入奇怪的东西，在检查一下保证全都是png文件
    total_png=[]
    for png in temp_png:
        if png.endswith('.png'):
            total_png.append(png)
    # 计算总样本量以供计算验证集比例
    total_num=len(total_png)
    trainval_num=int(total_num*trainval_rate)
    train_num=int(trainval_num*train_rate)
    test_num=int(total_num*test_rate)
    list=range(total_num)
    # 通过随机采样确定训练集和验证集的文件列表，此处可以通过指定随机数种子来保证结果复现
    # 使用random模块的sample函数从list中随机选择trainval_num个元素，返回一个新列表
    random.seed(0)
    trainval=random.sample(list,trainval_num)
    train=random.sample(trainval,train_num)
    # 剩余未被选中的索引即为测试集，此处也可以不要这行代码，在后续写入过程中可以直接else进行排除
    test = np.setdiff1d(list, trainval)
    # 打印训练集和验证集的大小
    print('trainval:',len(trainval))
    print('train:',len(train))
    print('test:',len(test))
    # 打开要写入的txt文件，并且指定写入模式为'w'，代表覆盖写入，如果不存在则新建
    ftrainval=open(os.path.join(save_path,'trainval.txt'),'w')
    ftrain=open(os.path.join(save_path,'train.txt'),'w')
    fval=open(os.path.join(save_path,'val.txt'),'w')
    ftest=open(os.path.join(save_path,'test.txt'),'w')
    # 写入文件
    for i in list:
        # 取到从后往前数第四个字符，也就是删除掉了'.png'字符，这里写入的是不带后缀的文件名。
        name=total_png[i][:-4]+'\n'
        if i in trainval:
            ftrainval.write(name)
            if i in train:
                ftrain.write(name)
            else:
                fval.write(name)
        if i in test:
            ftest.write(name)
    # 关闭文件
    ftrainval.close()
    ftrain.close()
    fval.close()
    ftest.close()

    pass

至此划分完毕。