import os
import random

train_ratio = 0.9 #训练集比例
test_ratio = 1-train_ratio #测试集比例

rootdata = r"dataset"

train_list, test_list = [],[]

class_flag = -1#初始类别标签
t=1
for a,b,c in os.walk(rootdata):#目录路径，子目录名，文件名
    # if t==1:
    #     # print(a)
    #     print(b)
    #     print(len(c))
    #     t=-1

    for i in range(0, int(len(c)*train_ratio)):#根据训练集比例确定训练集的文件数量
        train_data = os.path.join(a, c[i])+'\t'+str(class_flag)+'\n'
        # print('666'+train_data)
        train_list.append(train_data)

    for i in range(int(len(c) * train_ratio), len(c)):
        test_data = os.path.join(a, c[i]) + '\t' + str(class_flag)+'\n'
        test_list.append(test_data)

    class_flag += 1

random.shuffle(train_list)#随机打乱训练集列表 train_list 中的元素顺序
random.shuffle(test_list)

with open('train.txt','w',encoding='UTF-8') as f:
    for train_img in train_list:
        f.write(str(train_img))#将训练集中的文件路径和类别标签写入文件 'train.txt'，将其转换为字符串形式并写入文件。

with open('test.txt','w',encoding='UTF-8') as f:
    for test_img in test_list:
        f.write(test_img)