Files
calibration/imu_lidar/imu_io.py

73 lines
2.5 KiB
Python

"""IMU adapters for the V1 standard intermediate format.
Accepted inputs
---------------
1. CSV with header:
t,gx,gy,gz,ax,ay,az
- ``t`` in seconds on the IMU clock
- gyro in rad/s
- accel in m/s^2
2. NPZ with arrays:
t, gyro, acc
shapes: (N,), (N,3), (N,3)
"""
from __future__ import annotations
from pathlib import Path
import numpy as np
from .contracts import ImuSeries
def load_imu_samples(path: Path | str) -> ImuSeries:
"""Load normalized IMU samples from CSV or NPZ."""
source = Path(path)
if not source.exists():
raise FileNotFoundError(source)
if source.suffix.lower() == ".csv":
return _load_imu_csv(source)
if source.suffix.lower() == ".npz":
return _load_imu_npz(source)
raise ValueError(f"unsupported IMU format '{source.suffix}' (use .csv or .npz)")
def _load_imu_csv(path: Path) -> ImuSeries:
data = np.genfromtxt(path, delimiter=",", names=True, dtype=float)
if data.ndim == 0:
data = np.array([data])
names = set(data.dtype.names or ())
required = {"t", "gx", "gy", "gz", "ax", "ay", "az"}
if not required.issubset(names):
raise ValueError(f"IMU CSV must contain columns {sorted(required)}, got {sorted(names)}")
t = np.asarray(data["t"], dtype=float).reshape(-1)
gyro = np.column_stack([data["gx"], data["gy"], data["gz"]]).astype(float)
acc = np.column_stack([data["ax"], data["ay"], data["az"]]).astype(float)
order = np.argsort(t)
return ImuSeries(t_s=t[order], gyro_rad_s=gyro[order], acc_m_s2=acc[order])
def _load_imu_npz(path: Path) -> ImuSeries:
with np.load(path) as payload:
keys = set(payload.files)
if not {"t", "gyro", "acc"}.issubset(keys):
raise ValueError(f"IMU NPZ must contain t, gyro, acc; got {sorted(keys)}")
t = np.asarray(payload["t"], dtype=float).reshape(-1)
gyro = np.asarray(payload["gyro"], dtype=float).reshape(-1, 3)
acc = np.asarray(payload["acc"], dtype=float).reshape(-1, 3)
order = np.argsort(t)
return ImuSeries(t_s=t[order], gyro_rad_s=gyro[order], acc_m_s2=acc[order])
def save_imu_csv(path: Path | str, imu: ImuSeries) -> None:
"""Write IMU samples to the standard CSV format."""
destination = Path(path)
destination.parent.mkdir(parents=True, exist_ok=True)
array = np.column_stack([imu.t_s, imu.gyro_rad_s, imu.acc_m_s2])
header = "t,gx,gy,gz,ax,ay,az"
np.savetxt(destination, array, delimiter=",", header=header, comments="")