Files
calibration/tools/visualize_pair_3d.py

621 lines
21 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""Interactive 3D inspection of LiDARIMU motion-pair registration.
Similar to the RTKLiDAR ``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
(rotation_only runs default to R conjug + t_B so bad IMU Δp
does not dominate the overlay)
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
import sys
from pathlib import Path
from typing import Any
import numpy as np
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
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_frame_at, load_lidar_frames
from imu_lidar.motion_pairs import build_motion_pairs
from imu_lidar.motion_pairs_io import (
load_motion_pairs,
pairs_for_session,
resolve_motion_pairs_path,
)
from imu_lidar.registration import register_lidar_pair
from imu_lidar.time_offset import lidar_time_to_imu_time
class _LazyFrameStore:
"""Load NPZ frames on demand; indices match mid-time-sorted ``load_lidar_frames``."""
def __init__(self, lidar_dir: Path, *, max_cached: int = 16):
self.lidar_dir = Path(lidar_dir)
self.max_cached = max_cached
self._cache: dict[int, object] = {}
self._order: list[int] = []
def __getitem__(self, index: int):
index = int(index)
if index in self._cache:
return self._cache[index]
frame = load_lidar_frame_at(self.lidar_dir, index)
self._cache[index] = frame
self._order.append(index)
while len(self._order) > self.max_cached:
old = self._order.pop(0)
self._cache.pop(old, None)
return frame
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, dict[str, Any]]:
summary = json.loads(summary_path.read_text(encoding="utf-8"))
t_block = summary.get("T_IMU_lidar")
meta: dict[str, Any] = {
"rotation_only": False,
"translation_accepted": False,
"status": str(summary.get("status") or ""),
}
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), meta
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]
joint = session.get("joint") or {}
bias = np.asarray(
(session.get("imu_audit") or {}).get("gyro_bias_rad_s")
or joint.get("gyro_bias_rad_s")
or [0.0, 0.0, 0.0],
dtype=float,
).reshape(3)
status = str(summary.get("status") or "")
translation_accepted = bool(
joint.get("translation_accepted")
or (summary.get("details") or {}).get("joint", {}).get("translation_accepted")
)
rotation_only = ("rotation_only" in status) or (
not translation_accepted and float(np.linalg.norm(t_mat[:3, 3])) < 1e-9
)
meta.update(
{
"rotation_only": rotation_only,
"translation_accepted": translation_accepted,
"status": status,
}
)
return t_mat, dt, bias, meta
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,
*,
mode4_translation: str = "imu",
) -> dict[str, np.ndarray]:
"""Build overlay transforms.
``mode4_translation``:
- ``imu``: full SE3 conjug ``X^{-1} A X`` (needs trustworthy IMU Δp)
- ``gicp``: rotation conjug only; translation taken from LiDAR B
(correct check for rotation_only calibrations)
"""
calibrated = inverse_transform(x) @ a_ij @ x
if mode4_translation == "gicp":
calibrated = make_transform(b_gicp[:3, 3], calibrated[:3, :3])
elif mode4_translation != "imu":
raise ValueError(f"unknown mode4_translation={mode4_translation!r}")
return {
MODE_NAMES[0]: np.eye(4),
MODE_NAMES[1]: a_ij.copy(),
MODE_NAMES[2]: b_gicp.copy(),
MODE_NAMES[3]: calibrated,
}
def _resolve_pair(
frames,
pairs,
pair_index: int,
x: np.ndarray,
*,
mode4_translation: str = "imu",
):
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, mode4_translation=mode4_translation
)
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("LiDARIMU 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,
mode4_translation: str = "imu",
) -> 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, mode4_translation=mode4_translation
)
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, mode4_translation=mode4_translation
)
viewer = o3d.visualization.VisualizerWithKeyCallback()
viewer.create_window("LiDARIMU 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, mode4_translation=mode4_translation
)
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",
type=Path,
default=None,
help="IMU CSV (only needed when rebuilding pairs without motion_pairs.json)",
)
parser.add_argument(
"--summary",
required=True,
type=Path,
help="summary.json (or T_IMU_lidar.json) from a calibration run",
)
parser.add_argument(
"--motion-pairs",
type=Path,
default=None,
help="Cached motion_pairs.json (default: next to --summary)",
)
parser.add_argument(
"--session-id",
default=None,
help="Session id inside multi-session motion_pairs.json",
)
parser.add_argument(
"--rebuild-pairs",
action="store_true",
help="Ignore cache and rebuild pairs from IMU/LiDAR (slow)",
)
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)",
)
parser.add_argument(
"--mode4-translation",
choices=("auto", "gicp", "imu"),
default="auto",
help=(
"Mode-4 translation source: gicp=R conjug + t_B (rotation check); "
"imu=full X^-1 A X; auto=gicp for rotation_only summaries"
),
)
args = parser.parse_args(argv)
x, delta_t_s, gyro_bias, extr_meta = _load_extrinsic(args.summary)
if args.mode4_translation == "auto":
mode4_translation = "gicp" if extr_meta.get("rotation_only") else "imu"
else:
mode4_translation = args.mode4_translation
cache_path = args.motion_pairs or resolve_motion_pairs_path(args.summary)
use_cache = (not args.rebuild_pairs) and cache_path is not None and args.frame_i is None
frames = None
pairs: tuple = ()
fixed_single_pair = None
label = ""
b_gicp = np.eye(4)
transforms: dict[str, np.ndarray] = {}
frame_i = frame_j = None
if use_cache:
payload = load_motion_pairs(cache_path)
pair_list = pairs_for_session(payload, args.session_id)
if not pair_list:
raise SystemExit(f"no pairs in cache: {cache_path}")
if not 0 <= args.pair_index < len(pair_list):
raise SystemExit(
f"pair-index {args.pair_index} outside [0, {len(pair_list) - 1}] "
f"({len(pair_list)} pairs in cache)"
)
frames = _LazyFrameStore(args.lidar)
pairs = tuple(pair_list)
frame_i, frame_j, a_ij, b_gicp, transforms, label = _resolve_pair(
frames, pairs, args.pair_index, x, mode4_translation=mode4_translation
)
print(f"loaded {len(pairs)} cached pairs from {cache_path}")
else:
if args.imu is None:
raise SystemExit("--imu is required when motion_pairs.json is missing (or use --rebuild-pairs with --imu)")
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,
)
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, mode4_translation=mode4_translation
)
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, mode4_translation=mode4_translation
)
print(f"rebuilt {len(pairs)} pairs from {len(keyframes.indices)} keyframes")
print(
f"mode4 translation={mode4_translation} "
f"(status={extr_meta.get('status') or 'n/a'}, "
f"rotation_only={bool(extr_meta.get('rotation_only'))})"
)
if mode4_translation == "gicp":
print(
"note: mode4 uses R conjug + t_B; IMU Δp is ignored "
"(typical for rotation_only — raw Δp often has large Z drift)."
)
if mode4_translation == "imu":
print(
"note: mode4 uses full X^-1 A X. If clouds stack vertically, "
"IMU Δp is likely bad; retry with --mode4-translation gicp."
)
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,
mode4_translation=mode4_translation,
)
return 0
if __name__ == "__main__":
raise SystemExit(main())