83 lines
3.0 KiB
Python
83 lines
3.0 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 load_lidar_frames(path: Path | str) -> list[LidarFrame]:
|
||
|
|
"""Load all LiDAR frames listed by ``frames_index.csv`` under ``path``."""
|
||
|
|
|
||
|
|
root = Path(path)
|
||
|
|
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)}"
|
||
|
|
)
|
||
|
|
|
||
|
|
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")
|