Use new public methods in samplers

Author: NicolasHug

src/torchcodec/samplers/_common.py

@@ -1,6 +1,5 @@
 from typing import Callable, Union
 
-from torch import Tensor
 from torchcodec import FrameBatch
 
 _LIST_OF_INT_OR_FLOAT = Union[list[int], list[float]]
@@ -58,17 +57,15 @@ def _validate_common_params(*, decoder, num_frames_per_clip, policy):
         )
 
 
-def _make_5d_framebatch(
+def _reshape_4d_framebatch_into_5d(
     *,
-    data: Tensor,
-    pts_seconds: Tensor,
-    duration_seconds: Tensor,
+    frames: FrameBatch,
     num_clips: int,
     num_frames_per_clip: int,
 ) -> FrameBatch:
-    last_3_dims = data.shape[-3:]
+    last_3_dims = frames.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),
+        data=frames.data.view(num_clips, num_frames_per_clip, *last_3_dims),
+        pts_seconds=frames.pts_seconds.view(num_clips, num_frames_per_clip),
+        duration_seconds=frames.duration_seconds.view(num_clips, num_frames_per_clip),
     )

src/torchcodec/samplers/_index_based.py

@@ -4,11 +4,10 @@
 
 from torchcodec import FrameBatch
 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,
+    _reshape_4d_framebatch_into_5d,
     _validate_common_params,
 )
 
@@ -177,16 +176,9 @@ 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,
-    )
-    return _make_5d_framebatch(
-        data=frames,
-        pts_seconds=pts_seconds,
-        duration_seconds=duration_seconds,
+    frames = decoder.get_frames_at(indices=all_clips_indices)
+    return _reshape_4d_framebatch_into_5d(
+        frames=frames,
         num_clips=num_clips,
         num_frames_per_clip=num_frames_per_clip,
     )

src/torchcodec/samplers/_time_based.py

@@ -3,11 +3,10 @@
 import torch
 
 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,
+    _reshape_4d_framebatch_into_5d,
     _validate_common_params,
 )
 
@@ -210,16 +209,9 @@ 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,
-        timestamps=all_clips_timestamps,
-    )
-    return _make_5d_framebatch(
-        data=frames,
-        pts_seconds=pts_seconds,
-        duration_seconds=duration_seconds,
+    frames = decoder.get_frames_displayed_at(seconds=all_clips_timestamps)
+    return _reshape_4d_framebatch_into_5d(
+        frames=frames,
         num_clips=num_clips,
         num_frames_per_clip=num_frames_per_clip,
     )