This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new be77b25d40 [python] Read LeRobot video datasets for training (#9757)
be77b25d40 is described below
commit be77b25d4034ad21bb415e2b9a48d55bf4646fae
Author: XiaoHongbo <[email protected]>
AuthorDate: Sun Sep 13 16:28:48 2026 +0800
[python] Read LeRobot video datasets for training (#9757)
---
docs/docs/pypaimon/lerobot.md | 19 +-
.../pypaimon/multimodal/lerobot/dataset.py | 293 +++++++++++++++++++--
.../pypaimon/tests/multimodal_lerobot_test.py | 281 ++++++++++++++++++++
3 files changed, 564 insertions(+), 29 deletions(-)
diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md
index 722d63f3f2..b1e487fff1 100644
--- a/docs/docs/pypaimon/lerobot.md
+++ b/docs/docs/pypaimon/lerobot.md
@@ -109,9 +109,8 @@ Scalars map to scalar types, vectors to `VECTOR`,
higher-rank tensors to nested
Video features map to `BLOB`. Frame rows reference MP4 payloads copied once per
aligned file group. Video imports use the video grouping policy and check
-rolling before each Episode. They require a bucket-unaware table. Read them
-with a Paimon scan and `VideoFrameCollator`; `PaimonLeRobotDataset` currently
-supports image features only.
+rolling before each Episode. They require a bucket-unaware table. Use
+`VideoFrameCollator` for scans or `PaimonLeRobotDataset` for training.
## Capture LeRobot frames directly into Paimon
@@ -208,10 +207,9 @@ pin one named snapshot on every component.
## Train with Paimon LeRobot data
-For map-style training, read a tagged table group created by
-`load_from_lerobot` directly from Paimon. `PaimonLeRobotDataset` requires the
-complete table group; a frame-only table created by `PaimonLeRobotWriter` is
-not sufficient.
+For map-style training, pass an image- or video-backed table group created by
+`load_from_lerobot` to `PaimonLeRobotDataset`. A frame-only table created by
+`PaimonLeRobotWriter` is not sufficient.
```python
from torch.utils.data import DataLoader
@@ -224,6 +222,7 @@ dataset = PaimonLeRobotDataset(
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)
```
-If `tag_name` is omitted, the latest snapshots are used. Metadata is available
-through `dataset.meta`. Frame lookups use the BTree on `index`; payload columns
-remain lazy.
+Without `tag_name`, the latest snapshots are used. Frame lookups use the BTree
+on `index`; payloads remain lazy. Video decoding prefers TorchCodec, falls back
+to PyAV, and reuses a bounded decoder cache. Set `video_backend` to force
+either decoder.
diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
index 49e26097e0..7a13731e5f 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -24,6 +24,8 @@ import math
import operator
import os
import sys
+from collections import OrderedDict
+from functools import partial
import pyarrow as pa
@@ -42,6 +44,7 @@ from pypaimon.multimodal.lerobot.schema import (
_validate_lerobot_schema,
)
from pypaimon.multimodal.table import _target_schema, _time_travel_table
+from pypaimon.multimodal.video import VideoFrameCollator
from pypaimon.read.query_auth_split import QueryAuthSplit
@@ -78,7 +81,7 @@ class PaimonLeRobotDataset:
LeRobot metadata is resolved from the Paimon table group and remains
available through :attr:`meta`.
- Set ``return_uint8=True`` to keep 8-bit images in their decoded
+ Set ``return_uint8=True`` to keep 8-bit visual frames in their decoded
``torch.uint8`` representation instead of normalizing them to float32.
Higher-bit-depth images retain the existing float32 behavior.
"""
@@ -93,6 +96,7 @@ class PaimonLeRobotDataset:
delta_timestamps=None,
tolerance_s=1e-4,
blob_parallelism=16,
+ video_backend=None,
return_uint8=False):
if sys.version_info < (3, 10):
raise RuntimeError(
@@ -109,6 +113,10 @@ class PaimonLeRobotDataset:
raise ValueError("tolerance_s must be finite and non-negative.")
self.blob_parallelism = _positive_int(
blob_parallelism, "blob_parallelism")
+ if video_backend not in (None, "torchcodec", "pyav"):
+ raise ValueError(
+ "video_backend must be None, 'torchcodec', or 'pyav'.")
+ self.video_backend = video_backend
if not isinstance(return_uint8, bool):
raise TypeError("return_uint8 must be a boolean.")
self.return_uint8 = return_uint8
@@ -130,15 +138,11 @@ class PaimonLeRobotDataset:
name for name, feature in self._features.items()
if feature.get("dtype") == "image"
]
- video_keys = [
+ self._video_keys = [
name for name, feature in self._features.items()
if feature.get("dtype") == "video"
]
- if video_keys:
- raise NotImplementedError(
- "PaimonLeRobotDataset currently supports image-backed "
- "features only; video features are not yet supported: %s"
- % video_keys)
+ self._visual_keys = self._image_keys + self._video_keys
self._total_frames = int(
_metadata_member(
@@ -230,6 +234,18 @@ class PaimonLeRobotDataset:
self._read_table, snapshot, splits)
self._validation_context = validation_context
self._file_io = self._read_table.file_io
+ self._video_collators = [
+ VideoFrameCollator(
+ self._read_table,
+ video_column=key,
+ decoder_factory=partial(
+ _open_video_decoder, backend=self.video_backend),
+ decode_fn=_decode_video_frame,
+ output_column=key,
+ collate_fn=_identity,
+ )
+ for key in self._video_keys
+ ]
self._task_names = validation_context["task_names"]
self._subtask_names = validation_context["subtask_names"]
self._delta_projection = None
@@ -319,21 +335,30 @@ class PaimonLeRobotDataset:
self._image_keys,
self.blob_parallelism,
)
- converted = {
- position: _torch_row(
- row, self._features, self.return_uint8)
- for position, row in base_rows.items()
- }
- converted.update({
- position: _torch_row(
- row, self._features, self.return_uint8)
- for position, row in delta_rows.items()
- })
+ _decode_image_rows(
+ row_groups,
+ self._image_keys,
+ self._features,
+ self.return_uint8,
+ )
break
except OSError:
if attempt + 1 == _IMAGE_READ_ATTEMPTS:
raise
+ _decode_video_rows(
+ row_groups, getattr(self, "_video_collators", ()))
+ converted = {
+ position: _torch_row(
+ row, self._features, self.return_uint8)
+ for position, row in base_rows.items()
+ }
+ converted.update({
+ position: _torch_row(
+ row, self._features, self.return_uint8)
+ for position, row in delta_rows.items()
+ })
+
import torch
duplicates = _duplicate_indices(plans)
result = []
@@ -350,11 +375,34 @@ class PaimonLeRobotDataset:
])
item.update(plan["padding"])
if self.image_transforms is not None:
- for key in self._image_keys:
+ for key in self._visual_keys:
item[key] = self.image_transforms(item[key])
result.append(item)
return result
+ def close(self):
+ first_error = None
+ locator = getattr(self, "_frame_locator", None)
+ if locator is not None:
+ try:
+ locator.close()
+ except Exception as error:
+ first_error = error
+ for collator in getattr(self, "_video_collators", ()):
+ try:
+ collator.close()
+ except Exception as error:
+ if first_error is None:
+ first_error = error
+ if first_error is not None:
+ raise first_error
+
+ def __del__(self):
+ try:
+ self.close()
+ except Exception:
+ pass
+
def _read_rows(
self, indices, projection, splits=None, needs_filter=True):
if not indices:
@@ -1086,6 +1134,15 @@ def _resolve_image_blobs(
row[key] = body
+def _decode_image_rows(row_groups, image_keys, features, return_uint8):
+ for rows in row_groups:
+ for row in rows.values():
+ for key in image_keys:
+ if key in row:
+ row[key] = _image_tensor(
+ row[key], features[key], return_uint8=return_uint8)
+
+
def _image_blob_sources(row_groups, image_keys):
return [
(row, key, row[key])
@@ -1118,9 +1175,12 @@ def _torch_row(row, features, return_uint8=False):
if key not in result:
continue
value = result[key]
- if feature.get("dtype") == "image":
+ if feature.get("dtype") == "image" and not torch.is_tensor(value):
result[key] = _image_tensor(
value, feature, return_uint8=return_uint8)
+ elif feature.get("dtype") == "video":
+ result[key] = _video_tensor(
+ value, feature, return_uint8=return_uint8)
elif feature.get("dtype") != "string" and not torch.is_tensor(value):
dtype = getattr(torch, _TORCH_DTYPE_NAMES[feature.get("dtype")])
result[key] = torch.tensor(value, dtype=dtype)
@@ -1164,6 +1224,201 @@ def _image_tensor(payload, feature, return_uint8=False):
return tensor.div_(255) if normalize else tensor
+def _video_tensor(frame, feature, return_uint8=False):
+ import torch
+
+ if not torch.is_tensor(frame):
+ raise ValueError("LeRobot video decoder must return a Torch tensor.")
+ expected_shape = _feature_shape(feature, "video")
+ if len(expected_shape) != 3:
+ raise ValueError("LeRobot video feature must have three dimensions.")
+ names = feature.get("names") or []
+ output_shape = expected_shape if names and names[0] in (
+ "channel", "channels"
+ ) else expected_shape[2:] + expected_shape[:2]
+ if tuple(frame.shape) != output_shape:
+ raise ValueError(
+ "LeRobot video frame has shape %s, expected %s."
+ % (tuple(frame.shape), output_shape)
+ )
+ if frame.dtype == torch.uint8 and not return_uint8:
+ return frame.float().div_(255)
+ return frame
+
+
+def _open_video_decoder(stream, backend=None):
+ if backend in (None, "torchcodec"):
+ try:
+ return _open_torchcodec_decoder(stream)
+ except (ImportError, OSError, RuntimeError):
+ if backend == "torchcodec":
+ raise
+ stream.seek(0)
+ return _PyAVVideoDecoder(stream)
+
+
+def _open_torchcodec_decoder(stream):
+ try:
+ from torchcodec.decoders import VideoDecoder
+ except (ImportError, RuntimeError) as error:
+ raise ImportError(
+ "Video-backed PaimonLeRobotDataset requires TorchCodec from "
+ "'pypaimon[lerobot]'."
+ ) from error
+ try:
+ return VideoDecoder(stream, seek_mode="exact")
+ except TypeError:
+ # TorchCodec 0.2 accepts bytes but not seekable file-like objects.
+ stream.seek(0)
+ return VideoDecoder(stream.read(), seek_mode="exact")
+
+
+class _PyAVVideoDecoder:
+
+ # Reuse common overlapping delta windows without retaining a whole video.
+ _FRAME_CACHE_SIZE = 8
+
+ def __init__(self, stream):
+ try:
+ import av
+ except ImportError as error:
+ raise ImportError(
+ "Video-backed PaimonLeRobotDataset requires PyAV from "
+ "'pypaimon[lerobot]'."
+ ) from error
+ self._container = av.open(stream)
+ self._stream = self._container.streams.video[0]
+ self._next_index = 0
+ self._timestamps = []
+ self._keyframes = []
+ self._cache = OrderedDict()
+ self._frames = iter(self._container.decode(self._stream))
+
+ def __getitem__(self, index):
+ index = operator.index(index)
+ if index < 0:
+ raise IndexError("Video frame index %d is out of range." % index)
+ frame = self._cache.pop(index, None)
+ if frame is not None:
+ self._cache[index] = frame
+ return self._tensor(frame)
+
+ indexed = (
+ index > 0 and not self._timestamps
+ and self._index_packets()
+ )
+ at_frontier = self._next_index == len(self._timestamps)
+ if indexed:
+ self._seek(index)
+ elif index != self._next_index and not (
+ at_frontier and index >= self._next_index):
+ self._seek(index)
+ try:
+ while True:
+ frame = next(self._frames)
+ if frame.pts is None:
+ continue
+ timestamp = frame.pts * (
+ frame.time_base or self._stream.time_base)
+ position = bisect.bisect_left(self._timestamps, timestamp)
+ if (
+ position < len(self._timestamps)
+ and self._timestamps[position] == timestamp
+ ):
+ frame_index = position
+ elif self._next_index == len(self._timestamps):
+ frame_index = self._next_index
+ self._timestamps.append(timestamp)
+ if frame.key_frame:
+ self._keyframes.append(frame_index)
+ else:
+ continue
+ self._next_index = frame_index + 1
+ self._remember(frame_index, frame)
+ if frame_index == index:
+ return self._tensor(frame)
+ if frame_index > index:
+ break
+ except StopIteration as error:
+ raise IndexError(
+ "Video frame index %d is out of range." % index
+ ) from error
+ raise IndexError("Video frame index %d is out of range." % index)
+
+ def _index_packets(self):
+ entries = []
+ for packet in self._container.demux(self._stream):
+ if (packet.pts is None
+ or getattr(packet, "is_discard", False)):
+ continue
+ timestamp = packet.pts * (
+ packet.time_base or self._stream.time_base)
+ entries.append((timestamp, packet.is_keyframe))
+ if not entries:
+ self._container.seek(
+ 0, backward=True, any_frame=False, stream=self._stream)
+ self._frames = iter(self._container.decode(self._stream))
+ return False
+ entries.sort(key=lambda entry: entry[0])
+ self._timestamps = [timestamp for timestamp, unused in entries]
+ self._keyframes = [
+ index for index, (unused, keyframe) in enumerate(entries)
+ if keyframe
+ ]
+ return True
+
+ def _seek(self, index):
+ position = bisect.bisect_right(self._keyframes, index)
+ anchor = self._keyframes[position - 1] if position else 0
+ timestamp = self._timestamps[anchor]
+ self._container.seek(
+ round(timestamp / self._stream.time_base),
+ backward=True,
+ any_frame=False,
+ stream=self._stream,
+ )
+ self._next_index = None
+ self._frames = iter(self._container.decode(self._stream))
+
+ def _remember(self, index, frame):
+ self._cache.pop(index, None)
+ self._cache[index] = frame
+ if len(self._cache) > self._FRAME_CACHE_SIZE:
+ self._cache.popitem(last=False)
+
+ @staticmethod
+ def _tensor(frame):
+ import numpy as np
+ import torch
+ array = np.array(frame.to_ndarray(format="rgb24"), copy=True)
+ return torch.from_numpy(array).permute(2, 0, 1)
+
+ def close(self):
+ self._container.close()
+
+
+def _decode_video_frame(decoder, frame_index, unused_row):
+ return decoder[frame_index]
+
+
+def _identity(values):
+ return values
+
+
+def _decode_video_rows(row_groups, collators):
+ for collator in collators:
+ for rows in row_groups:
+ indices = [
+ index for index, row in rows.items()
+ if collator.video_column in row
+ ]
+ if not indices:
+ continue
+ decoded = collator([rows[index] for index in indices])
+ for index, row in zip(indices, decoded):
+ rows[index] = row
+
+
def _normalize_index(index, size):
index = operator.index(index)
if index < 0:
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
index 82e2b7631c..97b180d39e 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -17,6 +17,7 @@
import builtins
from array import array
from fractions import Fraction
+import importlib.util
import io
import json
import pickle
@@ -44,9 +45,11 @@ from pypaimon.multimodal.source_utils import _SourceFileIO
from pypaimon.multimodal.connection import MultimodalConnection
from pypaimon.multimodal.lerobot import load_from_lerobot
from pypaimon.multimodal.lerobot.dataset import (
+ _PyAVVideoDecoder,
_arrow_rows,
_image_tensor,
_index_names,
+ _open_video_decoder,
_selected_episodes,
_torch_row,
)
@@ -126,6 +129,193 @@ def _catalog_metadata(connection, name):
class LeRobotValidationTest(unittest.TestCase):
+ @unittest.skipUnless(
+ av is not None and importlib.util.find_spec("torch") is not None,
+ "PyAV and Torch are required for video decoding",
+ )
+ def test_pyav_decoder_reuses_windows_and_seeks_known_frames(self):
+ class Frame:
+
+ time_base = Fraction(1, 10)
+
+ def __init__(self, pts):
+ self.pts = pts
+ self.key_frame = pts % 10 == 0
+
+ def to_ndarray(self, format):
+ assert format == "rgb24"
+ return np.full((2, 2, 3), self.pts, dtype=np.uint8)
+
+ class Container:
+
+ def __init__(self):
+ self.stream = SimpleNamespace(time_base=Fraction(1, 10))
+ self.streams = SimpleNamespace(video=[self.stream])
+ self.position = 0
+ self.decoded = 0
+ self.seeks = []
+
+ def decode(self, stream):
+ assert stream is self.stream
+ while self.position < 120:
+ index = self.position
+ self.position += 1
+ self.decoded += 1
+ yield Frame(index)
+
+ def seek(self, offset, *, backward, any_frame, stream):
+ assert backward
+ assert not any_frame
+ assert stream is self.stream
+ self.seeks.append(offset)
+ self.position = offset
+
+ def close(self):
+ pass
+
+ container = Container()
+ with patch("av.open", return_value=container):
+ decoder = _PyAVVideoDecoder(io.BytesIO())
+ try:
+ for index in range(119):
+ decoder[index]
+ decoder[index + 1]
+ self.assertEqual(120, container.decoded)
+ self.assertEqual([], container.seeks)
+
+ decoder[5]
+ decoder[90]
+ decoder[8]
+ self.assertEqual(136, container.decoded)
+ self.assertEqual([0, 90, 0], container.seeks)
+ finally:
+ decoder.close()
+
+ def test_pyav_decoder_indexes_cold_random_reads(self):
+ class Frame:
+
+ time_base = Fraction(1, 10)
+
+ def __init__(self, pts):
+ self.pts = pts
+ self.key_frame = pts % 10 == 0
+
+ def to_ndarray(self, format):
+ assert format == "rgb24"
+ return np.full((2, 2, 3), self.pts, dtype=np.uint8)
+
+ class Container:
+
+ def __init__(self):
+ self.stream = SimpleNamespace(time_base=Fraction(1, 10))
+ self.streams = SimpleNamespace(video=[self.stream])
+ self.position = 0
+ self.decoded = 0
+ self.demuxed = 0
+ self.seeks = []
+
+ def demux(self, stream):
+ assert stream is self.stream
+ for index in range(120):
+ self.demuxed += 1
+ yield SimpleNamespace(
+ pts=index, time_base=Fraction(1, 10),
+ is_discard=False, is_keyframe=index % 10 == 0,
+ )
+
+ def decode(self, stream):
+ assert stream is self.stream
+ while self.position < 120:
+ index = self.position
+ self.position += 1
+ self.decoded += 1
+ yield Frame(index)
+
+ def seek(self, offset, *, backward, any_frame, stream):
+ assert backward
+ assert not any_frame
+ assert stream is self.stream
+ self.seeks.append(offset)
+ self.position = offset
+
+ def close(self):
+ pass
+
+ container = Container()
+ fake_av = SimpleNamespace(open=lambda unused_stream: container)
+ tensor = staticmethod(
+ lambda frame: frame.to_ndarray(format="rgb24"))
+ with patch.dict(sys.modules, {"av": fake_av}), patch.object(
+ _PyAVVideoDecoder, "_tensor", tensor):
+ decoder = _PyAVVideoDecoder(io.BytesIO())
+ try:
+ frame = decoder[95]
+ self.assertEqual(120, container.demuxed)
+ self.assertEqual([90], container.seeks)
+ self.assertEqual(6, container.decoded)
+ self.assertTrue((frame == 95).all())
+ finally:
+ decoder.close()
+
+ @unittest.skipUnless(
+ av is not None and importlib.util.find_spec("torch") is not None,
+ "PyAV and Torch are required for video decoding",
+ )
+ def test_pyav_decoder_seeks_before_b_frames(self):
+ output = io.BytesIO()
+ with av.open(output, mode="w", format="mp4") as container:
+ stream = container.add_stream("mpeg4", rate=30)
+ stream.width = 16
+ stream.height = 16
+ stream.pix_fmt = "yuv420p"
+ stream.gop_size = 12
+ stream.codec_context.max_b_frames = 2
+ for index in range(70):
+ image = np.full(
+ (16, 16, 3), index + 24, dtype=np.uint8)
+ frame = av.VideoFrame.from_ndarray(image, format="rgb24")
+ frame.pts = index
+ frame.time_base = Fraction(1, 30)
+ for packet in stream.encode(frame):
+ container.mux(packet)
+ for packet in stream.encode():
+ container.mux(packet)
+
+ payload = output.getvalue()
+ with av.open(io.BytesIO(payload)) as container:
+ expected = [
+ np.array(frame.to_ndarray(format="rgb24"), copy=True)
+ for frame in container.decode(video=0)
+ ]
+
+ decoder = _PyAVVideoDecoder(io.BytesIO(payload))
+ try:
+ for index in (69, 20, 35, 1, 68):
+ actual = decoder[index].permute(1, 2, 0).numpy()
+ np.testing.assert_array_equal(expected[index], actual)
+ finally:
+ decoder.close()
+
+ def test_default_video_backend_falls_back_on_os_error(self):
+ stream = Mock()
+ decoder = object()
+ with patch(
+ "pypaimon.multimodal.lerobot.dataset."
+ "_open_torchcodec_decoder",
+ side_effect=OSError("unavailable")), patch(
+ "pypaimon.multimodal.lerobot.dataset._PyAVVideoDecoder",
+ return_value=decoder) as pyav:
+ self.assertIs(decoder, _open_video_decoder(stream))
+ stream.seek.assert_called_once_with(0)
+ pyav.assert_called_once_with(stream)
+
+ with patch(
+ "pypaimon.multimodal.lerobot.dataset."
+ "_open_torchcodec_decoder",
+ side_effect=OSError("unavailable")):
+ with self.assertRaises(OSError):
+ _open_video_decoder(stream, backend="torchcodec")
+
def test_dataset_requires_supported_python(self):
with patch(
"pypaimon.multimodal.lerobot.dataset.sys.version_info",
@@ -1727,6 +1917,19 @@ class LeRobotValidationTest(unittest.TestCase):
@unittest.skipUnless(av is not None, "PyAV is required for MP4 decoding")
def test_imported_video_payload_can_be_decoded(self):
+ self._assert_imported_video_payload_can_be_decoded(False)
+
+ @unittest.skipUnless(
+ av is not None
+ and sys.version_info >= (3, 10)
+ and importlib.util.find_spec("datasets") is not None
+ and importlib.util.find_spec("torch") is not None,
+ "Video training reads require Python 3.10+, PyAV, datasets, and Torch",
+ )
+ def test_imported_video_payload_supports_training_reads(self):
+ self._assert_imported_video_payload_can_be_decoded(True)
+
+ def _assert_imported_video_payload_can_be_decoded(self, training_reads):
import pandas as pd
temp_dir = Path(tempfile.mkdtemp(prefix="pypaimon_lerobot_mp4_"))
@@ -1755,11 +1958,17 @@ class LeRobotValidationTest(unittest.TestCase):
"fps": 10.0,
},
"task_index": {"dtype": "int64", "shape": [1]},
+ "action": {"dtype": "float32", "shape": [1]},
"camera": {
"dtype": "video",
"shape": [16, 16, 3],
"video_info": {"video.fps": 10.0},
},
+ "camera_b": {
+ "dtype": "video",
+ "shape": [16, 16, 3],
+ "video_info": {"video.fps": 10.0},
+ },
},
}
episodes = [
@@ -1775,6 +1984,10 @@ class LeRobotValidationTest(unittest.TestCase):
"videos/camera/file_index": 0,
"videos/camera/from_timestamp": 0.5,
"videos/camera/to_timestamp": 0.7,
+ "videos/camera_b/chunk_index": 0,
+ "videos/camera_b/file_index": 0,
+ "videos/camera_b/from_timestamp": 0.5,
+ "videos/camera_b/to_timestamp": 0.7,
},
{
"episode_index": 1,
@@ -1788,6 +2001,10 @@ class LeRobotValidationTest(unittest.TestCase):
"videos/camera/file_index": 0,
"videos/camera/from_timestamp": 0.1,
"videos/camera/to_timestamp": 0.4,
+ "videos/camera_b/chunk_index": 0,
+ "videos/camera_b/file_index": 0,
+ "videos/camera_b/from_timestamp": 0.1,
+ "videos/camera_b/to_timestamp": 0.4,
},
]
physical_frame_values = [24, 56, 88, 120, 168, 216]
@@ -1824,6 +2041,10 @@ class LeRobotValidationTest(unittest.TestCase):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
+ camera_b_path = (
+ temp_dir / "videos/camera_b/chunk-000/file-000.mp4")
+ camera_b_path.parent.mkdir(parents=True)
+ shutil.copy2(video_path, camera_b_path)
class Dataset:
@@ -1841,6 +2062,8 @@ class LeRobotValidationTest(unittest.TestCase):
type=pa.float32(),
),
"task_index": pa.array([0] * 5, type=pa.int64()),
+ "action": pa.array(
+ [0.0, 1.0, 2.0, 3.0, 4.0], type=pa.float32()),
})
def __len__(self):
@@ -1915,6 +2138,64 @@ class LeRobotValidationTest(unittest.TestCase):
[0.5, 0.6, 0.1, 0.2, 0.3],
atol=1e-6,
)
+ if not training_reads:
+ return
+
+ dataset = pmm.PaimonLeRobotDataset(
+ table,
+ delta_timestamps={"camera": [0.0, 0.1]},
+ )
+ try:
+ last, first = dataset.__getitems__([4, 0])
+ self.assertEqual(
+ [2, 3, 16, 16], list(last["camera"].shape))
+ self.assertEqual(
+ [2, 3, 16, 16], list(first["camera"].shape))
+ self.assertEqual(
+ [3, 16, 16], list(first["camera_b"].shape))
+ self.assertEqual("torch.float32", str(last["camera"].dtype))
+ np.testing.assert_allclose(
+ [
+ float(last["camera"][0].mean()) * 255,
+ float(first["camera"][0].mean()) * 255,
+ float(first["camera"][1].mean()) * 255,
+ ],
+ [120, 168, 216],
+ atol=5,
+ )
+ self.assertEqual(
+ [False, True], last["camera_is_pad"].tolist())
+ self.assertEqual(
+ 1, len(dataset._video_collators[0]._decoders))
+
+ from torch.utils.data import DataLoader
+ worker_indices = []
+ for batch in DataLoader(
+ dataset,
+ batch_size=2,
+ shuffle=False,
+ num_workers=2,
+ multiprocessing_context="spawn"):
+ worker_indices.extend(batch["index"].tolist())
+ self.assertEqual(list(range(5)), worker_indices)
+ finally:
+ dataset.close()
+
+ action_dataset = pmm.PaimonLeRobotDataset(
+ table,
+ delta_timestamps={"action": [0.0, 0.1]},
+ )
+ try:
+ item = action_dataset[0]
+ self.assertEqual([2], list(item["action"].shape))
+ np.testing.assert_allclose(
+ [0.0, 1.0], item["action"].tolist())
+ self.assertEqual(
+ [3, 16, 16], list(item["camera"].shape))
+ self.assertEqual(
+ [3, 16, 16], list(item["camera_b"].shape))
+ finally:
+ action_dataset.close()
finally:
shutil.rmtree(temp_dir, ignore_errors=True)
