准备好图像样本和标注文件后,需对样本集进行划分,生成训练集、验证集和测试集,关于训练集划分参考:https://blog.csdn.net/kieven2008/article/details/81582591
我的目录结构:
——data
——image
——0.jpg
——1.jpg
.
.
.
——labels
——0.txt
——1.txt
.
.
.
将样本按 training : validation : testing = 8:1:1,实现在image同级目录下创建文件夹sample_set,将按比例生成的train.txt、val.txt、test.txt存放在sample_set中。
import math
import os
import random
root_path=os.getcwd()
labels_path=root_path+"/JPEGImages"
label_list=os.listdir(labels_path)
train_and_val_list=random.sample(label_list,int(math.floor(len(label_list)*9/10)))
train_and_val_list.sort()
train_list=random.sample(train_and_val_list,int(math.floor(len(train_and_val_list)*8/9)))
train_list.sort()
val_list=list(set(train_and_val_list).difference(set(train_list)))
val_list.sort()
test_list=list(set(label_list).difference(set(train_and_val_list)))
test_list.sort()
def create_set(input_list_name,input_txt_name,key):
sample_set_path=root_path+"/"+"sample_set"
if not os.path.exists(sample_set_path):
os.mkdir(sample_set_path)
txt_file=open(sample_set_path+"/"+key+".txt",'w+')
for item in eval(input_list_name):
txt_file.write(labels_path+'/'+item+'\n')
txt_file.close()
key_name=['train','val','test']
for name in key_name:
list_name=name+"_list"
txt_name=name+".txt"
create_set(list_name,txt_name,name)
python编程笔记:
random.sample(list,n): 从list中随机截取n个元素,以list形式返回,不改变原有list;
math.floor(x):返回x的向下取整值;
set():创建一个无序不重复元素集,可进行关系测试,删除重复数据,还可以计算交集、差集、并集等;
set().difference(set):求差集;
eval():执行一个字符串表达式,并返回表达式的值.