73 lines
2.5 KiB
Python
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="")
|