parse_imu_vts_lib.py

"""
parse_imu_vts_lib.py — Shared library for IMU/VTS binary parsing.
This is the canonical parser used by all other scripts.
See 01_parse_imu_vts.py for full format documentation.
"""

import struct
import numpy as np
from dataclasses import dataclass
from typing import List

# ── Data structures ────────────────────────────────────────────────────────────

@dataclass
class IMURecord:
    timestamp_ns: int
    accel: np.ndarray        # (3,) m/s²
    gyro: np.ndarray         # (3,) rad/s
    mag: np.ndarray          # (3,) µT
    temperature: float       # °C

    @property
    def timestamp_s(self) -> float:
        return self.timestamp_ns * 1e-9

    @property
    def timestamp_us(self) -> float:
        return self.timestamp_ns * 1e-3

@dataclass
class VTSRecord:
    frame_idx: int
    imu_idx: int
    timestamp_us: int        # microseconds

    @property
    def timestamp_s(self) -> float:
        return self.timestamp_us * 1e-6

    @property
    def timestamp_ns(self) -> int:
        return self.timestamp_us * 1000


# ── IMU Parser ─────────────────────────────────────────────────────────────────

IMU_MAGIC  = b'TRIMU001'
IMU_HEADER = 64
IMU_RECORD = 80
# Layout: u64 ts_ns, 3×f32 accel, 3×f32 gyro, 3×f32 mag, f32 temp, 32 bytes zeros
IMU_FMT    = '<Q3f3f3ff'

def parse_imu(path: str) -> List[IMURecord]:
    with open(path, 'rb') as f:
        data = f.read()
    if data[:8] != IMU_MAGIC:
        raise ValueError(f"Bad IMU magic in {path}")
    version = struct.unpack_from('<I', data, 8)[0]
    n_data  = len(data) - IMU_HEADER
    assert n_data % IMU_RECORD == 0, f"IMU data not divisible by {IMU_RECORD}"
    n = n_data // IMU_RECORD
    records = []
    for i in range(n):
        off = IMU_HEADER + i * IMU_RECORD
        ts, ax, ay, az, gx, gy, gz, mx, my, mz, temp = struct.unpack_from(IMU_FMT, data, off)
        records.append(IMURecord(
            timestamp_ns=ts,
            accel=np.array([ax, ay, az], dtype=np.float32),
            gyro=np.array([gx, gy, gz], dtype=np.float32),
            mag=np.array([mx, my, mz], dtype=np.float32),
            temperature=temp,
        ))
    rate = (n - 1) / (records[-1].timestamp_s - records[0].timestamp_s)
    print(f"[IMU] {n} records | ver={version} | ~{rate:.1f} Hz | {records[-1].timestamp_s - records[0].timestamp_s:.2f}s")
    return records


# ── VTS Parser ─────────────────────────────────────────────────────────────────

VTS_MAGIC  = b'TRIVTS01'
VTS_HEADER = 32
VTS_RECORD = 24
VTS_FMT    = '<IIIIIi'   # frame_idx, Hardwareclock, const, imu_idx, ts_us, reserved

def parse_vts(path: str) -> List[VTSRecord]:
    with open(path, 'rb') as f:
        data = f.read()
    if data[:8] != VTS_MAGIC:
        raise ValueError(f"Bad VTS magic in {path}")
    version   = struct.unpack_from('<I', data, 8)[0]
    fps_milli = struct.unpack_from('<I', data, 12)[0]
    n_data    = len(data) - VTS_HEADER
    assert n_data % VTS_RECORD == 0, f"VTS data not divisible by {VTS_RECORD}"
    n = n_data // VTS_RECORD
    records = []
    for i in range(n):
        off = VTS_HEADER + i * VTS_RECORD
        fi, col1, col2, imu_idx, ts_us, _ = struct.unpack_from(VTS_FMT, data, off)
        records.append(VTSRecord(frame_idx=fi, imu_idx=imu_idx, timestamp_us=ts_us))
    dur = records[-1].timestamp_s - records[0].timestamp_s
    actual_fps = (n - 1) / dur if dur > 0 else fps_milli / 1000
    print(f"[VTS] {n} records | ver={version} | fps={fps_milli/1000:.1f} | {dur:.2f}s")
    return records


# ── Sync helper ────────────────────────────────────────────────────────────────

def build_sync_index(imu_records: List[IMURecord], vts_records: List[VTSRecord]):
    """
    Match each VTS frame to the nearest IMU record by timestamp.
    Returns list of (VTSRecord, imu_index, sync_error_ms).
    """
    imu_ts_us = np.array([r.timestamp_us for r in imu_records])
    sync = []
    for vr in vts_records:
        frame_ts = vr.timestamp_us
        idx = int(np.searchsorted(imu_ts_us, frame_ts))
        idx = min(max(idx, 0), len(imu_records) - 1)
        if idx > 0:
            d_prev = abs(imu_ts_us[idx-1] - frame_ts)
            d_curr = abs(imu_ts_us[idx]   - frame_ts)
            if d_prev < d_curr:
                idx -= 1
        error_ms = (imu_ts_us[idx] - frame_ts) / 1000.0
        sync.append((vr, idx, error_ms))
    return sync