feat: pipeline-based sample building (#25)

yaak-ai · Nov 28, 2024 · cc5b915 · cc5b915
1 parent a7db0a9
commit cc5b915
Show file tree

Hide file tree

Showing 62 changed files with 1,072 additions and 1,233 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -42,7 +42,7 @@ jobs:
           only: ytt
 
       - name: setup uv
-        uses: astral-sh/setup-uv@v3
+        uses: astral-sh/setup-uv@v4
         with:
           version: "latest"
           enable-cache: true

diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml
@@ -33,7 +33,7 @@ jobs:
           fetch-depth: 0
 
       - name: build and inspect
-        uses: hynek/build-and-inspect-python-package@v2.9.0
+        uses: hynek/build-and-inspect-python-package@v2
         with:
           attest-build-provenance-github: "true"
 

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -11,14 +11,14 @@ repos:
       - id: pyupgrade
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.7.4
+    rev: v0.8.0
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
 
   - repo: https://github.com/DetachHead/basedpyright-pre-commit-mirror
-    rev: 1.21.1
+    rev: 1.22.0
     hooks:
       - id: basedpyright
 

diff --git a/config/_templates/dataset/carla.yaml b/config/_templates/dataset/carla.yaml
@@ -75,39 +75,65 @@ inputs:
 
       #@ end
 
-    table_builder:
-      _target_: rbyte.io.table.TableBuilder
-      _convert_: all
-      readers:
-        ego_logs:
-          path: ${data_dir}/(@=input_id@)/ego_logs.json
-          reader:
-            _target_: rbyte.io.JsonTableReader
-            _recursive_: false
-            fields:
-              records:
-                _idx_:
-                control.brake:
-                control.throttle:
-                control.steer:
-                state.velocity.value:
-                state.acceleration.value:
+    samples:
+      pipeline:
+        _target_: pipefunc.Pipeline
+        validate_type_annotations: false
+        functions:
+          - _target_: pipefunc.PipeFunc
+            bound:
+              path: ${data_dir}/(@=input_id@)/ego_logs.json
+            output_name: ego_logs
+            func:
+              _target_: rbyte.io.JsonDataFrameBuilder
+              fields:
+                records:
+                  control.brake:
+                  control.throttle:
+                  control.steer:
+                  state.velocity.value:
+                  state.acceleration.value:
 
-            transforms:
-              - _target_: rbyte.io.FpsResampler
-                source_fps: 20
-                target_fps: 30
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: ego_logs
+            output_name: data
+            func:
+              _target_: rbyte.io.DataFrameConcater
+              method: vertical
 
-      merger:
-        _target_: rbyte.io.TableConcater
-        method: vertical
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data
+            output_name: data_resampled
+            func:
+              _target_: rbyte.io.DataFrameFpsResampler
+              fps_in: 20
+              fps_out: 30
 
-      filter: |
-        `control.throttle` > 0.5
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_resampled
+            output_name: data_indexed
+            func:
+              _target_: rbyte.io.DataFrameIndexer
+              name: _idx_
 
-  #@ end
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_indexed
+            output_name: data_filtered
+            func:
+              _target_: rbyte.io.DataFrameFilter
+              predicate: |
+                `control.throttle` > 0.5
 
-sample_builder:
-  _target_: rbyte.RollingWindowSampleBuilder
-  index_column: _idx_
-  period: 1i
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_filtered
+            output_name: samples
+            func:
+              _target_: rbyte.RollingWindowSampleBuilder
+              index_column: _idx_
+              period: 1i
+  #@ end
diff --git a/config/_templates/dataset/mimicgen.yaml b/config/_templates/dataset/mimicgen.yaml
@@ -12,8 +12,8 @@
 #@ ]
 ---
 _target_: rbyte.Dataset
-_convert_: all
 _recursive_: false
+_convert_: all
 inputs:
   #@ for input_id, input_keys in inputs.items():
   #@ for input_key in input_keys:
@@ -28,27 +28,45 @@ inputs:
           key: (@=input_key@)/(@=frame_key@)
       #@ end
 
-    table_builder:
-      _target_: rbyte.io.TableBuilder
-      _convert_: all
-      readers:
-        hdf5:
-          path: "${data_dir}/(@=input_id@).hdf5"
-          reader:
-            _target_: rbyte.io.Hdf5TableReader
-            _recursive_: false
-            fields:
-              (@=input_key@):
-                _idx_:
-                obs/robot0_eef_pos:
+    samples:
+      pipeline:
+        _target_: pipefunc.Pipeline
+        validate_type_annotations: false
+        functions:
+          - _target_: pipefunc.PipeFunc
+            bound:
+              path: "${data_dir}/(@=input_id@).hdf5"
+            output_name: data
+            func:
+              _target_: rbyte.io.Hdf5DataFrameBuilder
+              fields:
+                (@=input_key@):
+                  obs/robot0_eef_pos:
+
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data
+            output_name: data_indexed
+            func:
+              _target_: rbyte.io.DataFrameIndexer
+              name: _idx_
+
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_indexed
+            output_name: data_concated
+            func:
+              _target_: rbyte.io.DataFrameConcater
+              method: vertical
+
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_concated
+            output_name: samples
+            func:
+              _target_: rbyte.RollingWindowSampleBuilder
+              index_column: _idx_
+              period: 1i
 
-      merger:
-        _target_: rbyte.io.TableConcater
-        method: vertical
   #@ end
   #@ end
-
-sample_builder:
-  _target_: rbyte.RollingWindowSampleBuilder
-  index_column: _idx_
-  period: 1i
diff --git a/config/_templates/dataset/nuscenes/mcap.yaml b/config/_templates/dataset/nuscenes/mcap.yaml
@@ -11,15 +11,15 @@
 #@ }
 ---
 _target_: rbyte.Dataset
-_convert_: all
 _recursive_: false
+_convert_: all
 inputs:
   #@ for input_id in inputs:
   (@=input_id@):
     sources:
       #@ for camera, topic in camera_topics.items():
       (@=camera@):
-        index_column: mcap/(@=topic@)/_idx_
+        index_column: (@=topic@)/_idx_
         source:
           _target_: rbyte.io.McapTensorSource
           path: "${data_dir}/(@=input_id@).mcap"
@@ -33,66 +33,85 @@ inputs:
             fastupsample: true
       #@ end
 
-    table_builder:
-      _target_: rbyte.io.TableBuilder
-      _convert_: all
-      readers:
-        mcap:
-          path: "${data_dir}/(@=input_id@).mcap"
-          reader:
-            _target_: rbyte.io.McapTableReader
-            _recursive_: false
-            decoder_factories:
-              - rbyte.utils.mcap.ProtobufDecoderFactory
-              - rbyte.utils.mcap.JsonDecoderFactory
+    samples:
+      pipeline:
+        _target_: pipefunc.Pipeline
+        validate_type_annotations: false
+        functions:
+          - _target_: pipefunc.PipeFunc
+            bound:
+              path: ${data_dir}/(@=input_id@).mcap
+            output_name: data
+            func:
+              _target_: rbyte.io.McapDataFrameBuilder
+              decoder_factories:
+                - rbyte.utils._mcap.ProtobufDecoderFactory
+                - rbyte.utils._mcap.JsonDecoderFactory
+              fields:
+                #@ for topic in camera_topics.values():
+                (@=topic@):
+                  log_time:
+                    _target_: polars.Datetime
+                    time_unit: ns
+                #@ end
 
-            fields:
-              #@ for topic in camera_topics.values():
-              (@=topic@):
-                _idx_:
-                log_time:
-                  _target_: polars.Datetime
-                  time_unit: ns
-              #@ end
+                /odom:
+                  log_time:
+                    _target_: polars.Datetime
+                    time_unit: ns
+                  vel.x:
 
-              /odom:
-                log_time:
-                  _target_: polars.Datetime
-                  time_unit: ns
-                vel.x:
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data
+            output_name: data_indexed
+            func:
+              _target_: rbyte.io.DataFrameIndexer
+              name: _idx_
 
-      merger:
-        _target_: rbyte.io.TableAligner
-        separator: "/"
-        merge:
-          mcap:
-            #@ topic = camera_topics.values()[0]
-            (@=topic@):
-              key: log_time
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_indexed
+            output_name: data_aligned
+            func:
+              _target_: rbyte.io.DataFrameAligner
+              separator: /
+              fields:
+                #@ topic = camera_topics.values()[0]
+                (@=topic@):
+                  key: log_time
 
-            #@ for topic in camera_topics.values()[1:]:
-            (@=topic@):
-              key: log_time
-              columns:
-                _idx_:
-                  method: asof
-                  tolerance: 40ms
-                  strategy: nearest
-            #@ end
+                #@ for topic in camera_topics.values()[1:]:
+                (@=topic@):
+                  key: log_time
+                  columns:
+                    _idx_:
+                      method: asof
+                      tolerance: 40ms
+                      strategy: nearest
+                #@ end
 
-            /odom:
-              key: log_time
-              columns:
-                vel.x:
-                  method: interp
+                /odom:
+                  key: log_time
+                  columns:
+                    vel.x:
+                      method: interp
 
-      filter: |
-        `mcap//odom/vel.x` >= 8
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_aligned
+            output_name: data_filtered
+            func:
+              _target_: rbyte.io.DataFrameFilter
+              predicate: |
+                `/odom/vel.x` >= 8
 
-      cache:
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_filtered
+            output_name: samples
+            func:
+              _target_: rbyte.RollingWindowSampleBuilder
+              index_column: (@=camera_topics.values()[0]@)/_idx_
+              period: 1i
   #@ end
-
-sample_builder:
-  _target_: rbyte.RollingWindowSampleBuilder
-  index_column: mcap/(@=camera_topics.values()[0]@)/_idx_
-  period: 1i