75 lines
2.5 KiB
Python
75 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
|
|
|
|
import csv
|
|
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:
|
|
required_order = ["t", "gx", "gy", "gz", "ax", "ay", "az"]
|
|
with path.open("r", encoding="utf-8-sig", newline="") as handle:
|
|
header = next(csv.reader(handle), [])
|
|
names = set(header)
|
|
if not set(required_order).issubset(names):
|
|
raise ValueError(f"IMU CSV must contain columns {sorted(required_order)}, got {sorted(names)}")
|
|
usecols = [header.index(name) for name in required_order]
|
|
data = np.loadtxt(path, delimiter=",", skiprows=1, usecols=usecols, ndmin=2)
|
|
t = np.asarray(data[:, 0], dtype=float).reshape(-1)
|
|
gyro = np.asarray(data[:, 1:4], dtype=float)
|
|
acc = np.asarray(data[:, 4:7], dtype=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="")
|