986 lines
30 KiB
Python
986 lines
30 KiB
Python
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"
|
||
# 优先存在 CAD 文件同目录下的 .cad_sample_cache/，跨任务复用；
|
||
# 目录不可写时回退到 output_dir/cad_cache/
|
||
cad_dir = os.path.dirname(os.path.abspath(cad_path))
|
||
shared_cache_dir = os.path.join(cad_dir, ".cad_sample_cache")
|
||
try:
|
||
os.makedirs(shared_cache_dir, exist_ok=True)
|
||
# 写权限探测
|
||
test_path = os.path.join(shared_cache_dir, ".write_test")
|
||
open(test_path, "w").close()
|
||
os.remove(test_path)
|
||
cache_dir = shared_cache_dir
|
||
except OSError:
|
||
cache_dir = os.path.join(output_dir, "cad_cache")
|
||
return os.path.join(cache_dir, 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"]))
|