#!/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())