453 lines
15 KiB
Python
453 lines
15 KiB
Python
#!/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())
|