python生成图像样本的训练集、验证集和测试集

       准备好图像样本和标注文件后，需对样本集进行划分，生成训练集、验证集和测试集，关于训练集划分参考：https://blog.youkuaiyun.com/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():执行一个字符串表达式，并返回表达式的值.