JingsongLi commented on code in PR #9757:
URL: https://github.com/apache/paimon/pull/9757#discussion_r3998881264
##########
paimon-python/pypaimon/multimodal/lerobot/dataset.py:
##########
@@ -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)
Review Comment:
[P2] Avoid decoding the full video prefix on cold random reads
When PyAV is selected explicitly or used as the fallback, a fresh decoder
skips seeking and decodes every frame from 0 through the first requested index.
The collator retains only eight videos per camera, so shuffled training over
more files can repeatedly discard the learned timestamp/keyframe index and pay
this cost again.
On a 6,000-frame, 62 MB H.264 MP4, reading frame 5,999 decoded all 6,000
frames and took about 1.4 seconds. Scanning packet timestamps took 12–15 ms,
and seeking to the final GOP and decoding its 30 frames took about 8 ms.
Please build or reuse an ordinal/keyframe index from demuxed packets so cold
random reads can seek to the required GOP while preserving exact frame
ordinals. The existing importer already uses packet timestamps to determine
frame ordinals.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
