feat: video/path dataframe cross-reference (#33)

- add VideoDataFrameBuilder
- update PathDataFrameBuilder
- update yaak dataset template

Author: egorchakov

.pre-commit-config.yaml

@@ -6,19 +6,19 @@ repos:
       - id: validate-pyproject
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.19.0
+    rev: v3.19.1
     hooks:
       - id: pyupgrade
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.8.1
+    rev: v0.9.2
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
 
   - repo: https://github.com/DetachHead/basedpyright-pre-commit-mirror
-    rev: 1.22.0
+    rev: 1.24.0
     hooks:
       - id: basedpyright
 

config/_templates/dataset/yaak.yaml

@@ -35,17 +35,14 @@ inputs:
           cache_dir: /tmp/rbyte-cache
         functions:
           - _target_: pipefunc.PipeFunc
-            scope: yaak_metadata
-            output_name: output
-            cache: true
             func:
               _target_: hydra.utils.get_method
               path: rbyte.io.build_yaak_metadata_dataframe
+            output_name: output
+            scope: metadata
+            cache: true
 
           - _target_: pipefunc.PipeFunc
-            bound:
-              path: ${data_dir}/(@=input_id@)/ai.mcap
-            output_name: mcap
             func:
               _target_: rbyte.io.McapDataFrameBuilder
               decoder_factories: [rbyte.utils._mcap.ProtobufDecoderFactory]
@@ -58,18 +55,21 @@ inputs:
                   score:
                     _target_: polars.Float32
 
+            output_name: output
+            bound:
+              path: ${data_dir}/(@=input_id@)/ai.mcap
+            scope: mcap
+
           - _target_: pipefunc.PipeFunc
             func:
               _target_: pipefunc.helpers.collect_kwargs
               parameters: [meta, mcap]
-            renames:
-              meta: yaak_metadata.output
             output_name: data
+            renames:
+              meta: metadata.output
+              mcap: mcap.output
 
           - _target_: pipefunc.PipeFunc
-            renames:
-              input: data
-            output_name: data_aligned
             func:
               _target_: rbyte.io.DataFrameAligner
               separator: /
@@ -111,18 +111,56 @@ inputs:
                         tolerance: 500ms
                         strategy: nearest
 
-          - _target_: pipefunc.PipeFunc
+            output_name: data_aligned
             renames:
-              input: data_aligned
-            output_name: data_filtered
+              input: data
+
+          - _target_: pipefunc.PipeFunc
             func:
               _target_: rbyte.io.DataFrameFilter
               predicate: |
                 `meta/VehicleMotion/speed` > 44
+            output_name: data_filtered
+            renames:
+              input: data_aligned
+
+          #@ for i, camera in enumerate(cameras):
+          - _target_: pipefunc.PipeFunc
+            func:
+              _target_: rbyte.io.VideoDataFrameBuilder
+              fields:
+                frame_idx:
+                  _target_: polars.Int32
+
+            output_name: data_(@=camera@)
+            bound:
+              path: "${data_dir}/(@=input_id@)/(@=camera@).pii.mp4"
+
+          - _target_: pipefunc.PipeFunc
+            func:
+              _target_: hydra.utils.get_method
+              path: polars.DataFrame.join
+            #@ if i == len(cameras) - 1:
+            output_name: data_joined
+            #@ else:
+            output_name: data_joined_(@=camera@)
+            #@ end
+            renames:
+              #@ if i == 0:
+              self: data_filtered
+              #@ else:
+              self: data_joined_(@=cameras[i-1]@)
+              #@ end
+              other: data_(@=camera@)
+            bound:
+              how: semi
+              left_on: meta/ImageMetadata.(@=camera@)/frame_idx
+              right_on: frame_idx
+          #@ end
 
           - _target_: pipefunc.PipeFunc
             renames:
-              input: data_filtered
+              input: data_joined
             output_name: samples
             func:
               _target_: rbyte.FixedWindowSampleBuilder
@@ -132,7 +170,7 @@ inputs:
               length: 6
 
       kwargs:
-        yaak_metadata:
+        metadata:
           path: ${data_dir}/(@=input_id@)/metadata.log
           fields:
             rbyte.io.yaak.proto.sensor_pb2.ImageMetadata:

justfile

@@ -62,24 +62,24 @@ visualize *ARGS: generate-config
         {{ ARGS }}
 
 [group('visualize')]
-visualize-mimicgen:
-    just visualize dataset=mimicgen logger=rerun/mimicgen ++data_dir={{ justfile_directory() }}/tests/data/mimicgen
+visualize-mimicgen *ARGS:
+    just visualize dataset=mimicgen logger=rerun/mimicgen ++data_dir={{ justfile_directory() }}/tests/data/mimicgen {{ ARGS }}
 
 [group('visualize')]
-visualize-yaak:
-    just visualize dataset=yaak logger=rerun/yaak ++data_dir={{ justfile_directory() }}/tests/data/yaak
+visualize-yaak *ARGS:
+    just visualize dataset=yaak logger=rerun/yaak ++data_dir={{ justfile_directory() }}/tests/data/yaak {{ ARGS }}
 
 [group('visualize')]
-visualize-zod:
-    just visualize dataset=zod logger=rerun/zod ++data_dir={{ justfile_directory() }}/tests/data/zod
+visualize-zod *ARGS:
+    just visualize dataset=zod logger=rerun/zod ++data_dir={{ justfile_directory() }}/tests/data/zod {{ ARGS }}
 
 [group('visualize')]
-visualize-nuscenes-mcap:
-    just visualize dataset=nuscenes/mcap logger=rerun/nuscenes/mcap ++data_dir={{ justfile_directory() }}/tests/data/nuscenes/mcap
+visualize-nuscenes-mcap *ARGS:
+    just visualize dataset=nuscenes/mcap logger=rerun/nuscenes/mcap ++data_dir={{ justfile_directory() }}/tests/data/nuscenes/mcap {{ ARGS }}
 
 [group('visualize')]
-visualize-nuscenes-rrd:
-    just visualize dataset=nuscenes/rrd logger=rerun/nuscenes/rrd ++data_dir={{ justfile_directory() }}/tests/data/nuscenes/rrd
+visualize-nuscenes-rrd *ARGS:
+    just visualize dataset=nuscenes/rrd logger=rerun/nuscenes/rrd ++data_dir={{ justfile_directory() }}/tests/data/nuscenes/rrd {{ ARGS }}
 
 # rerun server and viewer
 rerun bind="0.0.0.0" port="9876" ws-server-port="9877" web-viewer-port="9090":

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rbyte"
-version = "0.10.2"
+version = "0.11.0"
 description = "Multimodal PyTorch dataset library"
 authors = [{ name = "Evgenii Gorchakov", email = "[email protected]" }]
 maintainers = [{ name = "Evgenii Gorchakov", email = "[email protected]" }]
@@ -18,7 +18,7 @@ dependencies = [
   "parse>=1.20.2",
   "structlog>=24.4.0",
   "tqdm>=4.66.5",
-  "pipefunc>=0.46.0",
+  "pipefunc>=0.50.0",
 ]
 readme = "README.md"
 requires-python = ">=3.12,<3.13"
@@ -49,7 +49,7 @@ yaak = ["protobuf", "ptars>=0.0.3"]
 jpeg = ["simplejpeg>=1.7.6"]
 video = [
   "python-vali>=4.2.0.post0; sys_platform == 'linux'",
-  "video-reader-rs>=0.2.1",
+  "video-reader-rs>=0.2.2",
 ]
 hdf5 = ["h5py>=3.12.1"]
 rrd = ["rerun-sdk>=0.21.0", "pyarrow-stubs"]

src/rbyte/io/__init__.py

@@ -49,6 +49,13 @@
 else:
     __all__ += ["FfmpegFrameSource"]
 
+try:
+    from .video.dataframe_builder import VideoDataFrameBuilder
+except ImportError:
+    pass
+else:
+    __all__ += ["VideoDataFrameBuilder"]
+
 try:
     from .yaak import YaakMetadataDataFrameBuilder, build_yaak_metadata_dataframe
 except ImportError:

src/rbyte/io/path/dataframe_builder.py

@@ -42,10 +42,11 @@ def __call__(self, path: PathLike[str]) -> pl.DataFrame:
                 )
                 raise RuntimeError(msg)
 
-        parent = Path(os.path.commonpath([path, parser._expression]))  # pyright: ignore[reportPrivateUsage]  # noqa: SLF001
-        results = (parser.parse(p.as_posix()) for p in parent.rglob("*") if p.is_file())  # pyright: ignore[reportUnknownMemberType]
-
-        return pl.DataFrame(
-            data=(r.named for r in results if isinstance(r, parse.Result)),  # pyright: ignore[reportUnknownMemberType]
-            schema=self._fields,
+        parent = Path(
+            os.path.commonpath([path, parser._expression.replace("\\.", ".")])  # pyright: ignore[reportPrivateUsage]  # noqa: SLF001
         )
+        paths = map(Path.as_posix, parent.rglob("*"))
+        parsed = map(parser.parse, paths)  # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType]
+        data = (x.named for x in parsed if isinstance(x, parse.Result))  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
+
+        return pl.DataFrame(data=data, schema=self._fields)

src/rbyte/io/video/dataframe_builder.py

@@ -0,0 +1,33 @@
+from collections.abc import Mapping
+from os import PathLike
+from pathlib import Path
+from typing import Literal, final
+
+import polars as pl
+from polars._typing import PolarsIntegerType  # noqa: PLC2701
+from polars.datatypes import (
+    IntegerType,  # pyright: ignore[reportUnusedImport] # noqa: F401
+)
+from pydantic import ConfigDict, validate_call
+from video_reader import (
+    PyVideoReader,  # pyright: ignore[reportAttributeAccessIssue, reportUnknownVariableType]
+)
+
+type Fields = Mapping[Literal["frame_idx"], PolarsIntegerType]
+
+
+@final
+class VideoDataFrameBuilder:
+    __name__ = __qualname__
+
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+    def __init__(self, fields: Fields) -> None:
+        self._fields = fields
+
+    def __call__(self, path: PathLike[str]) -> pl.DataFrame:
+        vr = PyVideoReader(Path(path).resolve().as_posix())  # pyright: ignore[reportUnknownVariableType]
+        info = vr.get_info()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+        frame_count = int(info["frame_count"])  # pyright: ignore[reportUnknownArgumentType]
+        data = pl.arange(frame_count, eager=True)
+
+        return pl.DataFrame(data=data, schema=self._fields)  # pyright: ignore[reportArgumentType]

src/rbyte/io/yaak/idl-repo

@@ -14,11 +14,11 @@ def pad_dim(
     mode: str = "constant",
     value: float | None = None,
 ) -> Tensor:
-    _pad = [(0, 0) for _ in input.shape]
-    _pad[dim] = pad
-    _pad = list(mit.flatten(reversed(_pad)))
+    pad_ = [(0, 0) for _ in input.shape]
+    pad_[dim] = pad
+    pad_ = list(mit.flatten(reversed(pad_)))
 
-    return F.pad(input, _pad, mode=mode, value=value)
+    return F.pad(input, pad_, mode=mode, value=value)
 
 
 def pad_sequence(sequences: Sequence[Tensor], dim: int, value: float = 0.0) -> Tensor: