Merge branch 'main' of github.com:pytorch/torchcodec into samplers_hack

Author: NicolasHug

src/torchcodec/_frame.py

@@ -39,6 +39,8 @@ class Frame(Iterable):
     """The duration of the frame, in seconds (float)."""
 
     def __post_init__(self):
+        # This is called after __init__() when a Frame is created. We can run
+        # input validation checks here.
         if not self.data.ndim == 3:
             raise ValueError(f"data must be 3-dimensional, got {self.data.shape = }")
         self.pts_seconds = float(self.pts_seconds)
@@ -64,6 +66,8 @@ class FrameBatch(Iterable):
     """The duration of the frame, in seconds (1-D ``torch.Tensor`` of floats)."""
 
     def __post_init__(self):
+        # This is called after __init__() when a FrameBatch is created. We can
+        # run input validation checks here.
         if self.data.ndim < 4:
             raise ValueError(
                 f"data must be at least 4-dimensional. Got {self.data.shape = } "

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -299,7 +299,12 @@ class VideoDecoder {
  private:
   struct FrameInfo {
     int64_t pts = 0;
-    int64_t nextPts = 0;
+    // The value of this default is important: the last frame's nextPts will be
+    // INT64_MAX, which ensures that the allFrames vec contains FrameInfo
+    // structs with *increasing* nextPts values. That's a necessary condition
+    // for the binary searches on those values to work properly (as typically
+    // done during pts -> index conversions.)
+    int64_t nextPts = INT64_MAX;
   };
   struct FilterState {
     UniqueAVFilterGraph filterGraph;

test/decoders/test_video_decoder_ops.py

@@ -187,7 +187,7 @@ def test_get_frames_by_pts(self):
             assert_tensor_equal(frames[0], frames[-1])
 
     def test_pts_apis_against_index_ref(self):
-        # Non-regression test for https://github.com/pytorch/torchcodec/pull/286
+        # Non-regression test for https://github.com/pytorch/torchcodec/pull/287
         # Get all frames in the video, then query all frames with all time-based
         # APIs exactly where those frames are supposed to start. We assert that
         # we get the expected frame.
@@ -227,6 +227,20 @@ def test_pts_apis_against_index_ref(self):
         )
         assert_tensor_equal(pts_seconds, all_pts_seconds_ref)
 
+        _, pts_seconds, _ = zip(
+            *[
+                get_frames_by_pts_in_range(
+                    decoder,
+                    stream_index=stream_index,
+                    start_seconds=pts,
+                    stop_seconds=pts + 1e-4,
+                )
+                for pts in all_pts_seconds_ref
+            ]
+        )
+        pts_seconds = torch.tensor(pts_seconds)
+        assert_tensor_equal(pts_seconds, all_pts_seconds_ref)
+
         _, pts_seconds, _ = get_frames_by_pts(
             decoder, stream_index=stream_index, timestamps=all_pts_seconds_ref.tolist()
         )