parent
b23a75ac73
commit
6a437803be
@ -0,0 +1,68 @@
|
|||||||
|
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.')
|
Loading…
Reference in new issue