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
至此划分完毕。