From 6a437803bebd767e0a583dc6a63a5eba2a636cec Mon Sep 17 00:00:00 2001
From: hnu202401010612 <3117511861@qq.com>
Date: Thu, 26 Dec 2024 00:04:05 +0800
Subject: [PATCH] ADD file via upload

---
 getLabelList.py | 68 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 68 insertions(+)
 create mode 100644 getLabelList.py

diff --git a/getLabelList.py b/getLabelList.py
new file mode 100644
index 0000000..d8a026e
--- /dev/null
+++ b/getLabelList.py
@@ -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.')
\ No newline at end of file