python生成圖像樣本的訓練集、驗證集和測試集

       準備好圖像樣本和標註文件後，需對樣本集進行劃分，生成訓練集、驗證集和測試集，關於訓練集劃分參考：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():執行一個字符串表達式，並返回表達式的值.