#使用python处理分割数据集，并按比例随机划分训练集与测试集

对分割数据集进行打label处理

一些分割数据集会给出类似的图片作为label

我们要将其对应的灰度值转化为标签

获取文件列表

def get_files(files_dir):
    files_list = []      
    for file in os.listdir(files_dir):
        files_list.append(file)
    return  files_list

根据标签进行转换

files_dir = './label/'
files_list = get_files(files_dir) 
print(len(files_list))
for i in range(len(files_list)):
    a = Image.open("./label/" + files_list[i])
    #a.show()
    a =np.array(a)
    a[a == 0] = 0
    a[a == 200] = 1
    a[a == 204] = 1
    a[a == 213] = 255
    a[a == 209] = 1
    a[a == 206] = 255
    a[a == 207] = 255
    a[a == 201] = 2
    a[a == 203] = 2
    a[a == 211] = 255
    a[a == 208] = 255
    a[a == 216] = 255
    a[a == 217] = 3
    a[a == 215] = 255
    a[a == 218] = 255
    a[a == 219] = 255
    a[a == 210] = 5
    a[a == 232] = 255
    a[a == 214] = 6
    a[a == 202] = 255
    a[a == 220] = 7
    a[a == 221] = 7
    a[a == 222] = 7
    a[a == 231] = 255
    a[a == 224] = 7
    a[a == 225] = 7
    a[a == 226] = 7
    a[a == 230] = 255
    a[a == 228] = 255
    a[a == 229] = 255
    a[a == 233] = 255
    a[a == 205] = 8
    a[a == 212] = 255
    a[a == 227] = 8
    a[a == 223] = 255
    a[a == 250] = 8
    a[a == 249] = 255
    a[a == 255] = 255

    #print(a.shape)
    print(str(round(i*100/len(files_list) , 2 )) + "%") #打印处理进度
    a = Image.fromarray(a)
    a.save(("./label_mask/" + files_list[i]),"png")

至此，我们就完成了数据集label的转换

划分数据集

我们的目标是生成三个txt文件，里面按照比例随机分配数据集，分别是

• train.txt：所有训练集的文件名称
• trainval.txt：所有验证集的文件名称
• val.txt：所有测试集的文件名称

导入依赖库

import numpy as np
from PIL import Image
import os

获取指定目录下面数据集的list

def get_files(files_dir):
    files_list = []      
    for file in os.listdir(files_dir):
        files_list.append(file)
    return  files_list


files_dir = './image/'#图片文件的路径
files_list = get_files(files_dir) 
print("Total Sample:" + str(len(files_list)))

按照比例计算各个子集数量

ratio_train = 0.8 #训练集比例
ratio_trainval = 0.1 #验证集比例
ratio_val = 0.1 #测试集比例
assert (ratio_train + ratio_trainval + ratio_val) == 1.0,'Total ratio Not equal to 1' ##检查总比例是否等于1
np.random.shuffle(files_list) ##打乱文件列表
cnt_val = round(len(files_list) * ratio_val ,0)
cnt_trainval = round(len(files_list) * ratio_trainval ,0)
cnt_train = len(files_list) - cnt_val - cnt_trainval
print("val Sample:" + str(cnt_val))
print("trainval Sample:" + str(cnt_trainval))
print("train Sample:" + str(cnt_train))

这里训练集采用总数减去其他两个子集的处理方式，规避取整问题

随机将数据集分配到三个子集

np.random.shuffle(files_list) ##打乱文件列表
train_list = []
trainval_list = []
val_list = []

for i in range(int(cnt_train)):
    train_list.append(files_list[i])

for i in range(int(cnt_train) ,int(cnt_train + cnt_trainval)):
    trainval_list.append(files_list[i])

for i in range(int(cnt_train + cnt_trainval) ,int(cnt_train + cnt_trainval + cnt_val)):
    val_list.append(files_list[i])

保存索引文件（修改）

file = open('./index/train.txt','w')
for i in range(len(train_list)):   
    name = str(train_list[i])
    index = name.rfind('.')
    name = name[:index]
    file.write(name+'\n')
file.close()

file = open('./index/trainval.txt','w')
for i in range(len(trainval_list)):   
    name = str(trainval_list[i])
    index = name.rfind('.')
    name = name[:index]
    file.write(name+'\n')
file.close()

file = open('./index/val.txt','w')
for i in range(len(val_list)):   
    name = str(val_list[i])
    index = name.rfind('.')
    name = name[:index]
    file.write(name+'\n')
file.close()
print("Success! output:")
print('./index/train.txt')
print('./index/trainval.txt')
print('./index/val.txt')

总结

经过以上操作后，生成的TXT文件节选如下

# val.txt
1.jpg
.
.
.
xxxx.jpg

至此 完成我们的目标