python讀取資料集檔案下所有檔案並打亂劃分生成訓練測試txt檔案(生成train.txt、test.txt,順序隨機,預設比例8:2)

cv白嫖王發表於2020-11-26

因為每次讀取資料集都要寫載入程式碼,十分浪費時間,而我本人又比較懶。故寫了個python小工具,一鍵執行,讀取資料集檔案下所有檔案並打亂劃分生成訓練測試txt檔案(生成train.txt、test.txt,順序隨機,預設比例8:2),十分方便。

執行指令碼前需要將資料夾整理為這類的形式:
root = “C:\Users\hq\Desktop\HoldingObject\pokemon”(記得改為自己的根路徑)
路徑樹如下(資料夾、圖片數量無限制):
—pokemon
——0
(此處資料夾名為標籤)

————1,.jpg、2.jpg、3.jpg、。。。(此處命名無所謂,不影響)
——1 (此處資料夾名為標籤)
————1,.jpg、2.jpg、3.jpg、。。。(此處命名無所謂,不影響)
——2 (此處資料夾名為標籤)
————1,.jpg、2.jpg、3.jpg、。。。(此處命名無所謂,不影響)**

import os
import numpy as np

root = r"C:\Users\hq\Desktop\HoldingObject\pokemon"

#構建所有檔名的列表,dir為label
filename = []
#label = []
dirs = os.listdir(root)
for dir in dirs:
    dir_path = root + '\\' + dir
    names = os.listdir(dir_path)
    for n in names:
        filename.append(dir_path + '\\' + n + '\t' + dir)

#打亂檔名列表
np.random.shuffle(filename)
#劃分訓練集、測試集,預設比例4:1
train = filename[:int(len(filename)*0.8)]
test = filename[int(len(filename)*0.8):]

#分別寫入train.txt, test.txt	
with open('train.txt', 'w') as f1, open('test.txt', 'w') as f2:
    for i in train:
        f1.write(i + '\n')
    for j in test:
        f2.write(j + '\n')

print('成功!')

生成的檔案如下圖所示:

先佔個坑,等我回去補圖。。。

相關文章