Add code/cube3d/training/dat_mapping_merge.py

code/cube3d/training/dat_mapping_merge.py

@@ -0,0 +1,72 @@
+import os
+import json
+import re
+
+# 提取文件名中起始的连续数字（核心逻辑）
+def extract_starting_numbers(filename):
+    # 匹配字符串开头的连续数字
+    match = re.match(r'^\d+', filename)
+    if match:
+        return match.group()  # 返回起始数字字符串（如"4592"）
+    return filename  # 无起始数字则返回原文件名（小写处理后）
+
+# 处理单个ldr文件数据
+def process_ldr_data(lines, label_mapping, label_inverse_mapping, label_counter):
+    all_labels = []
+
+    for line in lines:
+        if line.startswith('1'):  # 只处理零件数据行
+            parts = line.split()
+            if len(parts) < 15:
+                print(f"Skipping invalid line: {line.strip()}")
+                continue
+
+            # 1. 统一转为小写
+            filename = parts[14].lower()
+            # 2. 提取起始连续数字作为标识
+            part_identifier = extract_starting_numbers(filename)
+            print(f"Processing file: {filename} → Identifier: {part_identifier}")
+            
+            if part_identifier not in label_mapping:
+                label_mapping[part_identifier] = label_counter
+                label_inverse_mapping[label_counter] = part_identifier
+                label_counter += 1
+            all_labels.append(label_mapping[part_identifier])
+
+    return label_mapping, label_inverse_mapping, label_counter
+
+# 处理一个文件夹中的所有ldr文件
+def process_all_ldr_in_folder(folder_path):
+    overall_label_mapping = {}
+    overall_label_inverse_mapping = {}
+    label_counter = 0
+
+    for root, dirs, files in os.walk(folder_path):
+        for file in files:
+            if file.endswith('.ldr'):
+                file_path = os.path.join(root, file)
+                with open(file_path, 'r') as f:
+                    lines = f.readlines()
+                    overall_label_mapping, overall_label_inverse_mapping, label_counter = process_ldr_data(
+                        lines, overall_label_mapping, overall_label_inverse_mapping, label_counter)
+
+    return overall_label_mapping, overall_label_inverse_mapping
+
+# 保存mapping到文件
+def save_mappings(label_mapping, label_inverse_mapping, output_dir):
+    os.makedirs(output_dir, exist_ok=True)
+    
+    with open(os.path.join(output_dir, 'label_mapping.json'), 'w') as f:
+        json.dump(label_mapping, f, indent=4)
+    
+    with open(os.path.join(output_dir, 'label_inverse_mapping.json'), 'w') as f:
+        json.dump(label_inverse_mapping, f, indent=4)
+
+
+if __name__ == "__main__":
+    input_folder = '/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300/ldr_rot_expand_trans'
+    output_folder = '/public/home/wangshuo/gap/assembly/data/car_1k/subset_bottom_300'
+
+    label_mapping, label_inverse_mapping = process_all_ldr_in_folder(input_folder)
+    save_mappings(label_mapping, label_inverse_mapping, output_folder)
+    print(f"Label mappings have been saved to {output_folder}")