calibration/imu_lidar/registration.py

"""LiDAR relative-motion registration.

Uses Open3D Generalized ICP when available; otherwise a NumPy point-to-point ICP.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from .contracts import LidarFrame
from .geometry import make_transform, orthonormalize_rotation, rotation_angle_deg, so3_log


@dataclass(frozen=True)
class RegistrationResult:
    transform: np.ndarray
    fitness: float
    rotation_deg: float
    translation_m: float
    backend: str
    ok: bool


def _voxel_downsample(points: np.ndarray, voxel: float) -> np.ndarray:
    if points.shape[0] == 0:
        return points
    quantized = np.floor(points / voxel).astype(np.int64)
    _, unique_indices = np.unique(quantized, axis=0, return_index=True)
    return points[np.sort(unique_indices)]


def _numpy_icp(
    source: np.ndarray,
    target: np.ndarray,
    *,
    max_iterations: int = 30,
    max_correspondence: float = 1.0,
) -> RegistrationResult:
    src = _voxel_downsample(source, 0.2)
    tgt = _voxel_downsample(target, 0.2)
    if src.shape[0] < 50 or tgt.shape[0] < 50:
        return RegistrationResult(np.eye(4), 0.0, 0.0, 0.0, "numpy_icp", False)

    # Subsample for speed.
    rng = np.random.default_rng(0)
    if src.shape[0] > 4000:
        src = src[rng.choice(src.shape[0], 4000, replace=False)]
    if tgt.shape[0] > 8000:
        tgt = tgt[rng.choice(tgt.shape[0], 8000, replace=False)]

    r = np.eye(3)
    t = np.zeros(3)
    last_error = 1e9
    inlier_ratio = 0.0
    for _ in range(max_iterations):
        transformed = src @ r.T + t
        # Nearest neighbour in target via brute force on chunks.
        diff = transformed[:, None, :] - tgt[None, :, :]
        dist2 = np.sum(diff * diff, axis=2)
        nn = np.argmin(dist2, axis=1)
        dist = np.sqrt(dist2[np.arange(src.shape[0]), nn])
        mask = dist < max_correspondence
        inlier_ratio = float(np.mean(mask))
        if np.count_nonzero(mask) < 30:
            break
        p = transformed[mask]
        q = tgt[nn[mask]]
        mu_p = p.mean(axis=0)
        mu_q = q.mean(axis=0)
        h = (p - mu_p).T @ (q - mu_q)
        u, _, vt = np.linalg.svd(h)
        r_delta = vt.T @ u.T
        if np.linalg.det(r_delta) < 0:
            vt[-1, :] *= -1
            r_delta = vt.T @ u.T
        t_delta = mu_q - r_delta @ mu_p
        # Update global transform: x' = r_delta (r x + t) + t_delta
        r = orthonormalize_rotation(r_delta @ r)
        t = r_delta @ t + t_delta
        mean_err = float(np.mean(dist[mask]))
        if abs(last_error - mean_err) < 1e-4:
            break
        last_error = mean_err

    transform = make_transform(t, r)
    return RegistrationResult(
        transform=transform,
        fitness=inlier_ratio,
        rotation_deg=rotation_angle_deg(r),
        translation_m=float(np.linalg.norm(t)),
        backend="numpy_icp",
        ok=inlier_ratio > 0.15,
    )


def _open3d_gicp(source: np.ndarray, target: np.ndarray) -> RegistrationResult | None:
    try:
        import open3d as o3d
    except ImportError:
        return None

    src = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(source))
    tgt = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(target))
    src = src.voxel_down_sample(0.2)
    tgt = tgt.voxel_down_sample(0.2)
    if len(src.points) < 50 or len(tgt.points) < 50:
        return RegistrationResult(np.eye(4), 0.0, 0.0, 0.0, "open3d_gicp", False)
    src.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=1.0, max_nn=30))
    tgt.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=1.0, max_nn=30))
    result = o3d.pipelines.registration.registration_generalized_icp(
        src,
        tgt,
        1.0,
        np.eye(4),
        o3d.pipelines.registration.TransformationEstimationForGeneralizedICP(),
        o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=50),
    )
    transform = np.asarray(result.transformation, dtype=float)
    return RegistrationResult(
        transform=transform,
        fitness=float(result.fitness),
        rotation_deg=rotation_angle_deg(transform[:3, :3]),
        translation_m=float(np.linalg.norm(transform[:3, 3])),
        backend="open3d_gicp",
        ok=float(result.fitness) > 0.15,
    )


def register_lidar_pair(source_points: np.ndarray, target_points: np.ndarray) -> RegistrationResult:
    """Register source -> target and return ``T_target_source``."""

    source = np.asarray(source_points, dtype=float).reshape(-1, 3)
    target = np.asarray(target_points, dtype=float).reshape(-1, 3)
    open3d_result = _open3d_gicp(source, target)
    if open3d_result is not None:
        return open3d_result
    return _numpy_icp(source, target)


def estimate_frame_rotations(
    frames: list[LidarFrame],
    *,
    stride: int = 1,
) -> tuple[list[np.ndarray], list[tuple[float, float]]]:
    """Estimate consecutive (or strided) LiDAR relative rotations for time sync."""

    rotations: list[np.ndarray] = []
    pair_times: list[tuple[float, float]] = []
    for index in range(0, len(frames) - stride, max(stride, 1)):
        a = frames[index]
        b = frames[index + stride]
        result = register_lidar_pair(b.points_xyz, a.points_xyz)
        if not result.ok:
            continue
        rotations.append(result.transform[:3, :3])
        pair_times.append((a.t_mid_s, b.t_mid_s))
    return rotations, pair_times