Add code/cube3d/training/merge_obj_by_prefix.py

code/cube3d/training/merge_obj_by_prefix.py

@@ -0,0 +1,94 @@
+import os
+import re
+import shutil
+import json  # 修复：补充json导入（原代码虽未用，但避免后续扩展报错）
+import pandas as pd  # 修复：补充pandas导入（同理，确保依赖完整）
+from collections import defaultdict
+
+def extract_merge_key(filename):
+    """提取合并键：开头连续数字为键（用于分组），字母开头返回原文件名并标记"""
+    name_without_ext = os.path.splitext(filename)[0]
+    # 匹配开头连续数字（如"123_bolt"→"123"，"A001"无匹配）
+    num_match = re.match(r'^\d+', name_without_ext)
+    if num_match:
+        return num_match.group(), 'num_start'  # 数字开头：返回数字键+类型
+    else:
+        return name_without_ext, 'letter_start'  # 字母开头：返回原名+类型
+
+def merge_obj_by_prefix(input_dir, output_dir):
+    """
+    按.obj前缀合并：每组复制排序后第一个文件到输出目录，字母开头文件单独打印
+    """
+    # 1. 初始化存储：数字前缀→文件路径列表，字母开头文件列表
+    num_groups = defaultdict(list)  # key: 数字前缀, value: [文件路径1, 文件路径2, ...]
+    letter_files = []               # 存储字母开头的文件名（仅记录，不复制）
+
+    # 2. 扫描输入目录所有.obj文件（递归包含子目录）
+    print(f"开始扫描 {input_dir} 下的.obj文件...\n")
+    for root, _, files in os.walk(input_dir):
+        for file in files:
+            if file.lower().endswith('.obj'):  # 兼容 .OBJ 大写后缀
+                file_path = os.path.join(root, file)
+                merge_key, start_type = extract_merge_key(file)
+                
+                if start_type == 'num_start':
+                    # 数字开头：加入对应分组
+                    num_groups[merge_key].append(file_path)
+                    print(f"  数字开头：{file} → 归入组【{merge_key}】")
+                else:
+                    # 字母开头：加入未合并列表
+                    letter_files.append(file)
+                    print(f"  字母开头：{file} → 不合并（仅记录）")
+
+    # 3. 检查是否有可处理的文件
+    if not num_groups and not letter_files:
+        print(f"警告：在 {input_dir} 中未找到任何.obj文件！")
+        return
+
+    # 4. 创建输出目录（清空原有内容，避免旧文件干扰）
+    if os.path.exists(output_dir):
+        shutil.rmtree(output_dir)  # 删除原有目录及文件
+    os.makedirs(output_dir, exist_ok=True)
+    print(f"\n已创建干净的输出目录：{output_dir}")
+
+    # 5. 处理每个数字分组：排序后复制第一个文件到输出目录
+    copied_count = 0
+    print(f"\n开始处理合并组（共 {len(num_groups)} 组）：")
+    for group_key in sorted(num_groups.keys()):  # 按数字前缀排序（如123→456）
+        file_paths = num_groups[group_key]
+        # 对组内文件按文件名排序（确保每次复制顺序一致）
+        sorted_files = sorted(file_paths, key=lambda x: os.path.basename(x))
+        first_file = sorted_files[0]  # 取排序后第一个文件
+        first_filename = os.path.basename(first_file)
+        
+        # 复制文件到输出目录（保留原文件名）
+        dest_path = os.path.join(output_dir, first_filename)
+        shutil.copy2(first_file, dest_path)  # copy2 保留文件元信息（修改时间等）
+        copied_count += 1
+        
+        print(f"  组【{group_key}】：复制排序后第一个文件 → {first_filename}")
+
+    # 6. 单独打印字母开头的未合并文件
+    print(f"\n=== 字母开头的未合并零件列表（共 {len(letter_files)} 个）===")
+    if letter_files:
+        for idx, file in enumerate(sorted(letter_files), 1):  # 排序后打印，更整齐
+            print(f"  {idx}. {file}")
+    else:
+        print("  无字母开头的未合并零件")
+
+    # 7. 输出最终统计
+    print(f"\n=== 处理完成统计 ===")
+    print(f"  合并组总数：{len(num_groups)} 组")
+    print(f"  成功复制文件数：{copied_count} 个（输出目录：{output_dir}）")
+    print(f"  未合并文件数：{len(letter_files)} 个（仅记录，未复制）")
+
+if __name__ == "__main__":
+    # --------------------------
+    # 配置参数（请根据实际情况修改）
+    # --------------------------
+    INPUT_DIR = "/public/home/wangshuo/gap/assembly/data/part_obj_300"  # .obj源文件目录（递归扫描）
+    OUTPUT_DIR = "/public/home/wangshuo/gap/assembly/data/part_obj_300_merge_by_prefix"  # 合并后文件输出目录（仅存复制的文件）
+    
+    # 执行合并逻辑
+    merge_obj_by_prefix(INPUT_DIR, OUTPUT_DIR)
+    print(f"\n✅ .obj零件合并复制完成！输出目录仅保留每组排序后第一个文件。")