Merge pull request #14 from open-world-agents/feat/get-frames-played-in-range

feat: add get_frames_played_in_range to video decoder hierarchy

Author: MilkClouds

src/mediaref/video_decoder/base.py

@@ -1,7 +1,7 @@
 """Base interface for video decoders."""
 
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Optional
 
 from .._typing import PathLike
 from .frame_batch import FrameBatch
@@ -59,6 +59,38 @@ def get_frames_played_at(self, seconds: List[float]) -> FrameBatch:
         """
         pass
 
+    @abstractmethod
+    def get_frames_played_in_range(
+        self, start_seconds: float, stop_seconds: float, fps: Optional[float] = None
+    ) -> FrameBatch:
+        """Return multiple frames in the given range.
+
+        Frames are in the half open range [start_seconds, stop_seconds). Each
+        returned frame's :term:`pts`, in seconds, is inside of the half open
+        range.
+
+        Args:
+            start_seconds: Time, in seconds, of the start of the range.
+            stop_seconds: Time, in seconds, of the end of the range.
+                As a half open range, the end is excluded.
+            fps: If specified, resample output to this frame rate by
+                duplicating or dropping frames as necessary. If None
+                (default), returns frames at the source video's frame rate.
+
+        Returns:
+            FrameBatch: The frames within the specified range.
+
+        Raises:
+            ValueError: If start_seconds > stop_seconds, or if the range
+                is outside the valid stream bounds.
+
+        Examples:
+            >>> with PyAVVideoDecoder("video.mp4") as decoder:
+            ...     batch = decoder.get_frames_played_in_range(0.0, 2.0)
+            ...     print(batch.data.shape)
+        """
+        pass
+
     @abstractmethod
     def close(self):
         """Release video decoder resources.

src/mediaref/video_decoder/pyav_decoder.py

@@ -2,7 +2,7 @@
 
 import gc
 from fractions import Fraction
-from typing import List
+from typing import List, Optional
 
 import av
 import cv2
@@ -44,6 +44,17 @@ def _frame_to_rgba(frame: av.VideoFrame) -> npt.NDArray[np.uint8]:
     return rgba_array
 
 
+def _convert_av_frames_to_nchw(av_frames: List[av.VideoFrame]) -> List[npt.NDArray[np.uint8]]:
+    """Convert a list of PyAV frames to NCHW numpy arrays (RGB)."""
+    frames = []
+    for frame in av_frames:
+        rgba_array = _frame_to_rgba(frame)
+        rgb_array = cv2.cvtColor(rgba_array, cv2.COLOR_RGBA2RGB)
+        frame_nchw = np.transpose(rgb_array, (2, 0, 1)).astype(np.uint8)
+        frames.append(frame_nchw)
+    return frames
+
+
 class PyAVVideoDecoder(BaseVideoDecoder):
     """Video decoder using PyAV with TorchCodec-compatible playback semantics.
 
@@ -128,6 +139,14 @@ def metadata(self) -> VideoStreamMetadata:
         """Access video stream metadata."""
         return self._metadata
 
+    def _create_empty_batch(self) -> FrameBatch:
+        """Create an empty FrameBatch with correct spatial dimensions."""
+        return FrameBatch(
+            data=np.empty((0, 3, self._metadata.height, self._metadata.width), dtype=np.uint8),
+            pts_seconds=np.array([], dtype=np.float64),
+            duration_seconds=np.array([], dtype=np.float64),
+        )
+
     def get_frames_played_at(self, seconds: List[float]) -> FrameBatch:
         """Retrieve frames that would be displayed at specific timestamps.
 
@@ -144,11 +163,7 @@ def get_frames_played_at(self, seconds: List[float]) -> FrameBatch:
             ValueError: If any timestamp is outside [begin_stream_seconds, end_stream_seconds)
         """
         if not seconds:
-            return FrameBatch(
-                data=np.empty((0, 3, self._metadata.height, self._metadata.width), dtype=np.uint8),
-                pts_seconds=np.array([], dtype=np.float64),
-                duration_seconds=np.array([], dtype=np.float64),
-            )
+            return self._create_empty_batch()
 
         # Validate timestamps per playback_semantics.md boundary conditions
         begin_stream = float(self._metadata.begin_stream_seconds)
@@ -163,12 +178,7 @@ def get_frames_played_at(self, seconds: List[float]) -> FrameBatch:
         av_frames = self._get_frames_played_at(seconds)
 
         # Convert to RGB numpy arrays in NCHW format
-        frames = []
-        for frame in av_frames:
-            rgba_array = _frame_to_rgba(frame)
-            rgb_array = cv2.cvtColor(rgba_array, cv2.COLOR_RGBA2RGB)
-            frame_nchw = np.transpose(rgb_array, (2, 0, 1)).astype(np.uint8)
-            frames.append(frame_nchw)
+        frames = _convert_av_frames_to_nchw(av_frames)
 
         pts_list = [float(frame.time) for frame in av_frames]
         duration = float(1.0 / self._metadata.average_rate)
@@ -179,6 +189,75 @@ def get_frames_played_at(self, seconds: List[float]) -> FrameBatch:
             duration_seconds=np.full(len(seconds), duration, dtype=np.float64),
         )
 
+    def get_frames_played_in_range(
+        self, start_seconds: float, stop_seconds: float, fps: Optional[float] = None
+    ) -> FrameBatch:
+        """Return multiple frames in the given range [start_seconds, stop_seconds).
+
+        Args:
+            start_seconds: Time, in seconds, of the start of the range.
+            stop_seconds: Time, in seconds, of the end of the range (excluded).
+            fps: If specified, resample output to this frame rate by
+                duplicating or dropping frames as necessary. If None,
+                returns frames at the source video's frame rate.
+
+        Returns:
+            FrameBatch with frame data in NCHW format.
+
+        Raises:
+            ValueError: If the range parameters are invalid.
+        """
+        begin_stream = float(self._metadata.begin_stream_seconds)
+        end_stream = float(self._metadata.end_stream_seconds)
+
+        if not start_seconds <= stop_seconds:
+            raise ValueError(
+                f"Invalid start seconds: {start_seconds}. "
+                f"It must be less than or equal to stop seconds ({stop_seconds})."
+            )
+        if not begin_stream <= start_seconds < end_stream:
+            raise ValueError(
+                f"Invalid start seconds: {start_seconds}. "
+                f"It must be greater than or equal to {begin_stream} "
+                f"and less than {end_stream}."
+            )
+        if not stop_seconds <= end_stream:
+            raise ValueError(f"Invalid stop seconds: {stop_seconds}. It must be less than or equal to {end_stream}.")
+
+        if fps is not None:
+            # Resample: generate timestamps at the given fps and get frames
+            timestamps = np.arange(start_seconds, stop_seconds, 1.0 / fps).tolist()
+            if not timestamps:
+                return self._create_empty_batch()
+            return self.get_frames_played_at(timestamps)
+
+        # Native frame rate: decode all frames with pts in [start_seconds, stop_seconds)
+        self._seek_to_or_before(start_seconds)
+
+        av_frames: List[av.VideoFrame] = []
+        for frame in self._container.decode(video=0):
+            if frame.time is None:
+                raise ValueError("Frame time is None")
+            frame_pts = float(frame.time)
+            if frame_pts >= stop_seconds:
+                break
+            if frame_pts >= start_seconds:
+                av_frames.append(frame)
+
+        if not av_frames:
+            return self._create_empty_batch()
+
+        frames = _convert_av_frames_to_nchw(av_frames)
+
+        pts_list = [float(frame.time) for frame in av_frames]
+        duration = float(1.0 / self._metadata.average_rate)
+
+        return FrameBatch(
+            data=np.stack(frames, axis=0),
+            pts_seconds=np.array(pts_list, dtype=np.float64),
+            duration_seconds=np.full(len(av_frames), duration, dtype=np.float64),
+        )
+
     def _get_frames_played_at(self, seconds: List[float]) -> List[av.VideoFrame]:
         """Get frames using TorchCodec playback semantics.
 

src/mediaref/video_decoder/torchcodec_decoder.py

@@ -1,6 +1,6 @@
 """TorchCodec-based video decoder."""
 
-from typing import ClassVar, List
+from typing import ClassVar, List, Optional
 
 import numpy as np
 from torchcodec.decoders import VideoDecoder
@@ -69,6 +69,47 @@ def get_frames_played_at(self, seconds: List[float]) -> FrameBatch:
             duration_seconds=torchcodec_batch.duration_seconds.numpy().astype(np.float64),
         )
 
+    def get_frames_played_in_range(
+        self, start_seconds: float, stop_seconds: float, fps: Optional[float] = None
+    ) -> FrameBatch:
+        """Return multiple frames in the given range.
+
+        Delegates to TorchCodec's native get_frames_played_in_range.
+
+        Args:
+            start_seconds: Time, in seconds, of the start of the range.
+            stop_seconds: Time, in seconds, of the end of the range (excluded).
+            fps: If specified, resample output to this frame rate. If None,
+                returns frames at the source video's frame rate.
+
+        Returns:
+            FrameBatch containing frame data and timing information.
+
+        Raises:
+            NotImplementedError: If ``fps`` is specified but the installed
+                TorchCodec version (<=0.10.0) does not support it.
+        """
+        if fps is not None:
+            try:
+                torchcodec_batch = VideoDecoder.get_frames_played_in_range(
+                    self, start_seconds=start_seconds, stop_seconds=stop_seconds, fps=fps
+                )
+            except TypeError:
+                raise NotImplementedError(
+                    "The installed version of TorchCodec (<=0.10.0) does not support "
+                    "the 'fps' parameter in get_frames_played_in_range. "
+                    "Upgrade TorchCodec or use fps=None."
+                )
+        else:
+            torchcodec_batch = VideoDecoder.get_frames_played_in_range(
+                self, start_seconds=start_seconds, stop_seconds=stop_seconds
+            )
+        return FrameBatch(
+            data=torchcodec_batch.data.numpy(),
+            pts_seconds=torchcodec_batch.pts_seconds.numpy().astype(np.float64),
+            duration_seconds=torchcodec_batch.duration_seconds.numpy().astype(np.float64),
+        )
+
     def close(self):
         """Release cache reference. Safe to call multiple times."""
         if hasattr(self, "_cache_key") and self._cache_key in self.cache and self.cache[self._cache_key].refs > 0: