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)