python读取数据集文件下所有文件并打乱划分生成训练测试txt文件（生成train.txt、test.txt，顺序随机，默认比例8:2）

因为每次读取数据集都要写加载代码，十分浪费时间，而我本人又比较懒。故写了个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