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

Author: NicolasHug

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -1119,12 +1119,18 @@ VideoDecoder::BatchDecodedOutput VideoDecoder::getFramesDisplayedByTimestamps(
 
     auto it = std::lower_bound(
         stream.allFrames.begin(),
-        stream.allFrames.end() - 1,
+        stream.allFrames.end(),
         framePts,
         [&stream](const FrameInfo& info, double framePts) {
           return ptsToSeconds(info.nextPts, stream.timeBase) <= framePts;
         });
     int64_t frameIndex = it - stream.allFrames.begin();
+    // If the frame index is larger than the size of allFrames, that means we
+    // couldn't match the pts value to the pts value of a NEXT FRAME. And
+    // that means that this timestamp falls during the time between when the
+    // last frame is displayed, and the video ends. Hence, it should map to the
+    // index of the last frame.
+    frameIndex = std::min(frameIndex, (int64_t)stream.allFrames.size() - 1);
     frameIndices[i] = frameIndex;
   }
 

test/decoders/test_video_decoder_ops.py

@@ -177,10 +177,10 @@ def test_get_frames_by_pts(self):
         for frame, expected_frame in zip(frames, expected_frames):
             assert_tensor_equal(frame, expected_frame)
 
-        # # first and last frame should be equal, at pts=2 [+ eps]. We then
-        # modify the # first frame and assert that it's now different from the
-        # last frame.  # This ensures a copy was properly made during the
-        # de-duplication logic.
+        # first and last frame should be equal, at pts=2 [+ eps]. We then modify
+        # the first frame and assert that it's now different from the last
+        # frame. This ensures a copy was properly made during the de-duplication
+        # logic.
         assert_tensor_equal(frames[0], frames[-1])
         frames[0] += 20
         with pytest.raises(AssertionError):