feat: add hdf5 support (#14)

yaak-ai · Oct 11, 2024 · aed93a1 · aed93a1
1 parent c1b6835
commit aed93a1
Show file tree

Hide file tree

Showing 33 changed files with 571 additions and 193 deletions.
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.24.6
+    rev: v1.26.0
     hooks:
       - id: typos
 
@@ -24,14 +24,14 @@ repos:
         exclude: examples/config
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.8
+    rev: v0.6.9
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
 
   - repo: https://github.com/DetachHead/basedpyright-pre-commit-mirror
-    rev: 1.18.0
+    rev: 1.18.3
     hooks:
       - id: basedpyright
 
@@ -40,15 +40,15 @@ repos:
       - id: just-format
         name: just-format
         language: system
-        stages: [commit]
+        stages: [pre-commit]
         entry: just --fmt --unstable
         pass_filenames: false
         always_run: true
 
       - id: generate-example-config
         name: generate-example-config
         language: system
-        stages: [commit]
+        stages: [pre-commit]
         entry: just generate-example-config
         pass_filenames: false
         always_run: true
diff --git a/README.md b/README.md
@@ -112,7 +112,7 @@ dataset:
     NuScenes-v1.0-mini-scene-0103:
       frame:
         /CAM_FRONT/image_rect_compressed:
-          index_column: /CAM_FRONT/image_rect_compressed/idx
+          index_column: /CAM_FRONT/image_rect_compressed/_idx_
           reader:
             _target_: rbyte.io.frame.mcap.McapFrameReader
             path: data/NuScenes-v1.0-mini-scene-0103.mcap
@@ -121,7 +121,7 @@ dataset:
             frame_decoder: ${frame_decoder}
 
         /CAM_FRONT_LEFT/image_rect_compressed:
-          index_column: /CAM_FRONT_LEFT/image_rect_compressed/idx
+          index_column: /CAM_FRONT_LEFT/image_rect_compressed/_idx_
           reader:
             _target_: rbyte.io.frame.mcap.McapFrameReader
             path: data/NuScenes-v1.0-mini-scene-0103.mcap
@@ -130,7 +130,7 @@ dataset:
             frame_decoder: ${frame_decoder}
 
         /CAM_FRONT_RIGHT/image_rect_compressed:
-          index_column: /CAM_FRONT_RIGHT/image_rect_compressed/idx
+          index_column: /CAM_FRONT_RIGHT/image_rect_compressed/_idx_
           reader:
             _target_: rbyte.io.frame.mcap.McapFrameReader
             path: data/NuScenes-v1.0-mini-scene-0103.mcap
@@ -151,22 +151,22 @@ dataset:
               - rbyte.utils.mcap.McapJsonDecoderFactory
             fields:
               /CAM_FRONT/image_rect_compressed:
+                _idx_:
                 log_time:
                   _target_: polars.Datetime
                   time_unit: ns
-                idx: null
 
               /CAM_FRONT_LEFT/image_rect_compressed:
+                _idx_:
                 log_time:
                   _target_: polars.Datetime
                   time_unit: ns
-                idx: null
 
               /CAM_FRONT_RIGHT/image_rect_compressed:
+                _idx_:
                 log_time:
                   _target_: polars.Datetime
                   time_unit: ns
-                idx: null
 
               /odom:
                 log_time:
@@ -175,28 +175,28 @@ dataset:
                 vel.x: null
 
           merger:
-            _target_: rbyte.io.table.TableMerger
+            _target_: rbyte.io.table.TableAligner
             separator: /
             merge:
               /CAM_FRONT/image_rect_compressed:
                 log_time:
                   method: ref
 
               /CAM_FRONT_LEFT/image_rect_compressed:
-                log_time:
-                  method: ref
-                idx:
+                _idx_:
                   method: asof
                   tolerance: 10ms
                   strategy: nearest
-
-              /CAM_FRONT_RIGHT/image_rect_compressed:
                 log_time:
                   method: ref
-                idx:
+
+              /CAM_FRONT_RIGHT/image_rect_compressed:
+                _idx_:
                   method: asof
                   tolerance: 10ms
                   strategy: nearest
+                log_time:
+                  method: ref
 
               /odom:
                 log_time:
@@ -211,7 +211,7 @@ dataset:
 
   sample_builder:
     _target_: rbyte.sample.builder.GreedySampleTableBuilder
-    index_column: /CAM_FRONT/image_rect_compressed/idx
+    index_column: /CAM_FRONT/image_rect_compressed/_idx_
 
 frame_decoder:
   _target_: simplejpeg.decode_jpeg
@@ -252,11 +252,11 @@ Batch(
         is_shared=False),
     table=TensorDict(
         fields={
-            /CAM_FRONT/image_rect_compressed/idx: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+            /CAM_FRONT/image_rect_compressed/_idx_: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
             /CAM_FRONT/image_rect_compressed/log_time: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
-            /CAM_FRONT_LEFT/image_rect_compressed/idx: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+            /CAM_FRONT_LEFT/image_rect_compressed/_idx_: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
             /CAM_FRONT_LEFT/image_rect_compressed/log_time: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
-            /CAM_FRONT_RIGHT/image_rect_compressed/idx: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+            /CAM_FRONT_RIGHT/image_rect_compressed/_idx_: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
             /CAM_FRONT_RIGHT/image_rect_compressed/log_time: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
             /odom/vel.x: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.float64, is_shared=False)},
         batch_size=torch.Size([1]),
@@ -273,24 +273,24 @@ logger:
   _target_: rbyte.viz.loggers.RerunLogger
   schema:
     frame:
-      /CAM_FRONT/image_rect_compressed: 
-        rerun.components.ImageBufferBatch:
+      /CAM_FRONT/image_rect_compressed:
+        Image:
           color_model: RGB
 
       /CAM_FRONT_LEFT/image_rect_compressed:
-        rerun.components.ImageBufferBatch:
+        Image:
           color_model: RGB
 
       /CAM_FRONT_RIGHT/image_rect_compressed:
-        rerun.components.ImageBufferBatch:
+        Image:
           color_model: RGB
 
     table:
-      /CAM_FRONT/image_rect_compressed/log_time: rerun.TimeNanosColumn
-      /CAM_FRONT/image_rect_compressed/idx: rerun.TimeSequenceColumn
-      /CAM_FRONT_LEFT/image_rect_compressed/idx: rerun.TimeSequenceColumn
-      /CAM_FRONT_RIGHT/image_rect_compressed/idx: rerun.TimeSequenceColumn
-      /odom/vel.x: rerun.components.ScalarBatch
+      /CAM_FRONT/image_rect_compressed/log_time: TimeNanosColumn
+      /CAM_FRONT/image_rect_compressed/_idx_: TimeSequenceColumn
+      /CAM_FRONT_LEFT/image_rect_compressed/_idx_: TimeSequenceColumn
+      /CAM_FRONT_RIGHT/image_rect_compressed/_idx_: TimeSequenceColumn
+      /odom/vel.x: Scalar
 ```
 
 Visualize the dataset:

diff --git a/examples/config_templates/build_table.yaml b/examples/config_templates/build_table.yaml
@@ -1,13 +1,10 @@
 ---
 defaults:
   - table_builder: !!null
+  - table_writer: !!null
   - _self_
 
 path: ???
-writer:
-  _target_: polars.DataFrame.write_csv
-  _partial_: true
-  file: ???
 
 hydra:
   output_subdir: !!null

diff --git a/examples/config_templates/dataset/hdf5.yaml b/examples/config_templates/dataset/hdf5.yaml
@@ -0,0 +1,64 @@
+#! https://sites.google.com/view/il-for-mm/datasets#h.cq0r3rd5nr9m
+
+#@yaml/map-key-override
+#@yaml/text-templated-strings
+
+#@ inputs = {
+#@     "table_setup_from_dishwasher_sample": [
+#@         "/data/demo_0",
+#@         "/data/demo_1",
+#@         "/data/demo_10",
+#@         "/data/demo_101",
+#@         "/data/demo_102",
+#@     ]
+#@ }
+
+#@ frame_keys = [
+#@     '/obs/rgb',
+#@     '/obs/depth',
+#@ ]
+---
+_target_: rbyte.Dataset
+_convert_: all
+_recursive_: false
+inputs:
+  #@ for input_id, input_keys in inputs.items():
+  #@ for input_key in input_keys:
+  (@=input_id@)(@=input_key@):
+    frame:
+      #@ for frame_key in frame_keys:
+      (@=frame_key@):
+        index_column: _idx_
+        reader:
+          _target_: rbyte.io.frame.hdf5.Hdf5FrameReader
+          path: "${data_dir}/(@=input_id@).hdf5"
+          key: (@=input_key@)(@=frame_key@)
+      #@ 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:
+
+        merger:
+          _target_: rbyte.io.table.TableConcater
+          separator: "/"
+  #@ end
+  #@ end
+
+sample_builder:
+  _target_: rbyte.sample.builder.GreedySampleTableBuilder
+  index_column: _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
@@ -21,7 +21,7 @@ inputs:
     frame:
       #@ for topic in camera_topics:
       (@=topic@):
-        index_column: (@=topic@)/idx
+        index_column: (@=topic@)/_idx_
         reader:
           _target_: rbyte.io.frame.mcap.McapFrameReader
           path: "${data_dir}/(@=input_id@).mcap"
@@ -54,7 +54,7 @@ inputs:
                 _target_: polars.Datetime
                 time_unit: ns
 
-              idx:
+              _idx_:
             #@ end
 
             /odom:
@@ -64,8 +64,8 @@ inputs:
               vel.x:
 
         merger:
-          _target_: rbyte.io.table.TableMerger
-          separator: /
+          _target_: rbyte.io.table.TableAligner
+          separator: "/"
           merge:
             (@=camera_topics[0]@):
               log_time:
@@ -75,7 +75,7 @@ inputs:
             (@=topic@):
               log_time:
                 method: ref
-              idx:
+              _idx_:
                 method: asof
                 tolerance: 10ms
                 strategy: nearest
@@ -98,7 +98,7 @@ inputs:
 
 sample_builder:
   _target_: rbyte.sample.builder.GreedySampleTableBuilder
-  index_column: (@=camera_topics[0]@)/idx
+  index_column: (@=camera_topics[0]@)/_idx_
   length: 1
   stride: 1
   min_step: 1

diff --git a/examples/config_templates/dataset/yaak.yaml b/examples/config_templates/dataset/yaak.yaml
@@ -64,7 +64,7 @@ inputs:
                 categories: ["0", "1", "2", "3"]
 
         merger:
-          _target_: rbyte.io.table.TableMerger
+          _target_: rbyte.io.table.TableAligner
           separator: "."
           merge:
             ImageMetadata.(@=cameras[0]@):

diff --git a/examples/config_templates/frame_reader/hdf5.yaml b/examples/config_templates/frame_reader/hdf5.yaml
@@ -0,0 +1,5 @@
+---
+_target_: rbyte.io.frame.Hdf5FrameReader
+_recursive_: true
+path: ???
+key: ???
diff --git a/examples/config_templates/logger/rerun/carla.yaml b/examples/config_templates/logger/rerun/carla.yaml
@@ -12,14 +12,14 @@ schema:
   frame:
     #@ for camera in cameras:
     (@=camera@):
-      rerun.components.ImageBufferBatch:
+      Image:
         color_model: RGB
     #@ end
 
   table:
-    frame_idx: rerun.TimeSequenceColumn
-    control.brake: rerun.components.ScalarBatch
-    control.steer: rerun.components.ScalarBatch
-    control.throttle: rerun.components.ScalarBatch
-    state.acceleration.value: rerun.components.ScalarBatch
-    state.velocity.value: rerun.components.ScalarBatch
+    frame_idx: TimeSequenceColumn
+    control.brake: Scalar
+    control.steer: Scalar
+    control.throttle: Scalar
+    state.acceleration.value: Scalar
+    state.velocity.value: Scalar
diff --git a/examples/config_templates/logger/rerun/hdf5.yaml b/examples/config_templates/logger/rerun/hdf5.yaml
@@ -0,0 +1,23 @@
+#@yaml/text-templated-strings
+
+#@ camera_topics = [
+#@     '/CAM_FRONT/image_rect_compressed',
+#@     '/CAM_FRONT_LEFT/image_rect_compressed',
+#@     '/CAM_FRONT_RIGHT/image_rect_compressed',
+#@ ]
+---
+_target_: rbyte.viz.loggers.RerunLogger
+schema:
+  frame:
+    /obs/rgb:
+      Image:
+        color_model: RGB
+
+    /obs/depth:
+      DepthImage:
+        color_model: L
+
+  table:
+    _idx_: TimeSequenceColumn
+    obs/object: Points3D
+    task_successes: Scalar