260 lines
8.5 KiB
Python
260 lines
8.5 KiB
Python
"""Rigid registration from scan point cloud to CAD reference.
|
|
|
|
Two-stage approach:
|
|
|
|
1. **Coarse**: centroid alignment + PCA axis matching with sign
|
|
disambiguation via per-axis skew (third central moment). Resolves
|
|
global orientation flips without requiring a good initial guess.
|
|
2. **Fine**: trimmed iterative closest point (ICP) on a downsampled CAD
|
|
target using ``scipy.spatial.cKDTree`` for nearest-neighbour
|
|
queries. Reflection-free rotation via SVD with the standard
|
|
determinant guard.
|
|
|
|
The result transforms scan points into the CAD coordinate frame:
|
|
``cad_point ≈ R @ scan_point + t``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
from scipy.spatial import cKDTree
|
|
|
|
from .preprocess import voxel_downsample
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class RegistrationResult:
|
|
R: np.ndarray
|
|
t: np.ndarray
|
|
rms_mm: float
|
|
iterations: int
|
|
inlier_ratio: float
|
|
converged: bool
|
|
coarse_R: np.ndarray
|
|
coarse_t: np.ndarray
|
|
initial_rms_mm: float
|
|
|
|
|
|
def register_to_cad(
|
|
scan: np.ndarray,
|
|
cad: np.ndarray,
|
|
config: dict,
|
|
) -> RegistrationResult:
|
|
"""Rigidly align ``scan`` to ``cad`` and return the transformation.
|
|
|
|
``config`` is a fully resolved config dict (must contain the
|
|
``registration`` section with ``icp_max_corr_mm`` already filled by
|
|
:func:`underframe_deflection.config.resolve_numeric_parameters`).
|
|
"""
|
|
if len(scan) < 3 or len(cad) < 3:
|
|
raise ValueError("register_to_cad requires at least 3 points in scan and cad.")
|
|
|
|
reg_cfg = config["registration"]
|
|
cad_cfg = config.get("cad", {})
|
|
|
|
# CAD-side downsampling — keeps NN queries cheap when the CAD has
|
|
# millions of points.
|
|
cad_voxel = float(cad_cfg.get("downsample_voxel_mm", 0.0) or 0.0)
|
|
target_n = int(reg_cfg.get("icp_downsample_target", 30000))
|
|
cad_sampled = cad
|
|
if cad_voxel > 0.0:
|
|
cad_sampled = voxel_downsample(cad_sampled, cad_voxel)
|
|
if target_n > 0 and len(cad_sampled) > target_n:
|
|
idx = np.linspace(0, len(cad_sampled) - 1, target_n, dtype=int)
|
|
cad_sampled = cad_sampled[idx]
|
|
|
|
R0, t0 = _coarse_align_pca(
|
|
scan,
|
|
cad_sampled,
|
|
disambiguate_sign=bool(reg_cfg.get("pca_sign_disambiguation", True)),
|
|
)
|
|
scan_coarse = scan @ R0.T + t0
|
|
|
|
max_corr = float(reg_cfg.get("icp_max_corr_mm", 25.0))
|
|
max_iter = int(reg_cfg.get("icp_max_iter", 60))
|
|
tol_mm = float(reg_cfg.get("icp_tol_mm", 0.05))
|
|
tol_rel = float(reg_cfg.get("icp_tol_rel", 1.0e-4))
|
|
trim = float(reg_cfg.get("trim_ratio", 0.2))
|
|
|
|
# ICP source downsampling — keeps query time bounded.
|
|
src_target = target_n if target_n > 0 else len(scan_coarse)
|
|
if len(scan_coarse) > src_target:
|
|
idx = np.linspace(0, len(scan_coarse) - 1, src_target, dtype=int)
|
|
scan_for_icp = scan_coarse[idx]
|
|
else:
|
|
scan_for_icp = scan_coarse
|
|
|
|
tree = cKDTree(cad_sampled)
|
|
R = np.eye(3, dtype=np.float64)
|
|
t = np.zeros(3, dtype=np.float64)
|
|
prev_rms = float("inf")
|
|
current = scan_for_icp.copy()
|
|
iterations_run = 0
|
|
inlier_ratio = 0.0
|
|
rms = float("inf")
|
|
converged = False
|
|
|
|
initial_rms = _correspondence_rms(current, tree, max_corr)
|
|
|
|
for iteration in range(1, max_iter + 1):
|
|
iterations_run = iteration
|
|
distances, neighbours = tree.query(current, k=1)
|
|
keep = distances <= max_corr
|
|
if trim > 0.0 and np.any(keep):
|
|
kept_d = distances[keep]
|
|
if len(kept_d) > 8:
|
|
cutoff = float(np.quantile(kept_d, 1.0 - trim))
|
|
keep = keep & (distances <= cutoff)
|
|
if np.count_nonzero(keep) < 3:
|
|
break
|
|
src = current[keep]
|
|
dst = cad_sampled[neighbours[keep]]
|
|
delta_R, delta_t = _kabsch(src, dst)
|
|
current = current @ delta_R.T + delta_t
|
|
R = delta_R @ R
|
|
t = delta_R @ t + delta_t
|
|
|
|
residuals = current - cad_sampled[neighbours]
|
|
residual_norms = np.linalg.norm(residuals[keep], axis=1)
|
|
rms = float(np.sqrt(np.mean(residual_norms ** 2))) if len(residual_norms) else float("inf")
|
|
inlier_ratio = float(np.count_nonzero(keep) / max(len(current), 1))
|
|
|
|
if abs(prev_rms - rms) <= tol_mm or (
|
|
prev_rms > 0.0 and abs(prev_rms - rms) / max(prev_rms, 1e-9) <= tol_rel
|
|
):
|
|
converged = True
|
|
break
|
|
prev_rms = rms
|
|
|
|
R_total = R @ R0
|
|
t_total = R @ t0 + t
|
|
return RegistrationResult(
|
|
R=R_total,
|
|
t=t_total,
|
|
rms_mm=rms,
|
|
iterations=iterations_run,
|
|
inlier_ratio=inlier_ratio,
|
|
converged=converged,
|
|
coarse_R=R0,
|
|
coarse_t=t0,
|
|
initial_rms_mm=initial_rms,
|
|
)
|
|
|
|
|
|
def apply_transform(points: np.ndarray, R: np.ndarray, t: np.ndarray) -> np.ndarray:
|
|
"""Apply ``R @ p + t`` to a point cloud (broadcasted)."""
|
|
return points @ R.T + t
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Internals
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _coarse_align_pca(
|
|
scan: np.ndarray,
|
|
cad: np.ndarray,
|
|
*,
|
|
disambiguate_sign: bool,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Return (R, t) that map ``scan`` into the CAD PCA frame."""
|
|
scan_axes, scan_center = _pca_basis(scan)
|
|
cad_axes, cad_center = _pca_basis(cad)
|
|
|
|
if disambiguate_sign:
|
|
scan_axes = _resolve_sign_by_skew(scan, scan_center, scan_axes, cad, cad_center, cad_axes)
|
|
|
|
# R maps scan → CAD: scan_axes^T rotates scan into its own PCA basis,
|
|
# cad_axes rotates back into the CAD frame.
|
|
R = cad_axes @ scan_axes.T
|
|
if np.linalg.det(R) < 0:
|
|
scan_axes[:, -1] *= -1.0
|
|
R = cad_axes @ scan_axes.T
|
|
t = cad_center - R @ scan_center
|
|
return R, t
|
|
|
|
|
|
def _pca_basis(points: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
|
center = points.mean(axis=0)
|
|
centered = points - center
|
|
covariance = centered.T @ centered / max(len(points), 1)
|
|
eigenvalues, eigenvectors = np.linalg.eigh(covariance)
|
|
order = np.argsort(eigenvalues)[::-1]
|
|
axes = eigenvectors[:, order]
|
|
# ``axes`` is (3, 3) where columns are unit eigenvectors sorted by
|
|
# decreasing variance. Treat columns as the new basis directions.
|
|
return axes, center
|
|
|
|
|
|
def _resolve_sign_by_skew(
|
|
scan: np.ndarray,
|
|
scan_center: np.ndarray,
|
|
scan_axes: np.ndarray,
|
|
cad: np.ndarray,
|
|
cad_center: np.ndarray,
|
|
cad_axes: np.ndarray,
|
|
) -> np.ndarray:
|
|
"""Pick eigenvector signs that match the chirality of ``cad``."""
|
|
scan_proj = (scan - scan_center) @ scan_axes
|
|
cad_proj = (cad - cad_center) @ cad_axes
|
|
cad_skew = np.array([_third_moment(cad_proj[:, i]) for i in range(3)])
|
|
best_signs = np.array([1.0, 1.0, 1.0])
|
|
best_score = np.inf
|
|
for s0 in (-1.0, 1.0):
|
|
for s1 in (-1.0, 1.0):
|
|
for s2 in (-1.0, 1.0):
|
|
signs = np.array([s0, s1, s2])
|
|
if np.prod(signs) < 0:
|
|
# Disallow reflections; the kabsch det guard will
|
|
# cover the residual cases.
|
|
continue
|
|
scan_skew = np.array(
|
|
[signs[i] ** 3 * _third_moment(scan_proj[:, i]) for i in range(3)]
|
|
)
|
|
score = float(np.sum(np.abs(scan_skew - cad_skew)))
|
|
if score < best_score:
|
|
best_score = score
|
|
best_signs = signs
|
|
return scan_axes * best_signs[None, :]
|
|
|
|
|
|
def _third_moment(values: np.ndarray) -> float:
|
|
if len(values) == 0:
|
|
return 0.0
|
|
std = float(np.std(values))
|
|
if std < 1e-9:
|
|
return 0.0
|
|
centered = values - values.mean()
|
|
return float(np.mean(centered ** 3) / (std ** 3))
|
|
|
|
|
|
def _kabsch(src: np.ndarray, dst: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Closed-form rigid-body fit (no reflections)."""
|
|
src_center = src.mean(axis=0)
|
|
dst_center = dst.mean(axis=0)
|
|
src_centered = src - src_center
|
|
dst_centered = dst - dst_center
|
|
H = src_centered.T @ dst_centered
|
|
U, _, Vt = np.linalg.svd(H)
|
|
D = np.eye(3, dtype=np.float64)
|
|
if np.linalg.det(Vt.T @ U.T) < 0:
|
|
D[-1, -1] = -1.0
|
|
R = Vt.T @ D @ U.T
|
|
t = dst_center - R @ src_center
|
|
return R, t
|
|
|
|
|
|
def _correspondence_rms(points: np.ndarray, tree: cKDTree, max_corr: float) -> float:
|
|
if len(points) == 0:
|
|
return float("inf")
|
|
distances, _ = tree.query(points, k=1)
|
|
keep = distances <= max_corr
|
|
if not np.any(keep):
|
|
return float(np.sqrt(np.mean(distances ** 2)))
|
|
return float(np.sqrt(np.mean(distances[keep] ** 2)))
|
|
|
|
|
|
__all__ = ["RegistrationResult", "register_to_cad", "apply_transform"]
|