Add code/cube3d/training/group_objs_by_geometry.py

code/cube3d/training/group_objs_by_geometry.py

@@ -0,0 +1,216 @@
+import os
+import numpy as np
+import open3d as o3d
+import shutil
+from collections import defaultdict
+import pandas as pd
+
+# --------------------------
+# 配置参数（根据需求调整）
+# --------------------------
+# 包围盒筛选阈值（差异比例，0-1）
+BBOX_DIMENSION_TOLERANCE = 0.2  # 长宽高单个维度差异不超过20%
+BBOX_VOLUME_TOLERANCE = 0.3     # 总体积差异不超过30%
+
+# ICP参数
+ICP_DISTANCE_THRESHOLD = 0.05   # ICP匹配距离阈值
+ICP_MAX_ITERATIONS = 100        # ICP最大迭代次数
+SIMILARITY_THRESHOLD = 0.9      # 几何相似度阈值（0-1）
+
+# 点云参数
+SAMPLE_POINT_NUM = 1000         # 点云采样点数
+
+
+def load_obj_and_calculate_bbox(obj_path):
+    """加载OBJ模型，计算包围盒（AABB）和点云（带归一化）"""
+    try:
+        # 加载网格模型
+        mesh = o3d.io.read_triangle_mesh(obj_path)
+        if not mesh.has_triangles():
+            print(f"警告：{os.path.basename(obj_path)} 无三角面，无法处理")
+            return None, None, None
+        
+        # 1. 计算轴对齐包围盒（AABB）
+        bbox = mesh.get_axis_aligned_bounding_box()
+        bbox_min = bbox.min_bound  # [min_x, min_y, min_z]
+        bbox_max = bbox.max_bound  # [max_x, max_y, max_z]
+        
+        # 计算包围盒尺寸（长宽高）和体积
+        bbox_dimensions = bbox_max - bbox_min  # [dx, dy, dz]
+        bbox_volume = np.prod(bbox_dimensions)  # 体积 = dx*dy*dz
+        
+        # 2. 生成点云并预处理（归一化，为ICP做准备）
+        pcd = mesh.sample_points_uniformly(number_of_points=SAMPLE_POINT_NUM)
+        pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
+        
+        # 归一化：平移到原点 + 缩放到单位球
+        pcd_center = pcd.get_center()
+        pcd.translate(-pcd_center)
+        pcd_scale = np.max(np.linalg.norm(np.asarray(pcd.points), axis=1))
+        if pcd_scale > 1e-6:
+            pcd.scale(1 / pcd_scale, center=np.zeros(3))
+        
+        return bbox_dimensions, bbox_volume, pcd
+    
+    except Exception as e:
+        print(f"错误：处理 {os.path.basename(obj_path)} 失败 - {str(e)}")
+        return None, None, None
+
+
+def is_bbox_similar(bbox_dim1, vol1, bbox_dim2, vol2):
+    """判断两个包围盒是否相似（尺寸和体积差异在阈值内）"""
+    # 检查单个维度差异（dx, dy, dz）
+    dim_diff = np.abs(bbox_dim1 - bbox_dim2) / np.maximum(bbox_dim1, bbox_dim2)
+    if np.any(dim_diff > BBOX_DIMENSION_TOLERANCE):
+        return False  # 任一维度差异过大
+    
+    # 检查体积差异
+    vol_diff = abs(vol1 - vol2) / max(vol1, vol2)
+    if vol_diff > BBOX_VOLUME_TOLERANCE:
+        return False  # 体积差异过大
+    
+    return True  # 包围盒相似
+
+
+def calculate_icp_similarity(pcd1, pcd2):
+    """ICP计算点云相似度"""
+    icp_result = o3d.pipelines.registration.registration_icp(
+        source=pcd1,
+        target=pcd2,
+        max_correspondence_distance=ICP_DISTANCE_THRESHOLD,
+        init=np.eye(4),
+        estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(),
+        criteria=o3d.pipelines.registration.ICPConvergenceCriteria(
+            relative_fitness=1e-6,
+            relative_rmse=1e-6,
+            max_iteration=ICP_MAX_ITERATIONS
+        )
+    )
+    
+    # 计算相似度（1-归一化距离）
+    avg_distance = icp_result.transformation_db
+    similarity = max(0.0, 1.0 - (avg_distance / ICP_DISTANCE_THRESHOLD))
+    return similarity
+
+
+def group_objs_by_geometry(input_dir):
+    """先通过包围盒筛选，再用ICP分组"""
+    obj_info = []
+    print(f"开始加载 {input_dir} 下的OBJ文件并计算包围盒...")
+    
+    # 1. 加载所有OBJ文件的信息（包围盒+点云）
+    for root, _, files in os.walk(input_dir):
+        for file in files:
+            if file.lower().endswith('.obj'):
+                obj_path = os.path.join(root, file)
+                bbox_dim, bbox_vol, pcd = load_obj_and_calculate_bbox(obj_path)
+                
+                if bbox_dim is not None and pcd is not None and len(pcd.points) > 100:
+                    obj_info.append({
+                        "path": obj_path,
+                        "name": file,
+                        "bbox_dim": bbox_dim,
+                        "bbox_vol": bbox_vol,
+                        "pcd": pcd
+                    })
+                    print(f"  已加载：{file} → 包围盒尺寸 {bbox_dim.round(2)}，体积 {bbox_vol:.2f}")
+    
+    if len(obj_info) < 2:
+        print(f"提示：仅找到 {len(obj_info)} 个有效文件，无需分组")
+        return {0: [obj_info[0]["path"]]} if obj_info else {}
+    
+    # 2. 按几何分组（先包围盒筛选，再ICP验证）
+    groups = defaultdict(list)
+    ungrouped = obj_info.copy()
+    group_id = 0
+    
+    print(f"\n开始分组（共 {len(ungrouped)} 个文件）...")
+    while ungrouped:
+        base = ungrouped.pop(0)
+        groups[group_id].append(base["path"])
+        print(f"\n组 {group_id}：以 {base['name']} 为基准（尺寸 {base['bbox_dim'].round(2)}）")
+        
+        to_remove = []
+        for idx, candidate in enumerate(ungrouped):
+            # 第一步：包围盒筛选（直接比较尺寸和体积）
+            if not is_bbox_similar(
+                base["bbox_dim"], base["bbox_vol"],
+                candidate["bbox_dim"], candidate["bbox_vol"]
+            ):
+                # 包围盒差异过大，跳过ICP
+                print(f"  包围盒不匹配：{candidate['name']}（尺寸 {candidate['bbox_dim'].round(2)}）→ 跳过")
+                continue
+            
+            # 第二步：ICP验证几何形状
+            similarity = calculate_icp_similarity(base["pcd"], candidate["pcd"])
+            print(f"  ICP匹配 {candidate['name']} → 相似度 {similarity:.3f}（阈值 {SIMILARITY_THRESHOLD}）")
+            
+            if similarity >= SIMILARITY_THRESHOLD:
+                groups[group_id].append(candidate["path"])
+                to_remove.append(idx)
+        
+        # 移除已加入组的文件
+        for idx in sorted(to_remove, reverse=True):
+            removed = ungrouped.pop(idx)
+            print(f"  加入组 {group_id}：{removed['name']}")
+        
+        group_id += 1
+    
+    return groups
+
+
+def merge_objs_by_geometry(input_dir, output_dir):
+    """合并几何相似的OBJ文件，每组保留第一个文件"""
+    groups = group_objs_by_geometry(input_dir)
+    if not groups:
+        print("未生成任何分组")
+        return
+    
+    # 清空并创建输出目录
+    if os.path.exists(output_dir):
+        shutil.rmtree(output_dir)
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # 复制每组第一个文件
+    report = []
+    for group_id, paths in sorted(groups.items()):
+        rep_path = paths[0]
+        rep_name = os.path.basename(rep_path)
+        shutil.copy2(rep_path, os.path.join(output_dir, rep_name))
+        report.append({
+            "group_id": group_id,
+            "representative": rep_name,
+            "count": len(paths),
+            "files": [os.path.basename(p) for p in paths]
+        })
+        print(f"组 {group_id} 保留代表性文件：{rep_name}（共 {len(paths)} 个）")
+    
+    # 生成报告
+    with open(os.path.join(output_dir, "merge_report.txt"), "w", encoding="utf-8") as f:
+        f.write("OBJ零件合并报告（包围盒+ICP）\n")
+        f.write(f"日期：{pd.Timestamp.now().strftime('%Y-%m-%d %H:%M')}\n")
+        f.write(f"包围盒阈值：尺寸差异≤{BBOX_DIMENSION_TOLERANCE*100}%，体积差异≤{BBOX_VOLUME_TOLERANCE*100}%\n")
+        f.write(f"ICP相似度阈值：≥{SIMILARITY_THRESHOLD}\n\n")
+        for item in report:
+            f.write(f"组 {item['group_id']}：\n")
+            f.write(f"  代表性文件：{item['representative']}\n")
+            f.write(f"  包含文件数：{item['count']}\n")
+            f.write(f"  文件列表：{', '.join(item['files'])}\n\n")
+    
+    print(f"\n合并完成！输出目录：{output_dir}，报告：merge_report.txt")
+
+
+if __name__ == "__main__":
+    # 路径配置
+    INPUT_DIR = "/public/home/wangshuo/gap/assembly/data/obj_merged"  # 输入OBJ目录
+    OUTPUT_DIR = "/public/home/wangshuo/gap/assembly/data/obj_geo_merged"  # 输出目录
+    
+    # 检查依赖
+    try:
+        import pandas as pd
+        o3d.__version__
+    except (ImportError, AttributeError):
+        print("请先安装依赖：pip install open3d pandas")
+        exit(1)
+    
+    merge_objs_by_geometry(INPUT_DIR, OUTPUT_DIR)