diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 722d63f3f2a2..b1e487fff1f1 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 49e26097e0e9..90bb1abece2f 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -24,6 +24,8 @@ import operator import os import sys +from collections import OrderedDict +from functools import partial import pyarrow as pa @@ -42,6 +44,7 @@ _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 @@ def __init__( 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 @@ def __init__( 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 @@ def _init_metadata(self): 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 @@ def _init_reader(self, raw_table, info): 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 @@ def __getitems__(self, indices): 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 @@ def __getitems__(self, indices): ]) 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,173 @@ 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) + + at_frontier = self._next_index == len(self._timestamps) + if 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 _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 82e2b7631c1a..17bddc553dfa 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.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,127 @@ 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() + + @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 +1851,19 @@ def sample_timestamps(unused_dataset, uri): @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 +1892,17 @@ def test_imported_video_payload_can_be_decoded(self): "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 +1918,10 @@ def test_imported_video_payload_can_be_decoded(self): "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 +1935,10 @@ def test_imported_video_payload_can_be_decoded(self): "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 +1975,10 @@ def test_imported_video_payload_can_be_decoded(self): 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 +1996,8 @@ class Dataset: 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 +2072,64 @@ def close(self): [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)