feat: 算法升级、整车配准集成与性能分析文档
算法更新: delivery v1.2.1、平整度v1.5、测距v2.2、整车配准v1.3 适配器: vehicle-registration 新增、各算法 manifest 和 adapter 更新 后端: taskCadFiles/source 字段迁移、derivedPointClouds 注册、取消任务联动 前端: ElapsedTime 时长显示、两步算法流程、canRun 逻辑修复 文档: 新增5份算法性能分析报告
This commit is contained in:
972
alg/整车配准v1.3/translate_align_pointcloud_to_cad.py
Normal file
972
alg/整车配准v1.3/translate_align_pointcloud_to_cad.py
Normal file
@@ -0,0 +1,972 @@
|
||||
import os
|
||||
import argparse
|
||||
import time
|
||||
import hashlib
|
||||
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"
|
||||
|
||||
|
||||
def format_elapsed(seconds):
|
||||
if seconds < 1.0:
|
||||
return f"{seconds * 1000.0:.1f} ms"
|
||||
|
||||
minutes, seconds = divmod(seconds, 60.0)
|
||||
if minutes < 1.0:
|
||||
return f"{seconds:.2f} s"
|
||||
|
||||
hours, minutes = divmod(minutes, 60.0)
|
||||
if hours < 1.0:
|
||||
return f"{int(minutes)} min {seconds:.2f} s"
|
||||
|
||||
return f"{int(hours)} h {int(minutes)} min {seconds:.2f} s"
|
||||
|
||||
|
||||
def log_elapsed(label, start_time):
|
||||
elapsed = time.perf_counter() - start_time
|
||||
print(f"[TIME] {label}: {format_elapsed(elapsed)}")
|
||||
return elapsed
|
||||
|
||||
|
||||
def build_cad_sample_cache_path(cad_path, output_dir, mesh_deflection, cad_sample_points):
|
||||
stat = os.stat(cad_path)
|
||||
cache_key = "|".join([
|
||||
os.path.abspath(cad_path),
|
||||
str(stat.st_size),
|
||||
str(stat.st_mtime_ns),
|
||||
f"{mesh_deflection:.12g}",
|
||||
str(cad_sample_points),
|
||||
])
|
||||
digest = hashlib.sha1(cache_key.encode("utf-8")).hexdigest()[:16]
|
||||
cad_name = os.path.splitext(os.path.basename(cad_path))[0] or "cad"
|
||||
return os.path.join(output_dir, "cad_cache", f"{cad_name}_{digest}.ply")
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 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, rng=None):
|
||||
"""
|
||||
为了避免百万点太多，这里随机采样少量点。
|
||||
如果点数不足，则直接返回原点云。
|
||||
"""
|
||||
|
||||
points = np.asarray(pcd.points)
|
||||
n = points.shape[0]
|
||||
|
||||
if n <= num_points:
|
||||
return pcd
|
||||
|
||||
if rng is None:
|
||||
rng = np.random.default_rng()
|
||||
|
||||
indices = rng.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 compute_alignment_error_stats(source_pcd, cad_pcd, sample_points=100000, rng=None):
|
||||
if sample_points is not None and sample_points > 0:
|
||||
eval_pcd = sample_farthest_points_simple(
|
||||
source_pcd,
|
||||
num_points=sample_points,
|
||||
rng=rng
|
||||
)
|
||||
else:
|
||||
eval_pcd = source_pcd
|
||||
|
||||
distances = np.asarray(
|
||||
eval_pcd.compute_point_cloud_distance(cad_pcd),
|
||||
dtype=np.float64
|
||||
)
|
||||
|
||||
if distances.shape[0] == 0:
|
||||
raise RuntimeError("No points available for alignment error statistics")
|
||||
|
||||
squared = distances * distances
|
||||
|
||||
return {
|
||||
"sample_count": int(distances.shape[0]),
|
||||
"mean": float(np.mean(distances)),
|
||||
"rmse": float(np.sqrt(np.mean(squared))),
|
||||
"median": float(np.median(distances)),
|
||||
"p95": float(np.percentile(distances, 95)),
|
||||
"p99": float(np.percentile(distances, 99)),
|
||||
"max": float(np.max(distances)),
|
||||
}
|
||||
|
||||
|
||||
def print_alignment_error_stats(stats):
|
||||
print("\n========== Alignment Error ==========")
|
||||
print("sample count:", stats["sample_count"])
|
||||
print("mean: ", stats["mean"])
|
||||
print("rmse: ", stats["rmse"])
|
||||
print("median: ", stats["median"])
|
||||
print("p95: ", stats["p95"])
|
||||
print("p99: ", stats["p99"])
|
||||
print("max: ", stats["max"])
|
||||
|
||||
|
||||
def save_alignment_error_stats(path, stats):
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
for key in ["sample_count", "mean", "rmse", "median", "p95", "p99", "max"]:
|
||||
f.write(f"{key}: {stats[key]}\n")
|
||||
|
||||
|
||||
def write_point_cloud_float32_xyz(path, pcd, chunk_size=1_000_000):
|
||||
points = np.asarray(pcd.points)
|
||||
if points.shape[0] == 0:
|
||||
raise RuntimeError(f"Cannot write empty point cloud: {path}")
|
||||
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
|
||||
header = (
|
||||
"ply\n"
|
||||
"format binary_little_endian 1.0\n"
|
||||
f"element vertex {points.shape[0]}\n"
|
||||
"property float x\n"
|
||||
"property float y\n"
|
||||
"property float z\n"
|
||||
"end_header\n"
|
||||
)
|
||||
|
||||
with open(path, "wb") as f:
|
||||
f.write(header.encode("ascii"))
|
||||
for start in range(0, points.shape[0], chunk_size):
|
||||
end = min(start + chunk_size, points.shape[0])
|
||||
f.write(np.asarray(points[start:end], dtype="<f4").tobytes(order="C"))
|
||||
|
||||
|
||||
def build_translation_matrix(t):
|
||||
T = np.eye(4, dtype=np.float64)
|
||||
T[:3, 3] = t
|
||||
return T
|
||||
|
||||
|
||||
def run_icp_refinement(
|
||||
source_pcd,
|
||||
cad_pcd,
|
||||
initial_transform,
|
||||
sample_points=200000,
|
||||
voxel_size=0.0,
|
||||
max_correspondence_distance=0.0,
|
||||
max_iteration=50,
|
||||
rng=None
|
||||
):
|
||||
if sample_points is not None and sample_points > 0:
|
||||
source_icp = sample_farthest_points_simple(
|
||||
source_pcd,
|
||||
num_points=sample_points,
|
||||
rng=rng
|
||||
)
|
||||
else:
|
||||
source_icp = o3d.geometry.PointCloud(source_pcd)
|
||||
|
||||
source_icp.transform(initial_transform)
|
||||
target_icp = cad_pcd
|
||||
|
||||
if voxel_size is not None and voxel_size > 0:
|
||||
source_icp = source_icp.voxel_down_sample(voxel_size=voxel_size)
|
||||
target_icp = cad_pcd.voxel_down_sample(voxel_size=voxel_size)
|
||||
|
||||
source_count = len(source_icp.points)
|
||||
target_count = len(target_icp.points)
|
||||
if source_count == 0 or target_count == 0:
|
||||
raise RuntimeError("ICP point cloud is empty after sampling/downsampling")
|
||||
|
||||
if max_correspondence_distance is None or max_correspondence_distance <= 0:
|
||||
initial_distances = np.asarray(
|
||||
source_icp.compute_point_cloud_distance(target_icp),
|
||||
dtype=np.float64
|
||||
)
|
||||
p95 = float(np.percentile(initial_distances, 95))
|
||||
max_correspondence_distance = max(p95 * 2.0, 1e-6)
|
||||
print("[INFO] ICP auto max correspondence distance:", max_correspondence_distance)
|
||||
|
||||
result = o3d.pipelines.registration.registration_icp(
|
||||
source_icp,
|
||||
target_icp,
|
||||
max_correspondence_distance,
|
||||
np.eye(4, dtype=np.float64),
|
||||
o3d.pipelines.registration.TransformationEstimationPointToPoint(),
|
||||
o3d.pipelines.registration.ICPConvergenceCriteria(
|
||||
max_iteration=max_iteration
|
||||
)
|
||||
)
|
||||
|
||||
icp_transform = result.transformation
|
||||
final_transform = icp_transform @ initial_transform
|
||||
|
||||
stats = {
|
||||
"fitness": float(result.fitness),
|
||||
"inlier_rmse": float(result.inlier_rmse),
|
||||
"max_correspondence_distance": float(max_correspondence_distance),
|
||||
"source_points": int(source_count),
|
||||
"target_points": int(target_count),
|
||||
"transform": icp_transform,
|
||||
}
|
||||
|
||||
return final_transform, stats
|
||||
|
||||
|
||||
def print_icp_stats(stats):
|
||||
print("\n========== ICP Refinement ==========")
|
||||
print("source points:", stats["source_points"])
|
||||
print("target points:", stats["target_points"])
|
||||
print("max correspondence distance:", stats["max_correspondence_distance"])
|
||||
print("fitness:", stats["fitness"])
|
||||
print("inlier rmse:", stats["inlier_rmse"])
|
||||
print("ICP delta transform:")
|
||||
print(stats["transform"])
|
||||
|
||||
|
||||
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,
|
||||
use_cad_cache=True,
|
||||
save_debug_samples=False,
|
||||
random_seed=0,
|
||||
error_sample_points=100000,
|
||||
compute_error_stats=True,
|
||||
save_aligned_float32=True,
|
||||
use_icp=True,
|
||||
icp_sample_points=200000,
|
||||
icp_voxel_size=0.0,
|
||||
icp_max_correspondence_distance=0.0,
|
||||
icp_max_iteration=50
|
||||
):
|
||||
total_start = time.perf_counter()
|
||||
stage_start = total_start
|
||||
rng = np.random.default_rng(random_seed) if random_seed is not None else np.random.default_rng()
|
||||
|
||||
if random_seed is not None:
|
||||
try:
|
||||
o3d.utility.random.seed(int(random_seed))
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# 1. CAD STEP 转 mesh，再采样成 CAD 点云
|
||||
# ------------------------------------------------------------
|
||||
|
||||
cad_pcd = None
|
||||
cad_pcd_path = os.path.join(output_dir, "cad_sampled.ply")
|
||||
|
||||
if use_cad_cache:
|
||||
cad_pcd_path = build_cad_sample_cache_path(
|
||||
cad_path,
|
||||
output_dir,
|
||||
mesh_deflection,
|
||||
cad_sample_points
|
||||
)
|
||||
if os.path.exists(cad_pcd_path):
|
||||
cad_pcd = o3d.io.read_point_cloud(cad_pcd_path)
|
||||
if len(cad_pcd.points) > 0:
|
||||
print("[OK] CAD sample cache loaded:", cad_pcd_path)
|
||||
else:
|
||||
print("[WARN] CAD sample cache is empty, rebuilding:", cad_pcd_path)
|
||||
cad_pcd = None
|
||||
|
||||
if cad_pcd is None:
|
||||
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
|
||||
)
|
||||
|
||||
os.makedirs(os.path.dirname(cad_pcd_path), exist_ok=True)
|
||||
o3d.io.write_point_cloud(cad_pcd_path, cad_pcd)
|
||||
print("[OK] CAD sample saved:", cad_pcd_path)
|
||||
|
||||
log_elapsed("CAD sample load/build", stage_start)
|
||||
stage_start = time.perf_counter()
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# 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)
|
||||
log_elapsed("source load", stage_start)
|
||||
stage_start = time.perf_counter()
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# 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
|
||||
)
|
||||
log_elapsed("crop regions", stage_start)
|
||||
stage_start = time.perf_counter()
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# 5. 前后区域分别随机取少量点
|
||||
# ------------------------------------------------------------
|
||||
|
||||
cad_front_sample = sample_farthest_points_simple(
|
||||
cad_front,
|
||||
num_points=local_sample_points,
|
||||
rng=rng
|
||||
)
|
||||
|
||||
cad_tail_sample = sample_farthest_points_simple(
|
||||
cad_tail,
|
||||
num_points=local_sample_points,
|
||||
rng=rng
|
||||
)
|
||||
|
||||
source_front_sample = sample_farthest_points_simple(
|
||||
source_front,
|
||||
num_points=local_sample_points,
|
||||
rng=rng
|
||||
)
|
||||
|
||||
source_tail_sample = sample_farthest_points_simple(
|
||||
source_tail,
|
||||
num_points=local_sample_points,
|
||||
rng=rng
|
||||
)
|
||||
|
||||
# 保存局部采样结果，方便检查
|
||||
if save_debug_samples:
|
||||
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)
|
||||
print("[OK] Debug sample point clouds saved")
|
||||
else:
|
||||
print("[INFO] Debug sample point clouds skipped")
|
||||
log_elapsed("local sampling", stage_start)
|
||||
stage_start = time.perf_counter()
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# 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========== Initial Translation ==========")
|
||||
print("t_initial:", t_final)
|
||||
print("T_initial:")
|
||||
print(T_final)
|
||||
log_elapsed("translation estimate", stage_start)
|
||||
stage_start = time.perf_counter()
|
||||
|
||||
icp_stats = None
|
||||
if use_icp:
|
||||
T_final, icp_stats = run_icp_refinement(
|
||||
source_pcd,
|
||||
cad_pcd,
|
||||
T_final,
|
||||
sample_points=icp_sample_points,
|
||||
voxel_size=icp_voxel_size,
|
||||
max_correspondence_distance=icp_max_correspondence_distance,
|
||||
max_iteration=icp_max_iteration,
|
||||
rng=rng
|
||||
)
|
||||
print_icp_stats(icp_stats)
|
||||
print("\n========== Final Rigid Transform ==========")
|
||||
print("final translation:", T_final[:3, 3])
|
||||
print("T_final:")
|
||||
print(T_final)
|
||||
else:
|
||||
print("[INFO] ICP refinement skipped")
|
||||
|
||||
log_elapsed("ICP refinement", stage_start)
|
||||
stage_start = time.perf_counter()
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# 9. 应用平移到整个点云
|
||||
# ------------------------------------------------------------
|
||||
|
||||
aligned_pcd = source_pcd
|
||||
aligned_pcd.transform(T_final)
|
||||
|
||||
if use_icp:
|
||||
aligned_path = os.path.join(output_dir, "aligned_pointcloud.ply")
|
||||
matrix_path = os.path.join(output_dir, "T_rigid_icp.txt")
|
||||
else:
|
||||
aligned_path = os.path.join(output_dir, "aligned_pointcloud_translation_only.ply")
|
||||
matrix_path = os.path.join(output_dir, "T_translation_only.txt")
|
||||
|
||||
if save_aligned_float32:
|
||||
write_point_cloud_float32_xyz(aligned_path, aligned_pcd)
|
||||
print("[INFO] Aligned point cloud saved as compact float32 XYZ PLY")
|
||||
else:
|
||||
o3d.io.write_point_cloud(aligned_path, aligned_pcd, write_ascii=False)
|
||||
print("[INFO] Aligned point cloud saved with Open3D default writer")
|
||||
np.savetxt(matrix_path, T_final, fmt="%.10f")
|
||||
|
||||
print("\n[OK] Saved aligned point cloud:", aligned_path)
|
||||
print("[OK] Saved transformation matrix:", matrix_path)
|
||||
log_elapsed("save outputs", stage_start)
|
||||
stage_start = time.perf_counter()
|
||||
|
||||
error_stats = None
|
||||
error_stats_path = os.path.join(output_dir, "alignment_error_stats.txt")
|
||||
|
||||
if compute_error_stats:
|
||||
error_stats = compute_alignment_error_stats(
|
||||
aligned_pcd,
|
||||
cad_pcd,
|
||||
sample_points=error_sample_points,
|
||||
rng=rng
|
||||
)
|
||||
print_alignment_error_stats(error_stats)
|
||||
save_alignment_error_stats(error_stats_path, error_stats)
|
||||
print("[OK] Saved alignment error stats:", error_stats_path)
|
||||
else:
|
||||
error_stats_path = None
|
||||
print("[INFO] Alignment error stats skipped")
|
||||
|
||||
log_elapsed("alignment error stats", stage_start)
|
||||
|
||||
# ------------------------------------------------------------
|
||||
# 10. 可视化
|
||||
# ------------------------------------------------------------
|
||||
|
||||
if visualize:
|
||||
visualize_start = time.perf_counter()
|
||||
# 为了不卡，显示时降采样
|
||||
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"
|
||||
)
|
||||
log_elapsed("visualization", visualize_start)
|
||||
|
||||
total_elapsed = log_elapsed("total", total_start)
|
||||
|
||||
return {
|
||||
"translation": T_final[:3, 3],
|
||||
"initial_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)),
|
||||
"icp_stats": icp_stats,
|
||||
"aligned_path": aligned_path,
|
||||
"matrix_path": matrix_path,
|
||||
"error_stats": error_stats,
|
||||
"error_stats_path": error_stats_path,
|
||||
"elapsed_seconds": total_elapsed
|
||||
}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 命令行入口
|
||||
# ============================================================
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="高铁点云与 CAD 的仅平移配准：取车头车尾少量点估计整体平移"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--cad",
|
||||
"-c",
|
||||
default="datasets/CAD/m000004802418tc01_asm.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="不显示可视化窗口"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--no_cad_cache",
|
||||
action="store_true",
|
||||
help="Disable CAD sampled point cloud cache"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--save_debug_samples",
|
||||
action="store_true",
|
||||
help="Save front/tail debug sample point clouds"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--random_seed",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Random seed for repeatable sampling; set -1 to disable deterministic sampling"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--error_sample_points",
|
||||
type=int,
|
||||
default=100000,
|
||||
help="Sample count for alignment error stats; set 0 to use all aligned source points"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--no_error_stats",
|
||||
action="store_true",
|
||||
help="Skip point-to-CAD nearest-neighbor error statistics"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--save_aligned_float64",
|
||||
action="store_true",
|
||||
help="Use Open3D default float64-style aligned PLY output instead of compact float32 XYZ"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--no_icp",
|
||||
action="store_true",
|
||||
help="Skip rigid ICP refinement and use translation-only alignment"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--icp_sample_points",
|
||||
type=int,
|
||||
default=200000,
|
||||
help="Source sample count used for ICP; set 0 to use all source points"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--icp_voxel_size",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Voxel size for ICP downsampling; 0 disables voxel downsampling"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--icp_max_correspondence_distance",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="ICP max correspondence distance; 0 uses an automatic value from initial distances"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--icp_max_iteration",
|
||||
type=int,
|
||||
default=50,
|
||||
help="Maximum ICP iterations"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
random_seed = None if args.random_seed < 0 else args.random_seed
|
||||
|
||||
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,
|
||||
use_cad_cache=not args.no_cad_cache,
|
||||
save_debug_samples=args.save_debug_samples,
|
||||
random_seed=random_seed,
|
||||
error_sample_points=args.error_sample_points,
|
||||
compute_error_stats=not args.no_error_stats,
|
||||
save_aligned_float32=not args.save_aligned_float64,
|
||||
use_icp=not args.no_icp,
|
||||
icp_sample_points=args.icp_sample_points,
|
||||
icp_voxel_size=args.icp_voxel_size,
|
||||
icp_max_correspondence_distance=args.icp_max_correspondence_distance,
|
||||
icp_max_iteration=args.icp_max_iteration
|
||||
)
|
||||
|
||||
print("\n========== Result ==========")
|
||||
print("initial translation:", result["initial_translation"])
|
||||
print("final translation:", result["translation"])
|
||||
print("front-tail translation difference:", result["front_tail_translation_difference"])
|
||||
if result["icp_stats"] is not None:
|
||||
print("icp fitness:", result["icp_stats"]["fitness"])
|
||||
print("icp inlier rmse:", result["icp_stats"]["inlier_rmse"])
|
||||
print("aligned path:", result["aligned_path"])
|
||||
print("matrix path:", result["matrix_path"])
|
||||
if result["error_stats"] is not None:
|
||||
print("error stats path:", result["error_stats_path"])
|
||||
print("error rmse:", result["error_stats"]["rmse"])
|
||||
print("error p95:", result["error_stats"]["p95"])
|
||||
print("elapsed:", format_elapsed(result["elapsed_seconds"]))
|
||||
Reference in New Issue
Block a user