calibration/tools/visualize_pair_3d.py

#!/usr/bin/env python3
"""Interactive 3D inspection of LiDAR–IMU motion-pair registration.

Similar to the RTK–LiDAR ``visualize_pair_3d`` viewer, but A comes from IMU
preintegration and X is ``T_IMU_lidar``.

Modes (keyboard):
  1          raw source (no transform)
  2          IMU prediction with X=I  (B_pred = A)
  3          LiDAR registration B     (reference)
  4          calibrated prediction    B_pred = X^{-1} A X
  N / ]      next motion pair
  P / [      previous motion pair
  Q / Esc    exit

Blue = target keyframe i; orange = source keyframe j after the selected transform.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from imu_lidar.geometry import (
    inverse_transform,
    make_transform,
    rotation_angle_deg,
    rpy_deg_xyz,
    transform_points,
)
from imu_lidar.imu_io import load_imu_samples
from imu_lidar.imu_preintegration import preintegrate_imu
from imu_lidar.keyframes import build_keyframes
from imu_lidar.lidar_io import load_lidar_frames
from imu_lidar.motion_pairs import build_motion_pairs
from imu_lidar.registration import register_lidar_pair
from imu_lidar.time_offset import lidar_time_to_imu_time


COLORS = {
    "target": [0.10, 0.65, 1.00],
    "source": [1.00, 0.35, 0.05],
}

MODE_NAMES = (
    "1 raw",
    "2 IMU (X=I)",
    "3 LiDAR B",
    "4 calibrated X^-1 A X",
)


def _load_extrinsic(summary_path: Path) -> tuple[np.ndarray, float, np.ndarray]:
    summary = json.loads(summary_path.read_text(encoding="utf-8"))
    t_block = summary.get("T_IMU_lidar")
    if t_block is None:
        matrix = summary.get("matrix")
        if matrix is not None:
            return np.asarray(matrix, dtype=float), 0.0, np.zeros(3)
        raise ValueError(f"no T_IMU_lidar in {summary_path}")
    t_mat = np.asarray(t_block["matrix"], dtype=float)
    dt = float(summary.get("time_offset_s") or 0.0)
    session = (summary.get("details") or {}).get("sessions", [{}])[0]
    bias = np.asarray(
        (session.get("imu_audit") or {}).get("gyro_bias_rad_s")
        or (session.get("joint") or {}).get("gyro_bias_rad_s")
        or [0.0, 0.0, 0.0],
        dtype=float,
    ).reshape(3)
    return t_mat, dt, bias


def _delta_components(reference: np.ndarray, candidate: np.ndarray) -> dict:
    delta = inverse_transform(reference) @ candidate
    translation = np.asarray(delta[:3, 3], dtype=float)
    return {
        "translation_xyz_cm": (translation * 100.0).tolist(),
        "translation_norm_cm": float(np.linalg.norm(translation) * 100.0),
        "rotation_rpy_deg_xyz": rpy_deg_xyz(delta[:3, :3]).tolist(),
        "rotation_angle_deg": rotation_angle_deg(delta[:3, :3]),
    }


def _print_delta(name: str, reference: np.ndarray, candidate: np.ndarray) -> dict:
    item = _delta_components(reference, candidate)
    tx, ty, tz = item["translation_xyz_cm"]
    roll, pitch, yaw = item["rotation_rpy_deg_xyz"]
    print(
        f"{name}: B^-1*motion  "
        f"t_xyz=[{tx:+.3f}, {ty:+.3f}, {tz:+.3f}] cm  "
        f"rpy=[{roll:+.3f}, {pitch:+.3f}, {yaw:+.3f}] deg  "
        f"|t|={item['translation_norm_cm']:.3f} cm  "
        f"|R|={item['rotation_angle_deg']:.4f} deg"
    )
    return item


def _cloud(o3d, points: np.ndarray, color, voxel: float):
    item = o3d.geometry.PointCloud()
    item.points = o3d.utility.Vector3dVector(points)
    if voxel > 0:
        item = item.voxel_down_sample(voxel)
    item.paint_uniform_color(color)
    return item


def _set_cloud_points(cloud, points: np.ndarray, color, voxel: float, o3d) -> None:
    tmp = _cloud(o3d, points, color, voxel)
    cloud.points = tmp.points
    cloud.colors = tmp.colors


def _build_pair_list(
    *,
    lidar_dir: Path,
    imu_path: Path,
    delta_t_s: float,
    gyro_bias: np.ndarray,
    min_rotation_deg: float,
    min_translation_m: float,
):
    frames = load_lidar_frames(lidar_dir)
    imu = load_imu_samples(imu_path)
    keyframes = build_keyframes(
        frames,
        min_translation_m=min_translation_m,
        min_rotation_deg=min_rotation_deg,
    )
    pair_set = build_motion_pairs(
        session_id="viz",
        keyframes=list(keyframes.frames),
        keyframe_indices=keyframes.indices,
        imu=imu,
        delta_t_s=delta_t_s,
        gyro_bias_rad_s=gyro_bias,
        min_rotation_deg=min_rotation_deg,
        min_translation_m=min_translation_m,
    )
    return frames, imu, keyframes, pair_set


def _pair_from_indices(
    frames,
    imu,
    *,
    i: int,
    j: int,
    delta_t_s: float,
    gyro_bias: np.ndarray,
):
    frame_i = frames[i]
    frame_j = frames[j]
    reg = register_lidar_pair(frame_j.points_xyz, frame_i.points_xyz)
    t_i = lidar_time_to_imu_time(frame_i.t_mid_s, delta_t_s)
    t_j = lidar_time_to_imu_time(frame_j.t_mid_s, delta_t_s)
    preint = preintegrate_imu(
        imu.t_s,
        imu.gyro_rad_s,
        imu.acc_m_s2,
        t_i,
        t_j,
        gyro_bias,
        np.zeros(3),
    )
    a = make_transform(preint.delta_p, preint.delta_R)
    return frame_i, frame_j, a, reg.transform


def _transforms_for_pair(x: np.ndarray, a_ij: np.ndarray, b_gicp: np.ndarray) -> dict[str, np.ndarray]:
    return {
        MODE_NAMES[0]: np.eye(4),
        MODE_NAMES[1]: a_ij.copy(),
        MODE_NAMES[2]: b_gicp.copy(),
        MODE_NAMES[3]: inverse_transform(x) @ a_ij @ x,
    }


def _resolve_pair(frames, pairs, pair_index: int, x: np.ndarray):
    pair = pairs[pair_index]
    frame_i = frames[pair.i]
    frame_j = frames[pair.j]
    a_ij = make_transform(
        pair.t_A_m if pair.t_A_m is not None else np.zeros(3),
        pair.R_A,
    )
    b_gicp = make_transform(
        pair.t_B_m if pair.t_B_m is not None else np.zeros(3),
        pair.R_B,
    )
    transforms = _transforms_for_pair(x, a_ij, b_gicp)
    label = (
        f"pair {pair_index + 1}/{len(pairs)}  "
        f"frames {pair.i} <- {pair.j}  "
        f"rotB={rotation_angle_deg(pair.R_B):.2f} deg  "
        f"|tB|={0.0 if pair.t_B_m is None else float(np.linalg.norm(pair.t_B_m)):.3f} m"
    )
    return frame_i, frame_j, a_ij, b_gicp, transforms, label


def _print_pair_header(label: str, b_gicp: np.ndarray, transforms: dict[str, np.ndarray]) -> None:
    print("-" * 72)
    print(label)
    print("blue=target i | orange=source j")
    print("1-4: overlay mode | N/]: next pair | P/[: prev pair | Q/Esc: exit")
    baseline = _print_delta("mode4 minus mode3", b_gicp, transforms[MODE_NAMES[3]])
    roll, pitch, yaw = np.abs(baseline["rotation_rpy_deg_xyz"])
    if max(roll, pitch) > max(0.10, 2.0 * yaw):
        print("note: roll/pitch dominate yaw on this pair.")
    tx, ty, tz = np.abs(baseline["translation_xyz_cm"])
    if tz > max(tx, ty):
        print("note: largest translation component is Z for this pair.")


def _save_topdown_png(
    path: Path,
    target: np.ndarray,
    source: np.ndarray,
    transforms: dict[str, np.ndarray],
) -> None:
    import matplotlib.pyplot as plt

    def downsample(points: np.ndarray) -> np.ndarray:
        if points.shape[0] <= 8000:
            return points
        idx = np.linspace(0, points.shape[0] - 1, 8000).astype(int)
        return points[idx]

    names = list(transforms.keys())
    fig, axes = plt.subplots(2, 2, figsize=(10, 10), sharex=True, sharey=True)
    tgt = downsample(target)
    for ax, name in zip(axes.ravel(), names):
        src = downsample(transform_points(source, transforms[name]))
        ax.scatter(tgt[:, 0], tgt[:, 1], s=1, c="tab:blue", alpha=0.35, label="target i")
        ax.scatter(src[:, 0], src[:, 1], s=1, c="tab:orange", alpha=0.35, label="source j")
        ax.set_title(name)
        ax.set_aspect("equal", adjustable="box")
        ax.grid(alpha=0.3)
    axes[0, 0].legend(loc="upper right", markerscale=4)
    fig.suptitle("LiDAR–IMU pair overlay (XY top-down)")
    fig.tight_layout()
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=140)
    plt.close(fig)
    print(f"wrote {path}")


def _run_gui(
    *,
    frames,
    pairs,
    x: np.ndarray,
    start_index: int,
    voxel: float,
    fixed_single_pair: tuple | None,
) -> None:
    import open3d as o3d

    if fixed_single_pair is not None:
        frame_i, frame_j, a_ij, b_gicp = fixed_single_pair
        transforms = _transforms_for_pair(x, a_ij, b_gicp)
        label = f"fixed frames (no pair switching)"
        pair_index = 0
        n_pairs = 1
        use_list = False
    else:
        pair_index = int(np.clip(start_index, 0, len(pairs) - 1))
        n_pairs = len(pairs)
        use_list = True
        frame_i, frame_j, a_ij, b_gicp, transforms, label = _resolve_pair(
            frames, pairs, pair_index, x
        )

    viewer = o3d.visualization.VisualizerWithKeyCallback()
    viewer.create_window("LiDAR–IMU registration inspection", 1400, 900)
    target_cloud = _cloud(o3d, frame_i.points_xyz, COLORS["target"], voxel)
    source_cloud = _cloud(o3d, frame_j.points_xyz, COLORS["source"], voxel)
    viewer.add_geometry(target_cloud)
    viewer.add_geometry(source_cloud)
    viewer.add_geometry(o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0))
    viewer.get_render_option().background_color = np.array([0.02, 0.02, 0.02])
    viewer.get_render_option().point_size = 2.0

    state = {
        "pair_index": pair_index,
        "mode_name": MODE_NAMES[3],
        "current": np.eye(4),
        "transforms": transforms,
        "b_gicp": b_gicp,
        "frame_i": frame_i,
        "frame_j": frame_j,
    }

    def apply_mode(vis, mode_name: str, *, announce: bool = True) -> None:
        desired = state["transforms"][mode_name]
        source_cloud.transform(desired @ inverse_transform(state["current"]))
        state["current"] = desired
        state["mode_name"] = mode_name
        vis.update_geometry(source_cloud)
        if announce:
            if mode_name == MODE_NAMES[2]:
                print(f"{mode_name}: registration reference; delta = 0")
            else:
                _print_delta(mode_name + " minus mode3", state["b_gicp"], desired)

    def load_pair(vis, new_index: int) -> None:
        if not use_list:
            print("pair switching disabled in --frame-i/--frame-j mode")
            return
        new_index = int(new_index) % n_pairs
        frame_i, frame_j, _a, b_gicp, transforms, label = _resolve_pair(
            frames, pairs, new_index, x
        )
        state["pair_index"] = new_index
        state["transforms"] = transforms
        state["b_gicp"] = b_gicp
        state["frame_i"] = frame_i
        state["frame_j"] = frame_j
        state["current"] = np.eye(4)
        _set_cloud_points(target_cloud, frame_i.points_xyz, COLORS["target"], voxel, o3d)
        _set_cloud_points(source_cloud, frame_j.points_xyz, COLORS["source"], voxel, o3d)
        vis.update_geometry(target_cloud)
        vis.update_geometry(source_cloud)
        _print_pair_header(label, b_gicp, transforms)
        apply_mode(vis, state["mode_name"], announce=True)

    def make_mode_cb(mode_name: str):
        def callback(vis):
            apply_mode(vis, mode_name, announce=True)
            return False

        return callback

    def next_pair(vis):
        load_pair(vis, state["pair_index"] + 1)
        return False

    def prev_pair(vis):
        load_pair(vis, state["pair_index"] - 1)
        return False

    _print_pair_header(label, b_gicp, transforms)
    for key, name in zip((ord("1"), ord("2"), ord("3"), ord("4")), MODE_NAMES):
        viewer.register_key_callback(key, make_mode_cb(name))
    for key in (ord("N"), ord("n"), ord("]")):
        viewer.register_key_callback(key, next_pair)
    for key in (ord("P"), ord("p"), ord("[")):
        viewer.register_key_callback(key, prev_pair)

    apply_mode(viewer, MODE_NAMES[3], announce=False)
    viewer.run()
    viewer.destroy_window()


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--lidar", required=True, type=Path, help="LiDAR session directory")
    parser.add_argument("--imu", required=True, type=Path, help="IMU CSV")
    parser.add_argument(
        "--summary",
        required=True,
        type=Path,
        help="summary.json (or T_IMU_lidar.json) from a calibration run",
    )
    parser.add_argument("--pair-index", type=int, default=0, help="Starting motion-pair index")
    parser.add_argument("--frame-i", type=int, default=None, help="Optional explicit frame index i")
    parser.add_argument("--frame-j", type=int, default=None, help="Optional explicit frame index j")
    parser.add_argument("--voxel", type=float, default=0.12)
    parser.add_argument("--min-pair-rotation-deg", type=float, default=3.0)
    parser.add_argument("--min-pair-translation-m", type=float, default=0.3)
    parser.add_argument(
        "--save-png",
        type=Path,
        default=None,
        help="Write a 2x2 XY top-down comparison PNG for the starting pair",
    )
    parser.add_argument(
        "--no-gui",
        action="store_true",
        help="Skip Open3D window (use with --save-png)",
    )
    args = parser.parse_args(argv)

    x, delta_t_s, gyro_bias = _load_extrinsic(args.summary)
    frames, imu, keyframes, pair_set = _build_pair_list(
        lidar_dir=args.lidar,
        imu_path=args.imu,
        delta_t_s=delta_t_s,
        gyro_bias=gyro_bias,
        min_rotation_deg=args.min_pair_rotation_deg,
        min_translation_m=args.min_pair_translation_m,
    )

    fixed_single_pair = None
    if args.frame_i is not None and args.frame_j is not None:
        frame_i, frame_j, a_ij, b_gicp = _pair_from_indices(
            frames,
            imu,
            i=args.frame_i,
            j=args.frame_j,
            delta_t_s=delta_t_s,
            gyro_bias=gyro_bias,
        )
        transforms = _transforms_for_pair(x, a_ij, b_gicp)
        label = f"frames {args.frame_i} <- {args.frame_j}"
        fixed_single_pair = (frame_i, frame_j, a_ij, b_gicp)
        pairs = ()
    else:
        if not pair_set.pairs:
            raise SystemExit("no motion pairs rebuilt; loosen min-pair thresholds or check data")
        if not 0 <= args.pair_index < len(pair_set.pairs):
            raise SystemExit(
                f"pair-index {args.pair_index} outside [0, {len(pair_set.pairs) - 1}] "
                f"({len(pair_set.pairs)} pairs available)"
            )
        pairs = pair_set.pairs
        frame_i, frame_j, a_ij, b_gicp, transforms, label = _resolve_pair(
            frames, pairs, args.pair_index, x
        )
        print(f"rebuilt {len(pairs)} pairs from {len(keyframes.indices)} keyframes")

    if args.save_png is not None:
        _print_pair_header(label, b_gicp, transforms)
        _save_topdown_png(args.save_png, frame_i.points_xyz, frame_j.points_xyz, transforms)

    if args.no_gui:
        return 0

    try:
        import open3d  # noqa: F401
    except ImportError as exc:
        raise SystemExit(
            "Open3D is required for interactive view. "
            "Install with: python -m pip install -e \".[open3d]\" "
            "or use --no-gui --save-png out.png"
        ) from exc

    _run_gui(
        frames=frames,
        pairs=pairs,
        x=x,
        start_index=args.pair_index,
        voxel=args.voxel,
        fixed_single_pair=fixed_single_pair,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())