XiaoHongbo-Hope commented on code in PR #9757:
URL: https://github.com/apache/paimon/pull/9757#discussion_r3998942839
##########
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:
Fixed in 4cfabc7b34. A cold non-zero read now demuxes packet
timestamps/keyframe flags once, builds the complete ordinal/keyframe index,
seeks to the target GOP, and decodes only from that keyframe. The new
regression requests frame 95 from a fresh 120-frame decoder and verifies only
frames 90–95 are decoded; I also reran random reads on a real MP4 with B-frames.
--
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]
