因为每次读取数据集都要写加载代码,十分浪费时间,而我本人又比较懒。故写了个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('成功!')
生成的文件如下图所示:
train.txt
test.txt