code/cube3d/training/group_objs_by_geometry.py

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)