You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@paimon.apache.org by ho...@apache.org on 2026/09/20 04:02:05 UTC
(paimon) 16/20: [python] Batch TorchCodec frame decoding for LeRobot datasets (#9836)
This is an automated email from the ASF dual-hosted git repository.
XiaoHongbo-Hope pushed a commit to branch release-2.1
in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 1fe52c90377b5ee4e8991f06f12ae532b1672390
Author: QuakeWang <45...@users.noreply.github.com>
AuthorDate: Thu Sep 17 14:41:40 2026 +0800
[python] Batch TorchCodec frame decoding for LeRobot datasets (#9836)
---
.../pypaimon/multimodal/lerobot/dataset.py | 9 +-
paimon-python/pypaimon/multimodal/video.py | 40 ++++-
.../pypaimon/tests/multimodal_lerobot_test.py | 163 +++++++++++++++++++++
.../pypaimon/tests/multimodal_video_test.py | 132 ++++++++++++++++-
4 files changed, 332 insertions(+), 12 deletions(-)
diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
index 006ca6357b..062192150a 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -268,7 +268,7 @@ class PaimonDatasetReader(ABC):
video_column=key,
decoder_factory=partial(
_open_video_decoder, backend=self.video_backend),
- decode_fn=_decode_video_frame,
+ decode_batch_fn=_decode_video_frames,
output_column=key,
collate_fn=_identity,
)
@@ -1631,8 +1631,11 @@ class _PyAVVideoDecoder:
self._container.close()
-def _decode_video_frame(decoder, frame_index, unused_row):
- return decoder[frame_index]
+def _decode_video_frames(decoder, frame_indices, unused_rows):
+ get_frames_at = getattr(decoder, "get_frames_at", None)
+ if get_frames_at is not None:
+ return get_frames_at(indices=frame_indices).data
+ return [decoder[index] for index in frame_indices]
def _identity(values):
diff --git a/paimon-python/pypaimon/multimodal/video.py b/paimon-python/pypaimon/multimodal/video.py
index c027a91e41..ac2e7039e8 100644
--- a/paimon-python/pypaimon/multimodal/video.py
+++ b/paimon-python/pypaimon/multimodal/video.py
@@ -37,6 +37,11 @@ class VideoFrameCollator:
order to avoid unnecessary decoder seeks. Rows are restored to their input
order before ``collate_fn`` is called.
+ Optional ``decode_batch_fn`` replaces ``decode_fn`` for each payload group.
+ It receives the cached decoder, sorted frame indices, and corresponding row
+ dictionaries, and returns a sequence with one frame per row in that order.
+ At least one of ``decode_fn`` or ``decode_batch_fn`` must be provided.
+
The cache is process-local and keyed by physical video payload identity.
``collate_fn`` defaults to PyTorch's ``default_collate`` and may be replaced
for decoders that already return batched objects.
@@ -48,16 +53,22 @@ class VideoFrameCollator:
*,
video_column,
decoder_factory,
- decode_fn,
+ decode_fn=None,
output_column="frame",
max_open_videos=8,
- collate_fn=None):
+ collate_fn=None,
+ decode_batch_fn=None):
if not video_column:
raise ValueError("video_column is required.")
if not callable(decoder_factory):
raise ValueError("decoder_factory must be callable.")
- if not callable(decode_fn):
- raise ValueError("decode_fn must be callable.")
+ if decode_fn is not None and not callable(decode_fn):
+ raise ValueError("decode_fn must be callable or None.")
+ if decode_batch_fn is not None and not callable(decode_batch_fn):
+ raise ValueError("decode_batch_fn must be callable or None.")
+ if decode_fn is None and decode_batch_fn is None:
+ raise ValueError(
+ "At least one of decode_fn or decode_batch_fn is required.")
if (
isinstance(max_open_videos, bool)
or not isinstance(max_open_videos, int)
@@ -76,6 +87,7 @@ class VideoFrameCollator:
self.video_column = video_column
self.decoder_factory = decoder_factory
self.decode_fn = decode_fn
+ self.decode_batch_fn = decode_batch_fn
self.output_column = output_column
self.max_open_videos = max_open_videos
self.collate_fn = collate_fn
@@ -126,11 +138,23 @@ class VideoFrameCollator:
for payload, frames in grouped.items():
decoder = self._decoder(payload)
- for frame_index, position, output in sorted(
- frames, key=lambda frame: frame[0]):
- output[self.output_column] = self.decode_fn(
- decoder, frame_index, output
+ frames.sort(key=lambda frame: frame[0])
+ if self.decode_batch_fn is not None:
+ values = self.decode_batch_fn(
+ decoder,
+ [index for index, _, _ in frames],
+ [output for _, _, output in frames],
+ )
+ if len(values) != len(frames):
+ raise ValueError(
+ "decode_batch_fn must return one frame per row.")
+ else:
+ values = (
+ self.decode_fn(decoder, index, output)
+ for index, _, output in frames
)
+ for (_, position, output), value in zip(frames, values):
+ output[self.output_column] = value
decoded[position] = output
return decoded
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
index 7c07c0718a..cdd027c53c 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -48,6 +48,7 @@ from pypaimon.multimodal.lerobot import load_from_lerobot
from pypaimon.multimodal.lerobot.dataset import (
_PyAVVideoDecoder,
_arrow_rows,
+ _decode_video_frames,
_decode_video_rows,
_decode_video_windows,
_image_tensor,
@@ -390,6 +391,168 @@ class LeRobotValidationTest(unittest.TestCase):
with self.assertRaises(OSError):
_open_video_decoder(stream, backend="torchcodec")
+ def test_video_batches_include_delta_frames_and_preserve_backends(self):
+ try:
+ import torch
+ except ImportError as error:
+ self.skipTest(str(error))
+
+ metadata = {
+ "repo_id": "test/video-batches",
+ "info": {
+ "codebase_version": "v3.0", "fps": 10,
+ "total_frames": 3, "total_episodes": 1, "total_tasks": 1,
+ "features": {
+ "index": {"dtype": "int64", "shape": [1]},
+ "episode_index": {"dtype": "int64", "shape": [1]},
+ "frame_index": {"dtype": "int64", "shape": [1]},
+ "timestamp": {"dtype": "float32", "shape": [1]},
+ "task_index": {"dtype": "int64", "shape": [1]},
+ "camera": {"dtype": "video", "shape": [2, 2, 3],
+ "names": ["height", "width", "channels"]},
+ },
+ },
+ "episodes": [{"episode_index": 0, "dataset_from_index": 0,
+ "dataset_to_index": 3, "length": 3,
+ "tasks": ["pick"]}],
+ "tasks": ["pick"],
+ }
+
+ class Reader(pmm.PaimonDatasetReader):
+
+ def read_indices(self, indices, columns):
+ return pa.Table.from_pylist([{
+ "index": index, "episode_index": 0,
+ "frame_index": index, "timestamp": index / 10,
+ "task_index": 0,
+ "camera": pmm.VideoFrameDescriptor(
+ "file:///shared.video", 0, 5, index + 3).serialize(),
+ } for index in indices], schema=self.schema).select(columns)
+
+ for backend, batch in (("torchcodec", True), (None, True),
+ ("pyav", False), (None, False)):
+ with self.subTest(backend=backend, batch=batch):
+ calls = []
+
+ class Decoder:
+ def __getitem__(self, index):
+ calls.append(index)
+ return torch.full((3, 2, 2), index, dtype=torch.uint8)
+
+ class BatchDecoder(Decoder):
+ def __getitem__(self, index):
+ raise AssertionError("unexpected single-frame call")
+
+ def get_frames_at(self, *, indices):
+ calls.append(indices)
+ return SimpleNamespace(data=torch.stack([
+ torch.full((3, 2, 2), i, dtype=torch.uint8)
+ for i in indices
+ ]))
+
+ decoder = BatchDecoder() if batch else Decoder()
+ module = "pypaimon.multimodal.lerobot.dataset."
+ with patch(
+ module + "_open_torchcodec_decoder",
+ return_value=decoder,
+ side_effect=None if batch else OSError("unavailable"),
+ ) as open_torchcodec, patch(
+ module + "_PyAVVideoDecoder", return_value=decoder,
+ ) as open_pyav:
+ file_io = SimpleNamespace(
+ new_input_stream=lambda path: io.BytesIO(b"video"))
+ reader = Reader(
+ metadata, file_io=file_io, video_backend=backend,
+ delta_timestamps={"camera": [-0.1, 0.0, 0.1]},
+ )
+ try:
+ last, first, duplicate = reader.get_items([2, 0, 2])
+ self.assertEqual([[3, 4, 5]] if batch else [3, 4, 5], calls)
+ self.assertEqual([3, 3, 2, 2], list(last["camera"].shape))
+ torch.testing.assert_close(
+ last["camera"][:, 0, 0, 0],
+ torch.tensor([4, 5, 5], dtype=torch.float32) / 255)
+ torch.testing.assert_close(
+ first["camera"][:, 0, 0, 0],
+ torch.tensor([3, 3, 4], dtype=torch.float32) / 255)
+ self.assertEqual([False, False, True],
+ last["camera_is_pad"].tolist())
+ self.assertEqual([True, False, False],
+ first["camera_is_pad"].tolist())
+ last["camera"].zero_()
+ self.assertGreater(float(duplicate["camera"].sum()), 0)
+ reader.get_items([1])
+ self.assertEqual(
+ [[3, 4, 5]] * 2 if batch else [3, 4, 5] * 2, calls)
+ self.assertEqual(0 if backend == "pyav" else 1,
+ open_torchcodec.call_count)
+ self.assertEqual(0 if batch else 1, open_pyav.call_count)
+ finally:
+ reader.close()
+
+ def test_video_batch_decode_propagates_errors(self):
+ decoder = Mock()
+ decoder.get_frames_at.side_effect = RuntimeError("decode failed")
+ with self.assertRaisesRegex(RuntimeError, "decode failed"):
+ _decode_video_frames(decoder, [0, 1], [{}, {}])
+ decoder.get_frames_at.assert_called_once_with(indices=[0, 1])
+
+ @unittest.skipUnless(
+ av is not None and importlib.util.find_spec("torchcodec") is not None,
+ "PyAV and TorchCodec are required for batch video decoding",
+ )
+ def test_torchcodec_batch_matches_single_frame_decoding(self):
+ import torch
+ try:
+ from torchcodec.decoders import VideoDecoder
+ except (ImportError, OSError, RuntimeError) as error:
+ self.skipTest(str(error))
+
+ output = io.BytesIO()
+ with av.open(output, mode="w", format="mp4") as container:
+ stream = container.add_stream("libx264", rate=10)
+ stream.width = stream.height = 16
+ stream.pix_fmt = "yuv420p"
+ stream.gop_size = 4
+ stream.codec_context.max_b_frames = 2
+ for index in range(12):
+ frame = av.VideoFrame.from_ndarray(
+ np.full((16, 16, 3), index * 16, dtype=np.uint8),
+ format="rgb24")
+ for packet in stream.encode(frame):
+ container.mux(packet)
+ for packet in stream.encode():
+ container.mux(packet)
+
+ payload = output.getvalue()
+ decoder = _open_video_decoder(io.BytesIO(payload), backend="torchcodec")
+ self.assertIsInstance(decoder, VideoDecoder)
+ indices = [9, 1, 9, 5]
+ expected = torch.stack([decoder[index] for index in indices])
+ file_io = SimpleNamespace(
+ new_input_stream=lambda path: io.BytesIO(payload))
+ collator = pmm.VideoFrameCollator(
+ SimpleNamespace(file_io=file_io), video_column="video",
+ decoder_factory=lambda source: decoder,
+ decode_batch_fn=_decode_video_frames,
+ collate_fn=lambda rows: rows,
+ )
+ try:
+ with patch.object(
+ decoder, "get_frames_at", wraps=decoder.get_frames_at,
+ ) as decode_batch:
+ result = collator([{
+ "video": pmm.VideoFrameDescriptor(
+ "file:///episode.mp4", 0, len(payload), index).serialize(),
+ } for index in indices])
+ decode_batch.assert_called_once_with(indices=[1, 5, 9, 9])
+ actual = torch.stack([row["frame"] for row in result])
+ self.assertEqual(torch.uint8, actual.dtype)
+ self.assertEqual((4, 3, 16, 16), tuple(actual.shape))
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+ finally:
+ collator.close()
+
def test_dataset_requires_supported_python(self):
with patch(
"pypaimon.multimodal.lerobot.dataset.sys.version_info",
diff --git a/paimon-python/pypaimon/tests/multimodal_video_test.py b/paimon-python/pypaimon/tests/multimodal_video_test.py
index 06c19cb053..918e917a32 100644
--- a/paimon-python/pypaimon/tests/multimodal_video_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_video_test.py
@@ -62,6 +62,52 @@ class VideoFrameCollatorTest(unittest.TestCase):
def tearDown(self):
self.temp_dir.cleanup()
+ def test_accepts_batch_callback_without_single_frame_callback(self):
+ descriptors = [
+ self._descriptor("episode.mp4", b"video", index)
+ for index in (2, 0)
+ ]
+ for kwargs in ({}, {"decode_fn": None}):
+ with self.subTest(kwargs=kwargs):
+ collator = VideoFrameCollator(
+ self.table,
+ video_column="video",
+ decoder_factory=lambda stream: _Decoder(stream, []),
+ decode_batch_fn=lambda decoder, indices, rows: [
+ decoder.decode(index) for index in indices
+ ],
+ collate_fn=lambda rows: rows,
+ **kwargs,
+ )
+ try:
+ result = collator([{"video": value} for value in descriptors])
+ self.assertEqual([(b"video", 2), (b"video", 0)],
+ [row["frame"] for row in result])
+ finally:
+ collator.close()
+
+ def test_rejects_missing_or_invalid_decode_callbacks(self):
+ callback = lambda *args: None
+ cases = [
+ ({}, "At least one"),
+ ({"decode_fn": None, "decode_batch_fn": None}, "At least one"),
+ ({"decode_fn": False}, "decode_fn must be callable"),
+ ({"decode_fn": False, "decode_batch_fn": callback},
+ "decode_fn must be callable"),
+ ({"decode_batch_fn": False}, "decode_batch_fn must be callable"),
+ ({"decode_fn": callback, "decode_batch_fn": False},
+ "decode_batch_fn must be callable"),
+ ]
+ for kwargs, message in cases:
+ with self.subTest(kwargs=kwargs):
+ with self.assertRaisesRegex(ValueError, message):
+ VideoFrameCollator(
+ self.table,
+ video_column="video",
+ decoder_factory=lambda stream: _Decoder(stream, []),
+ **kwargs,
+ )
+
def test_reuses_decoder_for_rows_with_same_descriptor(self):
descriptors = [
self._descriptor("episode-1.mp4", b"video-one", frame)
@@ -71,7 +117,10 @@ class VideoFrameCollatorTest(unittest.TestCase):
def factory(stream):
factory_calls.append("open")
- return _Decoder(stream, factory_calls)
+ decoder = _Decoder(stream, factory_calls)
+ decoder.get_frames_at = lambda indices: self.fail(
+ "must preserve the custom single-frame callback")
+ return decoder
collator = VideoFrameCollator(
self.table,
@@ -174,6 +223,87 @@ class VideoFrameCollatorTest(unittest.TestCase):
self.assertEqual("decoded-base", row_groups[0][4]["video"])
self.assertEqual("decoded-delta", row_groups[1][1]["video"])
+ def test_batch_decode_groups_payload_ranges_and_restores_order(self):
+ path = os.path.join(self.temp_dir.name, "shared.video")
+ with open(path, "wb") as output:
+ output.write(b"firstsecond")
+ calls = []
+ created = []
+
+ def factory(stream):
+ decoder = _Decoder(stream, [])
+ created.append(decoder)
+ return decoder
+
+ def decode_batch(decoder, indices, rows):
+ calls.append((decoder.decode(0)[0], indices))
+ return [
+ (row["request"], decoder.decode(index))
+ for index, row in zip(indices, rows)
+ ]
+
+ def row(request, offset, length, index):
+ return {"request": request, "video": VideoFrameDescriptor(
+ path, offset, length, index).serialize()}
+
+ collator = VideoFrameCollator(
+ self.table,
+ video_column="video",
+ decoder_factory=factory,
+ decode_fn=lambda *args: self.fail("unexpected single decode"),
+ decode_batch_fn=decode_batch,
+ max_open_videos=2,
+ collate_fn=lambda rows: rows,
+ )
+ rows = [
+ row("a", 0, 5, 3), row("b", 5, 6, 2),
+ row("c", 0, 5, 1), row("d", 0, 5, 3),
+ {"request": "e", "video": None}, row("f", 0, 11, 0),
+ ]
+ try:
+ result = collator(rows)
+ self.assertEqual(
+ [(b"first", [1, 3, 3]), (b"second", [2]),
+ (b"firstsecond", [0])], calls)
+ self.assertEqual([r["request"] for r in rows],
+ [r["request"] for r in result])
+ self.assertEqual(
+ [("a", (b"first", 3)), ("b", (b"second", 2)),
+ ("c", (b"first", 1)), ("d", (b"first", 3)),
+ None, ("f", (b"firstsecond", 0))],
+ [r["frame"] for r in result])
+ self.assertTrue(all("frame" not in r for r in rows))
+ self.assertEqual(3, len(created))
+ self.assertTrue(created[0].closed)
+ self.assertEqual(2, len(collator._decoders))
+ self.assertEqual(
+ ("g", (b"second", 1)),
+ collator(row("g", 5, 6, 1))["frame"])
+ self.assertEqual(3, len(created))
+ self.assertEqual([], collator([]))
+ self.assertEqual(4, len(calls))
+ finally:
+ collator.close()
+ self.assertTrue(all(decoder.closed for decoder in created))
+
+ def test_batch_decode_rejects_wrong_result_count(self):
+ descriptor = self._descriptor("episode.mp4", b"video", 0)
+ for count in (0, 2):
+ with self.subTest(count=count):
+ collator = VideoFrameCollator(
+ self.table,
+ video_column="video",
+ decoder_factory=lambda stream: _Decoder(stream, []),
+ decode_fn=lambda *args: self.fail("unexpected fallback"),
+ decode_batch_fn=lambda *args: [None] * count,
+ collate_fn=lambda rows: rows,
+ )
+ try:
+ with self.assertRaisesRegex(ValueError, "one frame per row"):
+ collator([{"video": descriptor}])
+ finally:
+ collator.close()
+
def test_evicts_least_recently_used_decoder(self):
descriptors = [
self._descriptor("episode-%d.mp4" % index, bytes([index]), index)