# import os # import argparse # import numpy as np # import open3d as o3d # from OCC.Core.STEPControl import STEPControl_Reader # from OCC.Core.IFSelect import IFSelect_RetDone # from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh # from OCC.Core.TopExp import TopExp_Explorer # from OCC.Core.TopAbs import TopAbs_FACE # from OCC.Core.BRep import BRep_Tool # try: # from OCC.Core.TopoDS import topods # USE_TOPODS = "topods" # except ImportError: # from OCC.Core.TopoDS import TopoDS # USE_TOPODS = "TopoDS" # # ============================================================ # # STEP CAD -> Open3D Mesh # # ============================================================ # def read_step_file(step_path): # reader = STEPControl_Reader() # status = reader.ReadFile(step_path) # if status != IFSelect_RetDone: # raise RuntimeError(f"STEP 文件读取失败: {step_path}") # reader.TransferRoots() # shape = reader.OneShape() # print("[OK] STEP loaded:", step_path) # return shape # def occ_shape_to_numpy_mesh(shape): # vertices = [] # triangles = [] # explorer = TopExp_Explorer(shape, TopAbs_FACE) # face_count = 0 # triangle_count = 0 # while explorer.More(): # current_shape = explorer.Current() # if USE_TOPODS == "topods": # face = topods.Face(current_shape) # else: # face = TopoDS.Face(current_shape) # location = face.Location() # triangulation = BRep_Tool.Triangulation(face, location) # if triangulation is None: # explorer.Next() # continue # trsf = location.Transformation() # base_index = len(vertices) # for i in range(1, triangulation.NbNodes() + 1): # pnt = triangulation.Node(i) # pnt.Transform(trsf) # vertices.append([ # pnt.X(), # pnt.Y(), # pnt.Z() # ]) # for i in range(1, triangulation.NbTriangles() + 1): # tri = triangulation.Triangle(i) # n1, n2, n3 = tri.Get() # triangles.append([ # base_index + n1 - 1, # base_index + n2 - 1, # base_index + n3 - 1 # ]) # triangle_count += 1 # face_count += 1 # explorer.Next() # vertices = np.asarray(vertices, dtype=np.float64) # triangles = np.asarray(triangles, dtype=np.int32) # print("[OK] CAD mesh extracted") # print(" CAD faces:", face_count) # print(" vertices:", vertices.shape) # print(" triangles:", triangles.shape) # if vertices.shape[0] == 0 or triangles.shape[0] == 0: # raise RuntimeError("CAD 三角网格提取失败") # return vertices, triangles # def step_to_open3d_mesh(step_path, mesh_deflection=0.1): # shape = read_step_file(step_path) # print("[INFO] Meshing CAD...") # mesh_algo = BRepMesh_IncrementalMesh(shape, mesh_deflection) # mesh_algo.Perform() # vertices, triangles = occ_shape_to_numpy_mesh(shape) # mesh = o3d.geometry.TriangleMesh() # mesh.vertices = o3d.utility.Vector3dVector(vertices) # mesh.triangles = o3d.utility.Vector3iVector(triangles) # mesh.remove_duplicated_vertices() # mesh.remove_duplicated_triangles() # mesh.remove_degenerate_triangles() # mesh.compute_vertex_normals() # return mesh # # ============================================================ # # 基础工具 # # ============================================================ # def axis_to_index(axis): # axis = axis.lower() # if axis == "x": # return 0 # if axis == "y": # return 1 # if axis == "z": # return 2 # raise ValueError("axis 必须是 x / y / z") # def get_bounds(pcd): # points = np.asarray(pcd.points) # return points.min(axis=0), points.max(axis=0) # def crop_by_axis_ratio(pcd, axis="x", part="front", ratio=0.15): # """ # 按某个轴裁剪车头或车尾。 # axis: # 高铁长度方向，通常是 x # part: # front: 取最小端 # tail: 取最大端 # ratio: # 取前后多少比例，例如 0.15 表示前 15% 和后 15% # """ # axis_id = axis_to_index(axis) # points = np.asarray(pcd.points) # min_bound = points.min(axis=0) # max_bound = points.max(axis=0) # axis_min = min_bound[axis_id] # axis_max = max_bound[axis_id] # axis_len = axis_max - axis_min # if part == "front": # lower = axis_min # upper = axis_min + axis_len * ratio # elif part == "tail": # lower = axis_max - axis_len * ratio # upper = axis_max # else: # raise ValueError("part 必须是 front 或 tail") # crop_min = min_bound.copy() # crop_max = max_bound.copy() # crop_min[axis_id] = lower # crop_max[axis_id] = upper # bbox = o3d.geometry.AxisAlignedBoundingBox( # min_bound=crop_min, # max_bound=crop_max # ) # cropped = pcd.crop(bbox) # print(f"[OK] Crop {part}") # print(" axis:", axis) # print(" range:", lower, upper) # print(" points:", np.asarray(cropped.points).shape) # return cropped # def preprocess_point_cloud(pcd, voxel_size): # pcd_down = pcd.voxel_down_sample(voxel_size) # normal_radius = voxel_size * 2.0 # feature_radius = voxel_size * 5.0 # pcd_down.estimate_normals( # search_param=o3d.geometry.KDTreeSearchParamHybrid( # radius=normal_radius, # max_nn=30 # ) # ) # fpfh = o3d.pipelines.registration.compute_fpfh_feature( # pcd_down, # search_param=o3d.geometry.KDTreeSearchParamHybrid( # radius=feature_radius, # max_nn=100 # ) # ) # return pcd_down, fpfh # def global_registration(source_down, target_down, source_fpfh, target_fpfh, voxel_size): # distance_threshold = voxel_size * 1.5 # print("[INFO] RANSAC global registration") # print(" distance threshold:", distance_threshold) # result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching( # source_down, # target_down, # source_fpfh, # target_fpfh, # mutual_filter=True, # max_correspondence_distance=distance_threshold, # estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(False), # ransac_n=4, # checkers=[ # o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9), # o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold) # ], # criteria=o3d.pipelines.registration.RANSACConvergenceCriteria( # 100000, # 0.999 # ) # ) # print("[OK] RANSAC result") # print(" fitness:", result.fitness) # print(" rmse:", result.inlier_rmse) # print(" T:\n", result.transformation) # return result # def icp_refine(source, target, init_T, voxel_size, max_iteration=100): # distance_threshold = voxel_size * 0.5 # source.estimate_normals( # search_param=o3d.geometry.KDTreeSearchParamHybrid( # radius=voxel_size * 2.0, # max_nn=30 # ) # ) # target.estimate_normals( # search_param=o3d.geometry.KDTreeSearchParamHybrid( # radius=voxel_size * 2.0, # max_nn=30 # ) # ) # print("[INFO] ICP refinement") # print(" distance threshold:", distance_threshold) # result = o3d.pipelines.registration.registration_icp( # source, # target, # distance_threshold, # init_T, # o3d.pipelines.registration.TransformationEstimationPointToPlane(), # criteria=o3d.pipelines.registration.ICPConvergenceCriteria( # max_iteration=max_iteration # ) # ) # print("[OK] ICP result") # print(" fitness:", result.fitness) # print(" rmse:", result.inlier_rmse) # print(" T:\n", result.transformation) # return result # def register_local_region(source_local, target_local, voxel_size): # """ # 单个局部区域配准： # source_local: 扫描点云局部 # target_local: CAD采样点云局部 # """ # source_down, source_fpfh = preprocess_point_cloud(source_local, voxel_size) # target_down, target_fpfh = preprocess_point_cloud(target_local, voxel_size) # print("[INFO] source_down:", np.asarray(source_down.points).shape) # print("[INFO] target_down:", np.asarray(target_down.points).shape) # ransac_result = global_registration( # source_down, # target_down, # source_fpfh, # target_fpfh, # voxel_size # ) # icp_result = icp_refine( # source_local, # target_local, # ransac_result.transformation, # voxel_size # ) # return icp_result # # ============================================================ # # 变换融合 # # ============================================================ # def project_rotation_to_so3(R): # """ # 将平均后的矩阵重新投影到合法旋转矩阵 SO(3) # """ # U, S, Vt = np.linalg.svd(R) # R_proj = U @ Vt # if np.linalg.det(R_proj) < 0: # U[:, -1] *= -1 # R_proj = U @ Vt # return R_proj # def fuse_two_transformations(T_front, T_tail, weight_front=0.5, weight_tail=0.5): # """ # 融合车头和车尾配准结果。 # 注意： # 这里是简单加权平均旋转和平移。 # 对刚体配准来说，T_front 和 T_tail 理论上应该接近。 # 如果二者差异很大，说明前/后某一个局部配准可能失败。 # """ # weight_sum = weight_front + weight_tail # wf = weight_front / weight_sum # wt = weight_tail / weight_sum # R_front = T_front[:3, :3] # t_front = T_front[:3, 3] # R_tail = T_tail[:3, :3] # t_tail = T_tail[:3, 3] # R_avg = wf * R_front + wt * R_tail # R_final = project_rotation_to_so3(R_avg) # t_final = wf * t_front + wt * t_tail # T_final = np.eye(4) # T_final[:3, :3] = R_final # T_final[:3, 3] = t_final # return T_final # def transformation_difference(T1, T2): # """ # 查看前后两个配准结果差异。 # """ # R1 = T1[:3, :3] # R2 = T2[:3, :3] # R_delta = R1.T @ R2 # cos_theta = (np.trace(R_delta) - 1.0) / 2.0 # cos_theta = np.clip(cos_theta, -1.0, 1.0) # angle_rad = np.arccos(cos_theta) # angle_deg = np.degrees(angle_rad) # t1 = T1[:3, 3] # t2 = T2[:3, 3] # trans_diff = np.linalg.norm(t1 - t2) # return angle_deg, trans_diff # # ============================================================ # # 可视化 # # ============================================================ # def visualize_pair(source, target, T=None, title="Registration"): # source_vis = o3d.geometry.PointCloud(source) # target_vis = o3d.geometry.PointCloud(target) # source_vis.paint_uniform_color([1.0, 0.0, 0.0]) # target_vis.paint_uniform_color([0.0, 0.5, 1.0]) # if T is not None: # source_vis.transform(T) # o3d.visualization.draw_geometries( # [source_vis, target_vis], # window_name=title # ) # # ============================================================ # # 主流程：车头 + 车尾 + 整体融合 # # ============================================================ # def register_train_by_front_tail( # cad_path, # pointcloud_path, # output_dir="front_tail_registration_output", # axis="x", # crop_ratio=0.15, # mesh_deflection=0.1, # cad_sample_points=500000, # voxel_size=10.0, # scale_pointcloud=1.0, # final_icp=True, # visualize=True # ): # os.makedirs(output_dir, exist_ok=True) # # ------------------------------------------------------------ # # 1. CAD STEP -> CAD sampled point cloud # # ------------------------------------------------------------ # cad_mesh = step_to_open3d_mesh( # cad_path, # mesh_deflection=mesh_deflection # ) # print(f"[INFO] Sampling CAD mesh: {cad_sample_points} points") # cad_pcd = cad_mesh.sample_points_uniformly(cad_sample_points) # cad_pcd.estimate_normals( # search_param=o3d.geometry.KDTreeSearchParamHybrid( # radius=voxel_size * 2.0, # max_nn=30 # ) # ) # cad_sample_path = os.path.join(output_dir, "cad_sampled_whole.ply") # o3d.io.write_point_cloud(cad_sample_path, cad_pcd) # # ------------------------------------------------------------ # # 2. Load source point cloud # # ------------------------------------------------------------ # source_pcd = o3d.io.read_point_cloud(pointcloud_path) # if len(source_pcd.points) == 0: # raise RuntimeError(f"点云读取失败: {pointcloud_path}") # print("[OK] Source point cloud loaded") # print(" points:", np.asarray(source_pcd.points).shape) # if scale_pointcloud != 1.0: # source_pcd.scale(scale_pointcloud, center=(0, 0, 0)) # print("[INFO] Source scaled:", scale_pointcloud) # # ------------------------------------------------------------ # # 3. Crop front and tail # # ------------------------------------------------------------ # source_front = crop_by_axis_ratio( # source_pcd, # axis=axis, # part="front", # ratio=crop_ratio # ) # source_tail = crop_by_axis_ratio( # source_pcd, # axis=axis, # part="tail", # ratio=crop_ratio # ) # cad_front = crop_by_axis_ratio( # cad_pcd, # axis=axis, # part="front", # ratio=crop_ratio # ) # cad_tail = crop_by_axis_ratio( # cad_pcd, # axis=axis, # part="tail", # ratio=crop_ratio # ) # o3d.io.write_point_cloud(os.path.join(output_dir, "source_front.ply"), source_front) # o3d.io.write_point_cloud(os.path.join(output_dir, "source_tail.ply"), source_tail) # o3d.io.write_point_cloud(os.path.join(output_dir, "cad_front.ply"), cad_front) # o3d.io.write_point_cloud(os.path.join(output_dir, "cad_tail.ply"), cad_tail) # # ------------------------------------------------------------ # # 4. Front registration # # ------------------------------------------------------------ # print("\n========== Front Registration ==========") # front_result = register_local_region( # source_front, # cad_front, # voxel_size=voxel_size # ) # T_front = front_result.transformation # np.savetxt( # os.path.join(output_dir, "T_front.txt"), # T_front, # fmt="%.10f" # ) # # ------------------------------------------------------------ # # 5. Tail registration # # ------------------------------------------------------------ # print("\n========== Tail Registration ==========") # tail_result = register_local_region( # source_tail, # cad_tail, # voxel_size=voxel_size # ) # T_tail = tail_result.transformation # np.savetxt( # os.path.join(output_dir, "T_tail.txt"), # T_tail, # fmt="%.10f" # ) # # ------------------------------------------------------------ # # 6. Check difference between front and tail # # ------------------------------------------------------------ # angle_diff, trans_diff = transformation_difference(T_front, T_tail) # print("\n========== Front / Tail Difference ==========") # print(f"rotation difference: {angle_diff:.6f} deg") # print(f"translation difference: {trans_diff:.6f}") # # 根据 ICP 误差自动给权重 # # RMSE 越小，权重越大 # front_weight = 1.0 / (front_result.inlier_rmse + 1e-8) # tail_weight = 1.0 / (tail_result.inlier_rmse + 1e-8) # T_fused = fuse_two_transformations( # T_front, # T_tail, # weight_front=front_weight, # weight_tail=tail_weight # ) # np.savetxt( # os.path.join(output_dir, "T_fused_before_final_icp.txt"), # T_fused, # fmt="%.10f" # ) # print("\n========== Fused Transformation ==========") # print(T_fused) # # ------------------------------------------------------------ # # 7. Optional final whole-car low-resolution ICP # # ------------------------------------------------------------ # if final_icp: # print("\n========== Final Whole-Car ICP ==========") # source_whole_down = source_pcd.voxel_down_sample(voxel_size) # cad_whole_down = cad_pcd.voxel_down_sample(voxel_size) # final_result = icp_refine( # source_whole_down, # cad_whole_down, # T_fused, # voxel_size=voxel_size, # max_iteration=50 # ) # T_final = final_result.transformation # final_fitness = final_result.fitness # final_rmse = final_result.inlier_rmse # else: # T_final = T_fused # final_fitness = None # final_rmse = None # np.savetxt( # os.path.join(output_dir, "T_final.txt"), # T_final, # fmt="%.10f" # ) # # ------------------------------------------------------------ # # 8. Save aligned whole point cloud # # ------------------------------------------------------------ # aligned_source = o3d.geometry.PointCloud(source_pcd) # aligned_source.transform(T_final) # aligned_path = os.path.join(output_dir, "aligned_whole_pointcloud.ply") # o3d.io.write_point_cloud(aligned_path, aligned_source) # print("[OK] Saved aligned whole point cloud:") # print(" ", aligned_path) # # ------------------------------------------------------------ # # 9. Visualization # # ------------------------------------------------------------ # if visualize: # source_vis = aligned_source.voxel_down_sample(voxel_size) # cad_vis = cad_pcd.voxel_down_sample(voxel_size) # source_vis.paint_uniform_color([1.0, 0.0, 0.0]) # cad_vis.paint_uniform_color([0.0, 0.5, 1.0]) # o3d.visualization.draw_geometries( # [source_vis, cad_vis], # window_name="Final Whole Train Registration: red=aligned source, blue=CAD" # ) # return { # "T_front": T_front, # "T_tail": T_tail, # "T_fused": T_fused, # "T_final": T_final, # "front_fitness": front_result.fitness, # "front_rmse": front_result.inlier_rmse, # "tail_fitness": tail_result.fitness, # "tail_rmse": tail_result.inlier_rmse, # "front_tail_rotation_diff_deg": angle_diff, # "front_tail_translation_diff": trans_diff, # "final_fitness": final_fitness, # "final_rmse": final_rmse, # "aligned_path": aligned_path # } # # ============================================================ # # 命令行入口 # # ============================================================ # if __name__ == "__main__": # parser = argparse.ArgumentParser( # description="高铁点云与标准 CAD 的前后局部配准 + 整体刚体融合" # ) # parser.add_argument( # "--cad", # "-c", # default="datasets/CAD/11.step", # help="标准 CAD STEP 文件路径" # ) # parser.add_argument( # "--pointcloud", # "-p", # default="datasets/Noise_PointCloud/11_noisy.ply", # help="待配准点云 PLY 文件路径" # ) # parser.add_argument( # "--output_dir", # "-o", # default="front_tail_registration_output", # help="输出目录" # ) # parser.add_argument( # "--axis", # default="x", # choices=["x", "y", "z"], # help="高铁长度方向。默认 x" # ) # parser.add_argument( # "--crop_ratio", # type=float, # default=0.15, # help="车头/车尾各取整车长度的比例。默认 0.15" # ) # parser.add_argument( # "--mesh_deflection", # type=float, # default=0.1, # help="CAD 三角化精度。默认 0.1" # ) # parser.add_argument( # "--cad_sample_points", # type=int, # default=500000, # help="CAD 表面采样点数。默认 500000" # ) # parser.add_argument( # "--voxel_size", # type=float, # default=10.0, # help="配准体素大小。高铁大尺度建议 5~30。默认 10" # ) # parser.add_argument( # "--scale_pointcloud", # type=float, # default=1.0, # help="点云缩放。如果点云是米，CAD 是毫米，设为 1000" # ) # parser.add_argument( # "--no_final_icp", # action="store_true", # help="不做最后整车 ICP 微调" # ) # parser.add_argument( # "--no_visualize", # action="store_true", # help="不显示可视化窗口" # ) # args = parser.parse_args() # result = register_train_by_front_tail( # cad_path=args.cad, # pointcloud_path=args.pointcloud, # output_dir=args.output_dir, # axis=args.axis, # crop_ratio=args.crop_ratio, # mesh_deflection=args.mesh_deflection, # cad_sample_points=args.cad_sample_points, # voxel_size=args.voxel_size, # scale_pointcloud=args.scale_pointcloud, # final_icp=not args.no_final_icp, # visualize=not args.no_visualize # ) # print("\n========== Final Result ==========") # print("front fitness:", result["front_fitness"]) # print("front rmse:", result["front_rmse"]) # print("tail fitness:", result["tail_fitness"]) # print("tail rmse:", result["tail_rmse"]) # print("front-tail rotation diff deg:", result["front_tail_rotation_diff_deg"]) # print("front-tail translation diff:", result["front_tail_translation_diff"]) # print("final fitness:", result["final_fitness"]) # print("final rmse:", result["final_rmse"]) # print("aligned path:", result["aligned_path"]) # print("T_final:") # print(result["T_final"]) import os import argparse import numpy as np import open3d as o3d from OCC.Core.STEPControl import STEPControl_Reader from OCC.Core.IFSelect import IFSelect_RetDone from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh from OCC.Core.TopExp import TopExp_Explorer from OCC.Core.TopAbs import TopAbs_FACE from OCC.Core.BRep import BRep_Tool try: from OCC.Core.TopoDS import topods USE_TOPODS = "topods" except ImportError: from OCC.Core.TopoDS import TopoDS USE_TOPODS = "TopoDS" # ============================================================ # STEP CAD -> Open3D Mesh # ============================================================ def read_step_file(step_path): reader = STEPControl_Reader() status = reader.ReadFile(step_path) if status != IFSelect_RetDone: raise RuntimeError(f"STEP 文件读取失败: {step_path}") reader.TransferRoots() shape = reader.OneShape() print("[OK] STEP loaded:", step_path) return shape def occ_shape_to_numpy_mesh(shape): vertices = [] triangles = [] explorer = TopExp_Explorer(shape, TopAbs_FACE) face_count = 0 tri_count = 0 while explorer.More(): current_shape = explorer.Current() if USE_TOPODS == "topods": face = topods.Face(current_shape) else: face = TopoDS.Face(current_shape) location = face.Location() triangulation = BRep_Tool.Triangulation(face, location) if triangulation is None: explorer.Next() continue trsf = location.Transformation() base_index = len(vertices) for i in range(1, triangulation.NbNodes() + 1): pnt = triangulation.Node(i) pnt.Transform(trsf) vertices.append([ pnt.X(), pnt.Y(), pnt.Z() ]) for i in range(1, triangulation.NbTriangles() + 1): tri = triangulation.Triangle(i) n1, n2, n3 = tri.Get() triangles.append([ base_index + n1 - 1, base_index + n2 - 1, base_index + n3 - 1 ]) tri_count += 1 face_count += 1 explorer.Next() vertices = np.asarray(vertices, dtype=np.float64) triangles = np.asarray(triangles, dtype=np.int32) print("[OK] CAD mesh extracted") print(" CAD faces:", face_count) print(" mesh vertices:", vertices.shape) print(" mesh triangles:", triangles.shape) if vertices.shape[0] == 0 or triangles.shape[0] == 0: raise RuntimeError("CAD 三角网格提取失败，请检查 STEP 文件") return vertices, triangles def step_to_open3d_mesh(step_path, mesh_deflection=0.1): shape = read_step_file(step_path) print("[INFO] Meshing CAD...") mesh_algo = BRepMesh_IncrementalMesh(shape, mesh_deflection) mesh_algo.Perform() vertices, triangles = occ_shape_to_numpy_mesh(shape) mesh = o3d.geometry.TriangleMesh() mesh.vertices = o3d.utility.Vector3dVector(vertices) mesh.triangles = o3d.utility.Vector3iVector(triangles) mesh.remove_duplicated_vertices() mesh.remove_duplicated_triangles() mesh.remove_degenerate_triangles() mesh.compute_vertex_normals() return mesh # ============================================================ # 工具函数 # ============================================================ def axis_to_index(axis): axis = axis.lower() if axis == "x": return 0 if axis == "y": return 1 if axis == "z": return 2 raise ValueError("axis 必须是 x / y / z") def crop_end_region_by_ratio(pcd, axis="x", part="front", ratio=0.08): """ 按长度方向裁剪前端/后端区域。 axis: 高铁长度方向，一般是 x part: front: 轴最小端 tail: 轴最大端 ratio: 取前后端占整体长度的比例，例如 0.08 表示取前 8% 和后 8% """ points = np.asarray(pcd.points) axis_id = axis_to_index(axis) min_bound = points.min(axis=0) max_bound = points.max(axis=0) axis_min = min_bound[axis_id] axis_max = max_bound[axis_id] axis_len = axis_max - axis_min if part == "front": lower = axis_min upper = axis_min + axis_len * ratio elif part == "tail": lower = axis_max - axis_len * ratio upper = axis_max else: raise ValueError("part 必须是 front 或 tail") crop_min = min_bound.copy() crop_max = max_bound.copy() crop_min[axis_id] = lower crop_max[axis_id] = upper bbox = o3d.geometry.AxisAlignedBoundingBox( min_bound=crop_min, max_bound=crop_max ) cropped = pcd.crop(bbox) print(f"[OK] Crop {part}") print(" range:", lower, upper) print(" points:", np.asarray(cropped.points).shape) return cropped def sample_farthest_points_simple(pcd, num_points=5000): """ 为了避免百万点太多，这里随机采样少量点。 如果点数不足，则直接返回原点云。 """ points = np.asarray(pcd.points) n = points.shape[0] if n <= num_points: return pcd indices = np.random.choice(n, size=num_points, replace=False) sampled = pcd.select_by_index(indices.tolist()) return sampled def compute_centroid(pcd): points = np.asarray(pcd.points) if points.shape[0] == 0: raise RuntimeError("点云为空，无法计算质心") return points.mean(axis=0) def build_translation_matrix(t): T = np.eye(4, dtype=np.float64) T[:3, 3] = t return T def visualize_alignment(source_pcd, cad_pcd, title="Translation Alignment Result"): source_vis = o3d.geometry.PointCloud(source_pcd) cad_vis = o3d.geometry.PointCloud(cad_pcd) source_vis.paint_uniform_color([1.0, 0.0, 0.0]) cad_vis.paint_uniform_color([0.0, 0.5, 1.0]) o3d.visualization.draw_geometries( [source_vis, cad_vis], window_name=title ) # ============================================================ # 只估计平移的配准 # ============================================================ def resolve_pointcloud_path(pointcloud_path): if os.path.isdir(pointcloud_path): preferred = os.path.join(pointcloud_path, "shifted_whole_pointcloud.ply") if os.path.exists(preferred): return preferred candidates = [ os.path.join(pointcloud_path, name) for name in sorted(os.listdir(pointcloud_path)) if name.lower().endswith(".ply") ] if len(candidates) == 1: return candidates[0] if not candidates: raise FileNotFoundError(f"No PLY point cloud found in directory: {pointcloud_path}") raise ValueError( "Multiple PLY files found in directory; specify one file explicitly: " + ", ".join(os.path.basename(path) for path in candidates) ) return pointcloud_path def translate_align_by_front_tail( cad_path, pointcloud_path, output_dir="results", axis="x", crop_ratio=0.08, mesh_deflection=0.1, cad_sample_points=300000, local_sample_points=10000, scale_pointcloud=1.0, visualize=True ): os.makedirs(output_dir, exist_ok=True) # ------------------------------------------------------------ # 1. CAD STEP 转 mesh，再采样成 CAD 点云 # ------------------------------------------------------------ cad_mesh = step_to_open3d_mesh( cad_path, mesh_deflection=mesh_deflection ) print(f"[INFO] Sampling CAD mesh: {cad_sample_points} points") cad_pcd = cad_mesh.sample_points_uniformly( number_of_points=cad_sample_points ) cad_pcd_path = os.path.join(output_dir, "cad_sampled.ply") o3d.io.write_point_cloud(cad_pcd_path, cad_pcd) # ------------------------------------------------------------ # 2. 读取扫描点云 # ------------------------------------------------------------ pointcloud_path = resolve_pointcloud_path(pointcloud_path) source_pcd = o3d.io.read_point_cloud(pointcloud_path) if len(source_pcd.points) == 0: raise RuntimeError(f"点云读取失败: {pointcloud_path}") print("[OK] Source point cloud loaded:", pointcloud_path) print(" source points:", np.asarray(source_pcd.points).shape) if scale_pointcloud != 1.0: source_pcd.scale(scale_pointcloud, center=(0, 0, 0)) print("[INFO] Source scaled:", scale_pointcloud) # ------------------------------------------------------------ # 3. 裁剪 CAD 前端 / 后端 # ------------------------------------------------------------ cad_front = crop_end_region_by_ratio( cad_pcd, axis=axis, part="front", ratio=crop_ratio ) cad_tail = crop_end_region_by_ratio( cad_pcd, axis=axis, part="tail", ratio=crop_ratio ) # ------------------------------------------------------------ # 4. 裁剪扫描点云前端 / 后端 # ------------------------------------------------------------ source_front = crop_end_region_by_ratio( source_pcd, axis=axis, part="front", ratio=crop_ratio ) source_tail = crop_end_region_by_ratio( source_pcd, axis=axis, part="tail", ratio=crop_ratio ) # ------------------------------------------------------------ # 5. 前后区域分别随机取少量点 # ------------------------------------------------------------ cad_front_sample = sample_farthest_points_simple( cad_front, num_points=local_sample_points ) cad_tail_sample = sample_farthest_points_simple( cad_tail, num_points=local_sample_points ) source_front_sample = sample_farthest_points_simple( source_front, num_points=local_sample_points ) source_tail_sample = sample_farthest_points_simple( source_tail, num_points=local_sample_points ) # 保存局部采样结果，方便检查 o3d.io.write_point_cloud(os.path.join(output_dir, "cad_front_sample.ply"), cad_front_sample) o3d.io.write_point_cloud(os.path.join(output_dir, "cad_tail_sample.ply"), cad_tail_sample) o3d.io.write_point_cloud(os.path.join(output_dir, "source_front_sample.ply"), source_front_sample) o3d.io.write_point_cloud(os.path.join(output_dir, "source_tail_sample.ply"), source_tail_sample) # ------------------------------------------------------------ # 6. 计算前端和后端质心 # ------------------------------------------------------------ cad_front_center = compute_centroid(cad_front_sample) cad_tail_center = compute_centroid(cad_tail_sample) source_front_center = compute_centroid(source_front_sample) source_tail_center = compute_centroid(source_tail_sample) print("\n========== Centroids ==========") print("CAD front center: ", cad_front_center) print("CAD tail center: ", cad_tail_center) print("Source front center:", source_front_center) print("Source tail center: ", source_tail_center) # ------------------------------------------------------------ # 7. 分别根据车头和车尾估计平移 # ------------------------------------------------------------ t_front = cad_front_center - source_front_center t_tail = cad_tail_center - source_tail_center print("\n========== Translation From Front / Tail ==========") print("t_front:", t_front) print("t_tail: ", t_tail) print("front-tail translation difference:", np.linalg.norm(t_front - t_tail)) # ------------------------------------------------------------ # 8. 融合平移 # ------------------------------------------------------------ t_final = 0.5 * t_front + 0.5 * t_tail T_final = build_translation_matrix(t_final) print("\n========== Final Translation ==========") print("t_final:", t_final) print("T_final:") print(T_final) # ------------------------------------------------------------ # 9. 应用平移到整个点云 # ------------------------------------------------------------ aligned_pcd = o3d.geometry.PointCloud(source_pcd) aligned_pcd.transform(T_final) aligned_path = os.path.join(output_dir, "aligned_pointcloud_translation_only.ply") matrix_path = os.path.join(output_dir, "T_translation_only.txt") o3d.io.write_point_cloud(aligned_path, aligned_pcd) np.savetxt(matrix_path, T_final, fmt="%.10f") print("\n[OK] Saved aligned point cloud:", aligned_path) print("[OK] Saved transformation matrix:", matrix_path) # ------------------------------------------------------------ # 10. 可视化 # ------------------------------------------------------------ if visualize: # 为了不卡，显示时降采样 source_vis = aligned_pcd.voxel_down_sample(voxel_size=max(crop_ratio, 1.0)) cad_vis = cad_pcd.voxel_down_sample(voxel_size=max(crop_ratio, 1.0)) source_vis.paint_uniform_color([1.0, 0.0, 0.0]) cad_vis.paint_uniform_color([0.0, 0.5, 1.0]) o3d.visualization.draw_geometries( [source_vis, cad_vis], window_name="Translation Only Alignment: red=aligned source, blue=CAD" ) return { "translation": t_final, "T_final": T_final, "t_front": t_front, "t_tail": t_tail, "front_tail_translation_difference": float(np.linalg.norm(t_front - t_tail)), "aligned_path": aligned_path, "matrix_path": matrix_path } # ============================================================ # 命令行入口 # ============================================================ if __name__ == "__main__": parser = argparse.ArgumentParser( description="高铁点云与 CAD 的仅平移配准：取车头车尾少量点估计整体平移" ) parser.add_argument( "--cad", "-c", default="datasets/CAD/M02.stp", help="CAD STEP 文件路径" ) parser.add_argument( "--pointcloud", "-p", default="datasets/shift_test_output/", help="待配准点云 PLY 文件路径" ) parser.add_argument( "--output_dir", "-o", default="results", help="输出目录" ) parser.add_argument( "--axis", default="x", choices=["x", "y", "z"], help="高铁长度方向，默认 x" ) parser.add_argument( "--crop_ratio", type=float, default=0.01, help="车头/车尾各取总长度的比例，默认 0.01" ) parser.add_argument( "--mesh_deflection", type=float, default=0.1, help="CAD 三角化精度，默认 0.1" ) parser.add_argument( "--cad_sample_points", type=int, default=30000, help="从 CAD mesh 采样的点数，默认 30000" ) parser.add_argument( "--local_sample_points", type=int, default=1000, help="车头/车尾各随机采样多少点用于计算质心，默认 1000" ) parser.add_argument( "--scale_pointcloud", type=float, default=1.0, help="点云缩放。如果点云是米，CAD 是毫米，设为 1000" ) parser.add_argument( "--no_visualize", action="store_true", help="不显示可视化窗口" ) args = parser.parse_args() result = translate_align_by_front_tail( cad_path=args.cad, pointcloud_path=args.pointcloud, output_dir=args.output_dir, axis=args.axis, crop_ratio=args.crop_ratio, mesh_deflection=args.mesh_deflection, cad_sample_points=args.cad_sample_points, local_sample_points=args.local_sample_points, scale_pointcloud=args.scale_pointcloud, visualize=not args.no_visualize ) print("\n========== Result ==========") print("translation:", result["translation"]) print("front-tail translation difference:", result["front_tail_translation_difference"]) print("aligned path:", result["aligned_path"]) print("matrix path:", result["matrix_path"])