128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
"""LiDAR adapters for the V1 standard intermediate format.
|
|
|
|
Accepted input: a directory containing ``frames_index.csv`` and per-frame NPZ files.
|
|
|
|
frames_index.csv
|
|
----------------
|
|
frame_id,file,t_start,t_end
|
|
|
|
Each NPZ referenced by ``file`` must contain:
|
|
- points: float array shaped (N, 3) in LiDAR Cartesian coordinates (metres)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from .contracts import LidarFrame
|
|
|
|
|
|
def _read_frames_index(root: Path) -> tuple[np.ndarray, str]:
|
|
index_path = root / "frames_index.csv"
|
|
if not index_path.exists():
|
|
raise FileNotFoundError(f"missing frames_index.csv under {root}")
|
|
|
|
rows = np.genfromtxt(index_path, delimiter=",", names=True, dtype=None, encoding="utf-8")
|
|
if rows.ndim == 0:
|
|
rows = np.array([rows])
|
|
names = set(rows.dtype.names or ())
|
|
# NumPy may rename reserved name ``file`` to ``file_``.
|
|
file_key = "filename" if "filename" in names else ("file_" if "file_" in names else "file")
|
|
required = {"frame_id", "t_start", "t_end"}
|
|
if not required.issubset(names) or file_key not in names:
|
|
raise ValueError(
|
|
f"frames_index.csv must contain frame_id,{file_key}/filename,t_start,t_end; got {sorted(names)}"
|
|
)
|
|
return rows, file_key
|
|
|
|
|
|
def list_lidar_frame_entries(path: Path | str) -> list[tuple[str, float, float, Path]]:
|
|
"""Return ``(frame_id, t_start, t_end, npz_path)`` sorted by mid time (same as ``load_lidar_frames``)."""
|
|
|
|
root = Path(path)
|
|
rows, file_key = _read_frames_index(root)
|
|
entries: list[tuple[str, float, float, Path]] = []
|
|
for row in rows:
|
|
t0 = float(row["t_start"])
|
|
t1 = float(row["t_end"])
|
|
entries.append((str(row["frame_id"]), t0, t1, root / str(row[file_key])))
|
|
entries.sort(key=lambda item: 0.5 * (item[1] + item[2]))
|
|
return entries
|
|
|
|
|
|
def load_lidar_frame_at(root: Path | str, index: int) -> LidarFrame:
|
|
"""Load one frame by index in mid-time-sorted order (matches motion-pair ``i``/``j``)."""
|
|
|
|
entries = list_lidar_frame_entries(root)
|
|
if index < 0 or index >= len(entries):
|
|
raise IndexError(f"frame index {index} outside [0, {len(entries) - 1}] for {root}")
|
|
frame_id, t0, t1, npz_path = entries[index]
|
|
with np.load(npz_path) as payload:
|
|
if "points" not in payload.files:
|
|
raise ValueError(f"{npz_path} must contain array 'points'")
|
|
points = np.asarray(payload["points"], dtype=float)
|
|
if points.ndim != 2 or points.shape[1] < 3:
|
|
raise ValueError(f"{npz_path}: points must have shape (N, 3[+])")
|
|
return LidarFrame(
|
|
frame_id=frame_id,
|
|
t_start_s=t0,
|
|
t_end_s=t1,
|
|
points_xyz=points[:, :3],
|
|
path=npz_path,
|
|
)
|
|
|
|
|
|
def lidar_frame_count(path: Path | str) -> int:
|
|
return len(list_lidar_frame_entries(path))
|
|
|
|
|
|
def load_lidar_frames(path: Path | str) -> list[LidarFrame]:
|
|
"""Load all LiDAR frames listed by ``frames_index.csv`` under ``path``."""
|
|
|
|
root = Path(path)
|
|
rows, file_key = _read_frames_index(root)
|
|
|
|
frames: list[LidarFrame] = []
|
|
for row in rows:
|
|
frame_id = str(row["frame_id"])
|
|
rel = str(row[file_key])
|
|
npz_path = root / rel
|
|
with np.load(npz_path) as payload:
|
|
if "points" not in payload.files:
|
|
raise ValueError(f"{npz_path} must contain array 'points'")
|
|
points = np.asarray(payload["points"], dtype=float)
|
|
if points.ndim != 2 or points.shape[1] < 3:
|
|
raise ValueError(f"{npz_path}: points must have shape (N, 3[+])")
|
|
frames.append(
|
|
LidarFrame(
|
|
frame_id=frame_id,
|
|
t_start_s=float(row["t_start"]),
|
|
t_end_s=float(row["t_end"]),
|
|
points_xyz=points[:, :3],
|
|
path=npz_path,
|
|
)
|
|
)
|
|
frames.sort(key=lambda frame: frame.t_mid_s)
|
|
return frames
|
|
|
|
|
|
def save_lidar_session(
|
|
root: Path | str,
|
|
frames: list[LidarFrame],
|
|
*,
|
|
points_dirname: str = "frames",
|
|
) -> None:
|
|
"""Write a LiDAR session directory in the standard intermediate format."""
|
|
|
|
destination = Path(root)
|
|
frames_dir = destination / points_dirname
|
|
frames_dir.mkdir(parents=True, exist_ok=True)
|
|
index_rows: list[str] = ["frame_id,filename,t_start,t_end"]
|
|
for index, frame in enumerate(frames):
|
|
relative = f"{points_dirname}/frame_{index:05d}.npz"
|
|
np.savez_compressed(destination / relative, points=np.asarray(frame.points_xyz, dtype=float))
|
|
index_rows.append(f"{frame.frame_id},{relative},{frame.t_start_s:.9f},{frame.t_end_s:.9f}")
|
|
(destination / "frames_index.csv").write_text("\n".join(index_rows) + "\n", encoding="utf-8")
|