feat: mcap support

- add mcap frame and tabular readers - abstract tabular source processing - slimmer configs - add `read-frames` script
yaak-ai · Sep 18, 2024 · f894e58 · f894e58
1 parent 09d5b31
commit f894e58
Show file tree

Hide file tree

Showing 56 changed files with 1,663 additions and 1,058 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -12,8 +12,6 @@ on:
 jobs:
   pre-commit:
     runs-on: ubuntu-latest
-    env:
-      UV_CACHE_DIR: /tmp/.uv-cache
 
     steps:
       - name: setup ssh
@@ -33,16 +31,10 @@ jobs:
           python-version-file: "pyproject.toml"
 
       - name: install uv
-        run: curl -LsSf https://astral.sh/uv/install.sh | sh
-
-      - name: restore uv cache
-        uses: actions/cache@v4
+        uses: astral-sh/setup-uv@v3
         with:
-          path: /tmp/.uv-cache
-          key: uv-${{ runner.os }}-${{ hashFiles('uv.lock') }}
-          restore-keys: |
-            uv-${{ runner.os }}-${{ hashFiles('uv.lock') }}
-            uv-${{ runner.os }}
+          version: "latest"
+          enable-cache: true
 
       - name: uv sync
         run: uv sync --all-extras --dev
@@ -63,6 +55,3 @@ jobs:
         run: uvx pre-commit run --all-files --color=always
         env:
           SKIP: just-format,generate-example-config
-
-      - name: uv cache prune
-        run: uv cache prune --ci
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -8,7 +8,7 @@ repos:
       - id: validate-pyproject
 
   - repo: https://github.com/crate-ci/typos
-    rev: v1.23.7
+    rev: v1.24.6
     hooks:
       - id: typos
 
@@ -24,14 +24,14 @@ repos:
         exclude: examples/config
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.2
+    rev: v0.6.5
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
 
   - repo: https://github.com/DetachHead/basedpyright-pre-commit-mirror
-    rev: 1.17.0
+    rev: 1.17.4
     hooks:
       - id: basedpyright
 

diff --git a/README.md b/README.md
@@ -7,7 +7,7 @@ Multimodal dataset library.
 ## Installation
 
 ```bash
-uv add https://github.com/yaak-ai/rbyte/releases/latest/download/rbyte-X.Y.Z-py3-none-any.whl [--extra visualize]
+uv add https://github.com/yaak-ai/rbyte/releases/latest/download/rbyte-X.Y.Z-py3-none-any.whl [--extra visualize] [--extra jpeg] [--extra mcap]
 ```
 
 ## Usage

diff --git a/examples/config_templates/build_table.yaml b/examples/config_templates/build_table.yaml
@@ -0,0 +1,15 @@
+---
+defaults:
+  - table_builder: !!null
+  - _self_
+
+path: ???
+writer:
+  _target_: polars.DataFrame.write_csv
+  _partial_: true
+  file: ???
+
+hydra:
+  output_subdir: !!null
+  run:
+    dir: .
diff --git a/examples/config_templates/dataset/carla.yaml b/examples/config_templates/dataset/carla.yaml
@@ -1,44 +1,61 @@
+#@yaml/text-templated-strings
+
+#@ drives = [
+#@     'carla_id1',
+#@ ]
+
+#@ cameras = [
+#@     'cam_front_left',
+#@ ]
 ---
 _target_: rbyte.Dataset
 _recursive_: false
-config:
-  inputs:
-    - id: carla_id1
-      sources:
-        frame:
-          - id: cam_front_left
-            index_column: frame_idx
-            reader:
-              _target_: rbyte.io.frame.jpeg.JpegFrameReader
-              path: "${data_dir}/${.....id}/frames/${..id}.defish.mp4/576x324/{:09d}.jpg"
-
-        table:
-          path: ${data_dir}/${...id}/ego_logs.json
-          builder:
-            _target_: rbyte.io.table.carla.CarlaRecordsTableBuilder
-            _recursive_: false
-            config:
-              index_column: frame_idx
-              select:
-                - control.brake
-                - control.throttle
-                - control.steer
-                - state.velocity.value
-                - state.acceleration.value
+_convert_: all
+inputs:
+  #@ for input_id in drives:
+  (@=input_id@):
+    frame:
+      #@ for source_id in cameras:
+      (@=source_id@):
+        index_column: frame_idx
+        reader:
+          _target_: rbyte.io.frame.DirectoryFrameReader
+          path: "${data_dir}/(@=input_id@)/frames/(@=source_id@).defish.mp4/576x324/{:09d}.jpg"
+          frame_decoder:
+            _target_: simplejpeg.decode_jpeg
+            _partial_: true
+            colorspace: rgb
+            fastdct: true
+            fastupsample: true
 
-              filter: |
-                `control.throttle` > 0.5
+      #@ end
 
-              transforms:
-                - _target_: rbyte.io.table.transforms.FpsResampler
-                  source_fps: 20
-                  target_fps: 30
-  samples:
-    builder:
-      _target_: rbyte.sample.builder.GreedySampleTableBuilder
-      config:
+    table:
+      path: ${data_dir}/(@=input_id@)/ego_logs.json
+      builder:
+        _target_: rbyte.io.table.carla.CarlaRecordsTableBuilder
+        _convert_: all
         index_column: frame_idx
-        length: 1
-        stride: 1
-        min_step: 1
-        filter: !!null
+        select:
+          - control.brake
+          - control.throttle
+          - control.steer
+          - state.velocity.value
+          - state.acceleration.value
+
+        filter: |
+          `control.throttle` > 0.5
+
+        transforms:
+          - _target_: rbyte.io.table.transforms.FpsResampler
+            source_fps: 20
+            target_fps: 30
+    #@ end
+
+sample_builder:
+  _target_: rbyte.sample.builder.GreedySampleTableBuilder
+  index_column: frame_idx
+  length: 1
+  stride: 1
+  min_step: 1
+  filter: !!null
diff --git a/examples/config_templates/dataset/mcap.yaml b/examples/config_templates/dataset/mcap.yaml
@@ -0,0 +1,107 @@
+#@yaml/text-templated-strings
+
+#@ drives = [
+#@     'Niro102-HQ/2023-05-08--13-59-22',
+#@ ]
+
+#@ cameras = [
+#@     'cam_front_left',
+#@     'cam_left_forward',
+#@     'cam_right_forward',
+#@ ]
+---
+_target_: rbyte.Dataset
+_convert_: all
+_recursive_: false
+inputs:
+  #@ for input_id in drives:
+  (@=input_id@):
+    frame:
+      #@ for source_id in cameras:
+      (@=source_id@):
+        index_column: /camera/(@=source_id@)/frame_idx
+        reader:
+          _target_: rbyte.io.frame.mcap.McapFrameReader
+          path: "${data_dir}/(@=input_id@)/drive.mcap"
+          topic: /camera/(@=source_id@)
+          message_decoder_factory:
+            _target_: mcap_protobuf.decoder.DecoderFactory
+          frame_decoder:
+            _target_: simplejpeg.decode_jpeg
+            _partial_: true
+            colorspace: rgb
+            fastdct: true
+            fastupsample: true
+      #@ end
+
+    table:
+      path: "${data_dir}/(@=input_id@)/drive.mcap"
+      builder:
+        _target_: rbyte.io.table.TableBuilder
+        reader:
+          _target_: rbyte.io.table.mcap.McapProtobufTableReader
+          _recursive_: false
+          _convert_: all
+          fields:
+            #@ for camera in cameras:
+            /camera/(@=camera@):
+              timestamp:
+                _target_: polars.Datetime
+                time_unit: ns
+            #@ end
+
+            /sensor/vehicle_motion:
+              time_stamp:
+                _target_: polars.Datetime
+                time_unit: ns
+
+              speed: polars.Float32
+              gear:
+                _target_: polars.Enum
+                categories: ["0", "1", "2", "3"]
+
+        merger:
+          _target_: rbyte.io.table.TableMerger
+          separator: /
+          merge:
+            /camera/(@=cameras[0]@):
+              timestamp:
+                method: ref
+
+            #@ for camera in cameras[1:]:
+            /camera/(@=camera@):
+              timestamp:
+                method: ref
+
+              frame_idx:
+                method: asof
+                tolerance: 10ms
+                strategy: nearest
+            #@ end
+
+            /sensor/vehicle_motion:
+              time_stamp:
+                method: ref
+              speed:
+                method: interp
+              gear:
+                method: asof
+                tolerance: 100ms
+
+        filter: |
+          `/sensor/vehicle_motion/gear` == '3'
+
+        cache:
+          _target_: rbyte.utils.dataframe.DataframeDiskCache
+          directory: /tmp/rbyte-cache
+          size_limit: 1GiB
+  #@ end
+
+sample_builder:
+  _target_: rbyte.sample.builder.GreedySampleTableBuilder
+  index_column: /camera/cam_front_left/frame_idx
+  length: 6
+  stride: 1
+  min_step: 6
+  filter: |
+    array_mean(`/sensor/vehicle_motion/speed`) > 47