feat: support multiple table readers (#15)

- add support for multiple table readers - add a `protobuf -> pl.DataFrame` mcap decoder factory - replace `CarlaRecordsTableBuilder` with `JsonTableReader`
yaak-ai · Oct 21, 2024 · 5e99824 · 5e99824
1 parent 3d94f2b
commit 5e99824
Show file tree

Hide file tree

Showing 38 changed files with 684 additions and 423 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -3,7 +3,7 @@ fail_fast: true
 
 repos:
   - repo: https://github.com/abravalheri/validate-pyproject
-    rev: v0.20.2
+    rev: v0.21
     hooks:
       - id: validate-pyproject
 
@@ -13,7 +13,7 @@ repos:
       - id: typos
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.17.0
+    rev: v3.18.0
     hooks:
       - id: pyupgrade
 
@@ -31,7 +31,7 @@ repos:
       - id: ruff-format
 
   - repo: https://github.com/DetachHead/basedpyright-pre-commit-mirror
-    rev: 1.18.3
+    rev: 1.18.4
     hooks:
       - id: basedpyright
 

diff --git a/README.md b/README.md
@@ -139,40 +139,41 @@ dataset:
             frame_decoder: ${frame_decoder}
 
       table:
-        path: data/NuScenes-v1.0-mini-scene-0103.mcap
         builder:
           _target_: rbyte.io.table.TableBuilder
           _convert_: all
-          reader:
-            _target_: rbyte.io.table.mcap.McapTableReader
-            _recursive_: false
-            decoder_factories:
-              - mcap_protobuf.decoder.DecoderFactory
-              - rbyte.utils.mcap.McapJsonDecoderFactory
-            fields:
-              /CAM_FRONT/image_rect_compressed:
-                _idx_:
-                log_time:
-                  _target_: polars.Datetime
-                  time_unit: ns
-
-              /CAM_FRONT_LEFT/image_rect_compressed:
-                _idx_:
-                log_time:
-                  _target_: polars.Datetime
-                  time_unit: ns
-
-              /CAM_FRONT_RIGHT/image_rect_compressed:
-                _idx_:
-                log_time:
-                  _target_: polars.Datetime
-                  time_unit: ns
-
-              /odom:
-                log_time:
-                  _target_: polars.Datetime
-                  time_unit: ns
-                vel.x: null
+          readers:
+            - path: data/NuScenes-v1.0-mini-scene-0103.mcap
+              reader:
+                _target_: rbyte.io.table.mcap.McapTableReader
+                _recursive_: false
+                decoder_factories:
+                  - mcap_protobuf.decoder.DecoderFactory
+                  - rbyte.utils.mcap.McapJsonDecoderFactory
+                fields:
+                  /CAM_FRONT/image_rect_compressed:
+                    _idx_:
+                    log_time:
+                      _target_: polars.Datetime
+                      time_unit: ns
+
+                  /CAM_FRONT_LEFT/image_rect_compressed:
+                    _idx_:
+                    log_time:
+                      _target_: polars.Datetime
+                      time_unit: ns
+
+                  /CAM_FRONT_RIGHT/image_rect_compressed:
+                    _idx_:
+                    log_time:
+                      _target_: polars.Datetime
+                      time_unit: ns
+
+                  /odom:
+                    log_time:
+                      _target_: polars.Datetime
+                      time_unit: ns
+                    vel.x: null
 
           merger:
             _target_: rbyte.io.table.TableAligner

diff --git a/examples/config_templates/build_table.yaml b/examples/config_templates/build_table.yaml
@@ -1,11 +1,9 @@
 ---
 defaults:
   - table_builder: !!null
-  - table_writer: !!null
+  - table_writer: console
   - _self_
 
-path: ???
-
 hydra:
   output_subdir: !!null
   run:

diff --git a/examples/config_templates/dataset/carla.yaml b/examples/config_templates/dataset/carla.yaml
@@ -62,7 +62,7 @@ inputs:
     frame:
       #@ for source_id in cameras:
       (@=source_id@):
-        index_column: frame_idx
+        index_column: _idx_
         reader:
           _target_: rbyte.io.frame.DirectoryFrameReader
           path: "${data_dir}/(@=input_id@)/frames/(@=source_id@).defish.mp4/576x324/{:09d}.jpg"
@@ -76,28 +76,40 @@ inputs:
       #@ end
 
     table:
-      path: ${data_dir}/(@=input_id@)/ego_logs.json
       builder:
-        _target_: rbyte.io.table.carla.CarlaRecordsTableBuilder
+        _target_: rbyte.io.table.TableBuilder
         _convert_: all
-        index_column: frame_idx
-        select:
-          - control.brake
-          - control.throttle
-          - control.steer
-          - state.velocity.value
-          - state.acceleration.value
+        readers:
+          - path: ${data_dir}/(@=input_id@)/ego_logs.json
+            reader:
+              _target_: rbyte.io.table.JsonTableReader
+              _recursive_: false
+              fields:
+                records:
+                  _idx_:
+                  control.brake:
+                  control.throttle:
+                  control.steer:
+                  state.velocity.value:
+                  state.acceleration.value:
+
+              transforms:
+                - _target_: rbyte.io.table.transforms.FpsResampler
+                  source_fps: 20
+                  target_fps: 30
+
+        merger:
+          _target_: rbyte.io.table.TableConcater
+          method: vertical
+
+        filter: |
+          `control.throttle` > 0.5
 
-        filter: !!null
-        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
+  index_column: _idx_
   length: 1
   stride: 1
   min_step: 1

diff --git a/examples/config_templates/dataset/hdf5.yaml b/examples/config_templates/dataset/hdf5.yaml
@@ -36,22 +36,23 @@ inputs:
       #@ end
 
     table:
-      path: "${data_dir}/(@=input_id@).hdf5"
       builder:
         _target_: rbyte.io.table.TableBuilder
         _convert_: all
-        reader:
-          _target_: rbyte.io.table.hdf5.Hdf5TableReader
-          _recursive_: false
-          fields:
-            /data/demo_0:
-              _idx_:
-              obs/object:
-              task_successes:
+        readers:
+          - path: "${data_dir}/(@=input_id@).hdf5"
+            reader:
+              _target_: rbyte.io.table.hdf5.Hdf5TableReader
+              _recursive_: false
+              fields:
+                (@=input_key@):
+                  _idx_:
+                  obs/object:
+                  task_successes:
 
         merger:
           _target_: rbyte.io.table.TableConcater
-          separator: "/"
+          method: vertical
   #@ end
   #@ end
 

diff --git a/examples/config_templates/dataset/mcap.yaml b/examples/config_templates/dataset/mcap.yaml
@@ -36,32 +36,33 @@ inputs:
       #@ end
 
     table:
-      path: "${data_dir}/(@=input_id@).mcap"
       builder:
         _target_: rbyte.io.table.TableBuilder
         _convert_: all
-        reader:
-          _target_: rbyte.io.table.mcap.McapTableReader
-          _recursive_: false
-          decoder_factories:
-            - mcap_protobuf.decoder.DecoderFactory
-            - rbyte.utils.mcap.McapJsonDecoderFactory
+        readers:
+          - path: "${data_dir}/(@=input_id@).mcap"
+            reader:
+              _target_: rbyte.io.table.mcap.McapTableReader
+              _recursive_: false
+              decoder_factories:
+                - rbyte.utils.mcap.ProtobufDecoderFactory
+                - rbyte.utils.mcap.JsonDecoderFactory
 
-          fields:
-            #@ for topic in camera_topics:
-            (@=topic@):
-              log_time:
-                _target_: polars.Datetime
-                time_unit: ns
+              fields:
+                #@ for topic in camera_topics:
+                (@=topic@):
+                  log_time:
+                    _target_: polars.Datetime
+                    time_unit: ns
 
-              _idx_:
-            #@ end
+                  _idx_:
+                #@ end
 
-            /odom:
-              log_time:
-                _target_: polars.Datetime
-                time_unit: ns
-              vel.x:
+                /odom:
+                  log_time:
+                    _target_: polars.Datetime
+                    time_unit: ns
+                  vel.x:
 
         merger:
           _target_: rbyte.io.table.TableAligner