Adds support for IMU measurements.

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
 

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

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: