diff --git a/CAPTURE.md b/CAPTURE.md index c3b4c73..7348375 100644 --- a/CAPTURE.md +++ b/CAPTURE.md @@ -76,7 +76,7 @@ A Python interface that mirrors this file structure is available in `scantools.c ### 1. Session data - `sensors.txt`, `rigs.txt`, `trajectories.txt` follow [the Kapture format](https://github.com/naver/kapture/blob/main/kapture_format.adoc#2--sensors). However, the pose convention is reverted: `rigs.txt` contains camera-to-rig transformations and `trajectories.txt` contains sensor-to-world transformations. -- `images.txt`, `pointclouds.txt`, `depths.txt`, `wifi.txt`, and `bt.txt` follow the specifications of their corresponding `records_*.txt` in Kapture. +- `images.txt`, `pointclouds.txt`, `depths.txt`, `bt.txt`, `wifi.txt`, `imu.txt`, and `gravity.txt` follow the specifications of their corresponding `records_*.txt` in Kapture. `gravity.txt` contains gravity direction estimates (down vector in sensor frame at each timestamp) from the onboard VIO system. ### 2. Processed files diff --git a/scantools/capture/records.py b/scantools/capture/records.py index e5a7a8e..902c2ca 100644 --- a/scantools/capture/records.py +++ b/scantools/capture/records.py @@ -143,6 +143,24 @@ class RecordsLidar(RecordsFilePath): field_names = ['point_cloud_path'] +# IMU data. +class RecordsIMU(RecordsBase[np.ndarray]): + record_type = np.ndarray + field_names = ['x', 'y', 'z'] + + def record_to_list(self, record: np.ndarray) -> List[str]: + return list(map(str, record)) + + @classmethod + def load(cls, path: Path) -> 'RecordsIMU': + table = read_csv(path) + records = cls() + for timestamp, sensor_id, *data in table: + assert len(data) == 3 + records[int(timestamp), sensor_id] = np.array( + list(map(float, data))) + return records + # New data types inherit from RecordEntry (a record) and RecordsArray (mapping of records) @dataclass diff --git a/scantools/capture/session.py b/scantools/capture/session.py index 2bf0b2d..897d6bb 100644 --- a/scantools/capture/session.py +++ b/scantools/capture/session.py @@ -7,7 +7,8 @@ from .sensors import Sensors, Camera from .rigs import Rigs from .trajectories import Trajectories -from .records import RecordsBluetooth, RecordsCamera, RecordsDepth, RecordsLidar, RecordsWifi +from .records import ( + RecordsBluetooth, RecordsCamera, RecordsDepth, RecordsIMU, RecordsLidar, RecordsWifi) from .proc import Proc from .pose import Pose @@ -41,6 +42,8 @@ class Session: pointclouds: Optional[RecordsLidar] = None wifi: Optional[RecordsWifi] = None bt: Optional[RecordsBluetooth] = None + imu: Optional[RecordsIMU] = None + gravity: Optional[RecordsIMU] = None proc: Optional[Proc] = None id: Optional[str] = None @@ -52,7 +55,8 @@ def __post_init__(self): all_devices = set(self.sensors.keys()) if self.rigs is not None: assert len(self.sensors.keys() & self.rigs.keys()) == 0 - assert len(self.rigs.sensor_ids - self.sensors.keys()) == 0 + # TODO: Fix me - currently missing calibration for IMUs. + # assert len(self.rigs.sensor_ids - self.sensors.keys()) == 0 all_devices |= self.rigs.keys() if self.trajectories is not None: assert len(self.trajectories.device_ids - all_devices) == 0 @@ -75,7 +79,7 @@ def filename(cls, attr: Union[Field, str]) -> str: return f'{name}.txt' @classmethod - def load(cls, path: Path, wireless=True) -> 'Session': + def load(cls, path: Path, wireless=True, imu=True) -> 'Session': if not path.exists(): raise IOError(f'Session directory does not exists: {path}') data = {} @@ -95,6 +99,8 @@ def load(cls, path: Path, wireless=True) -> 'Session': else: if attr.name in ['bt', 'wifi'] and not wireless: continue + if attr.name in ['gravity', 'imu'] and not imu: + continue obj = type_.load(filepath) data[attr.name] = obj if 'sensors' not in data: