lipan
29f1d45d66
算法更新: delivery v1.2.1、平整度v1.5、测距v2.2、整车配准v1.3 适配器: vehicle-registration 新增、各算法 manifest 和 adapter 更新 后端: taskCadFiles/source 字段迁移、derivedPointClouds 注册、取消任务联动 前端: ElapsedTime 时长显示、两步算法流程、canRun 逻辑修复 文档: 新增5份算法性能分析报告
1351 lines
39 KiB
Python
1351 lines
39 KiB
Python
# 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"])
|