155 lines
5.0 KiB
Python
155 lines
5.0 KiB
Python
"""Serialize / deserialize motion pairs for fast visualization."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from .contracts import MotionPair
|
|
|
|
SCHEMA_VERSION = 2
|
|
|
|
# Keep visualization fields plus the compact 3x3 rotation metadata needed to
|
|
# rerun Phase-A without repeating LiDAR registration. Full 9x9 Phase-C matrices
|
|
# remain excluded from this cache.
|
|
_METADATA_KEEP = frozenset(
|
|
{
|
|
"backend",
|
|
"rotation_deg_A",
|
|
"rotation_deg_B",
|
|
"translation_m_B",
|
|
"weight",
|
|
"duration_s",
|
|
"mean_gyro_norm",
|
|
"preint_sigma_rad",
|
|
"cov",
|
|
"J_bg",
|
|
"phase_a_metadata_rehydrated",
|
|
"rehydrated_R_A_error_deg",
|
|
"t_i_imu_s",
|
|
"t_j_imu_s",
|
|
"gyro_bias0_rad_s",
|
|
"accel_bias0_m_s2",
|
|
"time_offset_s",
|
|
"keyframe_span",
|
|
"is_consecutive",
|
|
"modeling",
|
|
}
|
|
)
|
|
|
|
|
|
def _to_list(value: Any) -> Any:
|
|
if isinstance(value, np.ndarray):
|
|
return value.tolist()
|
|
if isinstance(value, (np.floating, np.integer, np.bool_)):
|
|
return value.item()
|
|
return value
|
|
|
|
|
|
def pair_to_dict(pair: MotionPair) -> dict[str, Any]:
|
|
meta = {
|
|
str(k): _to_list(v)
|
|
for k, v in (pair.metadata or {}).items()
|
|
if str(k) in _METADATA_KEEP
|
|
}
|
|
return {
|
|
"session_id": pair.session_id,
|
|
"i": int(pair.i),
|
|
"j": int(pair.j),
|
|
"t_i_s": float(pair.t_i_s),
|
|
"t_j_s": float(pair.t_j_s),
|
|
"R_A": np.asarray(pair.R_A, dtype=float).reshape(3, 3).tolist(),
|
|
"R_B": np.asarray(pair.R_B, dtype=float).reshape(3, 3).tolist(),
|
|
"t_A_m": None if pair.t_A_m is None else np.asarray(pair.t_A_m, dtype=float).reshape(3).tolist(),
|
|
"t_B_m": None if pair.t_B_m is None else np.asarray(pair.t_B_m, dtype=float).reshape(3).tolist(),
|
|
"fitness": float(pair.fitness),
|
|
"metadata": meta,
|
|
}
|
|
|
|
|
|
def pair_from_dict(payload: dict[str, Any]) -> MotionPair:
|
|
t_a = payload.get("t_A_m")
|
|
t_b = payload.get("t_B_m")
|
|
return MotionPair(
|
|
session_id=str(payload.get("session_id", "")),
|
|
i=int(payload["i"]),
|
|
j=int(payload["j"]),
|
|
t_i_s=float(payload["t_i_s"]),
|
|
t_j_s=float(payload["t_j_s"]),
|
|
R_A=np.asarray(payload["R_A"], dtype=float).reshape(3, 3),
|
|
R_B=np.asarray(payload["R_B"], dtype=float).reshape(3, 3),
|
|
t_A_m=None if t_a is None else np.asarray(t_a, dtype=float).reshape(3),
|
|
t_B_m=None if t_b is None else np.asarray(t_b, dtype=float).reshape(3),
|
|
fitness=float(payload.get("fitness", 0.0)),
|
|
metadata=dict(payload.get("metadata") or {}),
|
|
)
|
|
|
|
|
|
def build_motion_pairs_payload(
|
|
*,
|
|
prepared_sessions: list[dict[str, Any]],
|
|
) -> dict[str, Any]:
|
|
"""Build a JSON-serializable cache from pipeline ``prepared`` session dicts."""
|
|
|
|
sessions_out: list[dict[str, Any]] = []
|
|
for prep in prepared_sessions:
|
|
pairs = prep.get("pairs") or ()
|
|
sessions_out.append(
|
|
{
|
|
"session_id": prep.get("session_id"),
|
|
"delta_t_s": float(prep.get("time_offset_s", 0.0)),
|
|
"gyro_bias_rad_s": np.asarray(prep.get("gyro_bias_rad_s", np.zeros(3)), dtype=float)
|
|
.reshape(3)
|
|
.tolist(),
|
|
"pair_count": len(pairs),
|
|
"pairs": [pair_to_dict(pair) for pair in pairs],
|
|
}
|
|
)
|
|
return {
|
|
"schema_version": SCHEMA_VERSION,
|
|
"sessions": sessions_out,
|
|
"note": "Cached motion pairs for visualization; A=IMU preintegration, B=LiDAR registration",
|
|
}
|
|
|
|
|
|
def save_motion_pairs(path: Path | str, payload: dict[str, Any]) -> Path:
|
|
destination = Path(path)
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
destination.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
|
return destination
|
|
|
|
|
|
def load_motion_pairs(path: Path | str) -> dict[str, Any]:
|
|
payload = json.loads(Path(path).read_text(encoding="utf-8"))
|
|
version = int(payload.get("schema_version", 0))
|
|
if version not in {1, SCHEMA_VERSION}:
|
|
raise ValueError(
|
|
f"unsupported motion_pairs schema_version={version}; "
|
|
f"expected 1 or {SCHEMA_VERSION}"
|
|
)
|
|
return payload
|
|
|
|
|
|
def pairs_for_session(payload: dict[str, Any], session_id: str | None = None) -> list[MotionPair]:
|
|
sessions = payload.get("sessions") or []
|
|
if not sessions:
|
|
return []
|
|
if session_id is None:
|
|
chosen = sessions[0]
|
|
else:
|
|
chosen = next((s for s in sessions if s.get("session_id") == session_id), None)
|
|
if chosen is None:
|
|
raise KeyError(f"session_id {session_id!r} not found in motion_pairs cache")
|
|
return [pair_from_dict(item) for item in chosen.get("pairs") or []]
|
|
|
|
|
|
def resolve_motion_pairs_path(summary_path: Path | str) -> Path | None:
|
|
"""Return ``motion_pairs.json`` next to a summary if it exists."""
|
|
|
|
summary = Path(summary_path)
|
|
candidate = summary.parent / "motion_pairs.json"
|
|
return candidate if candidate.is_file() else None
|