minor refac

Author: NicolasHug

src/torchcodec/_frame.py

@@ -110,15 +110,5 @@ def __getitem__(self, key) -> Union["FrameBatch", Frame]:
     def __len__(self):
         return len(self.data)
 
-    def __getitem__(self, key):
-        return FrameBatch(
-            self.data[key],
-            self.pts_seconds[key],
-            self.duration_seconds[key],
-        )
-
-    def __len__(self):
-        return len(self.data)
-
     def __repr__(self):
         return _frame_repr(self)

src/torchcodec/samplers/_common.py

@@ -1,5 +1,8 @@
 from typing import Callable, Union
 
+from torch import Tensor
+from torchcodec import FrameBatch
+
 _LIST_OF_INT_OR_FLOAT = Union[list[int], list[float]]
 
 
@@ -53,3 +56,19 @@ def _validate_common_params(*, decoder, num_frames_per_clip, policy):
         raise ValueError(
             f"Invalid policy ({policy}). Supported values are {_POLICY_FUNCTIONS.keys()}."
         )
+
+
+def _make_5d_framebatch(
+    *,
+    data: Tensor,
+    pts_seconds: Tensor,
+    duration_seconds: Tensor,
+    num_clips: int,
+    num_frames_per_clip: int,
+) -> FrameBatch:
+    last_3_dims = data.shape[-3:]
+    return FrameBatch(
+        data=data.view(num_clips, num_frames_per_clip, *last_3_dims),
+        pts_seconds=pts_seconds.view(num_clips, num_frames_per_clip),
+        duration_seconds=duration_seconds.view(num_clips, num_frames_per_clip),
+    )

src/torchcodec/samplers/_index_based.py

@@ -6,6 +6,7 @@
 from torchcodec.decoders import VideoDecoder
 from torchcodec.decoders._core import get_frames_at_indices
 from torchcodec.samplers._common import (
+    _make_5d_framebatch,
     _POLICY_FUNCTION_TYPE,
     _POLICY_FUNCTIONS,
     _validate_common_params,
@@ -176,16 +177,18 @@ def _generic_index_based_sampler(
         policy_fun=_POLICY_FUNCTIONS[policy],
     )
 
+    # TODO: Use public method of decoder, when it exists
     frames, pts_seconds, duration_seconds = get_frames_at_indices(
         decoder._decoder,
         stream_index=decoder.stream_index,
         frame_indices=all_clips_indices,
     )
-    last_3_dims = frames.shape[-3:]
-    return FrameBatch(
-        data=frames.view(num_clips, num_frames_per_clip, *last_3_dims),
-        pts_seconds=pts_seconds.view(num_clips, num_frames_per_clip),
-        duration_seconds=duration_seconds.view(num_clips, num_frames_per_clip),
+    return _make_5d_framebatch(
+        data=frames,
+        pts_seconds=pts_seconds,
+        duration_seconds=duration_seconds,
+        num_clips=num_clips,
+        num_frames_per_clip=num_frames_per_clip,
     )
 
 

src/torchcodec/samplers/_time_based.py

@@ -5,6 +5,7 @@
 from torchcodec import FrameBatch
 from torchcodec.decoders._core import get_frames_by_pts
 from torchcodec.samplers._common import (
+    _make_5d_framebatch,
     _POLICY_FUNCTION_TYPE,
     _POLICY_FUNCTIONS,
     _validate_common_params,
@@ -209,17 +210,18 @@ def _generic_time_based_sampler(
         policy_fun=_POLICY_FUNCTIONS[policy],
     )
 
+    # TODO: Use public method of decoder, when it exists
     frames, pts_seconds, duration_seconds = get_frames_by_pts(
         decoder._decoder,
         stream_index=decoder.stream_index,
         frame_ptss=all_clips_timestamps,
     )
-    last_3_dims = frames.shape[-3:]
-
-    return FrameBatch(
-        data=frames.view(num_clips, num_frames_per_clip, *last_3_dims),
-        pts_seconds=pts_seconds.view(num_clips, num_frames_per_clip),
-        duration_seconds=duration_seconds.view(num_clips, num_frames_per_clip),
+    return _make_5d_framebatch(
+        data=frames,
+        pts_seconds=pts_seconds,
+        duration_seconds=duration_seconds,
+        num_clips=num_clips,
+        num_frames_per_clip=num_frames_per_clip,
     )
 
 

test/samplers/test_samplers.py

@@ -26,8 +26,6 @@ def _assert_output_type_and_shapes(
     video, clips, expected_num_clips, num_frames_per_clip
 ):
     assert isinstance(clips, FrameBatch)
-    # assert len(clips) == expected_num_clips
-    # assert all(isinstance(clip, FrameBatch) for clip in clips)
     expected_clips_data_shape = (
         expected_num_clips,
         num_frames_per_clip,