You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
68 lines
2.1 KiB
68 lines
2.1 KiB
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.') |