import os
import random
# 垃圾名到标签的映射
waste_to_label = {
    'bum': '0',
    'carton': '1',
    'egg_shells': '2',
    'plastic_bags':'3',
    'plastic_boxes':'4',
    'plastic_buckets':'5'
    # 假设这里有剩余的映射关系
    # '垃圾分类4': '3',
    # '垃圾分类5': '4',
    # ...
    # '垃圾分类10': '9'
}

# 要遍历的文件夹路径
waste_folder_path = 'D:\clsGarbageCode\garbage'

# 输出文件路径
output_file_paths = {
    'train': 'D:\clsGarbageCode\garbage\\train_list.txt',
    'val': 'D:\clsGarbageCode\garbage\\val_list.txt',
    'test': 'D:\clsGarbageCode\garbage\\test_list.txt'
    # 'label': 'D:\clsGarbageCode\garbage\\label_list.txt'
}

images_info = []

# 遍历 waste 文件夹
for folder_name, label in waste_to_label.items():
    # with open(output_file_paths['label'], 'w') as f:
    #     f.write(f'{folder_name}\n')
    folder_path = os.path.join(waste_folder_path, folder_name)
    
    # 遍历文件夹内的所有图片
    for img_name in os.listdir(folder_path):
        img_path = os.path.join(folder_name, img_name)
        images_info.append((img_path,label))

random.shuffle(images_info)   

# 计算每个部分应该包含的图片数量
total_images = len(images_info)
train_size = int(0.6 * total_images)
val_size = int(0.2 * total_images)
test_size = total_images - train_size - val_size

# 分配数据到不同的文件
train_info = images_info[:train_size]
val_info = images_info[train_size:train_size + val_size]
test_info = images_info[train_size + val_size:]

# 写入到不同的文件
for split, file_path in output_file_paths.items():
    with open(file_path, 'w') as f:
        if split == 'train':
            for img_path, label in train_info:
                f.write(f'{img_path} {label}\n')
        elif split == 'val':
            for img_path, label in val_info:
                f.write(f'{img_path} {label}\n')
        elif split == 'test':
            for img_path, label in test_info:
                f.write(f'{img_path} {label}\n')

print('train_list.txt, val_list.txt, and test_list.txt have been created successfully.')