Oops, fix

Author: NicolasHug

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -855,14 +855,6 @@ torch::Tensor VideoDecoder::getFramesPlayedInRangeAudio(
 
   StreamInfo& streamInfo = streamInfos_[activeStreamIndex_];
 
-  auto lastDecodedFrameIsPlayedAt = [&streamInfo](double seconds) {
-    auto pts = secondsToClosestPts(seconds, streamInfo.timeBase);
-    return (
-        streamInfo.lastDecodedAvFramePts <= pts and
-        pts <= streamInfo.lastDecodedAvFramePts +
-                streamInfo.lastDecodedAvFrameDuration);
-  };
-
   // TODO-AUDIO This essentially enforce that we don't need to seek (backwards).
   // We should remove it and seek back to the stream's beginning when needed.
   // See test_multiple_calls
@@ -871,8 +863,7 @@ torch::Tensor VideoDecoder::getFramesPlayedInRangeAudio(
        streamInfo.lastDecodedAvFrameDuration == 0) ||
           (streamInfo.lastDecodedAvFramePts +
                streamInfo.lastDecodedAvFrameDuration <=
-           secondsToClosestPts(startSeconds, streamInfo.timeBase)) ||
-          !lastDecodedFrameIsPlayedAt(startSeconds),
+           secondsToClosestPts(startSeconds, streamInfo.timeBase)),
       "The previous call's stop_seconds is larger than the current calls's start_seconds (roughly)");
 
   setCursorPtsInSeconds(startSeconds);
@@ -883,7 +874,23 @@ torch::Tensor VideoDecoder::getFramesPlayedInRangeAudio(
   std::vector<torch::Tensor> tensors;
 
   bool reachedEOF = false;
-  while (!lastDecodedFrameIsPlayedAt(stopSeconds) && !reachedEOF) {
+  auto shouldStopDecoding = [&streamInfo, stopSeconds, &reachedEOF]() {
+    if (reachedEOF) {
+      return true;
+    }
+    // Return true iff stopSeconds is in [begin, end] of the last decoded frame.
+    // We use [begin, end] and not [begin, end), which may seem counter
+    // intuitive, but this actually ensures that stopSeconds is an open upper
+    // bound, i.e. a frame that starts on stopSeconds won't be part of the
+    // output.
+    auto pts = secondsToClosestPts(stopSeconds, streamInfo.timeBase);
+    auto lastDecodedAvFrameEnd = streamInfo.lastDecodedAvFramePts +
+        streamInfo.lastDecodedAvFrameDuration;
+    return (streamInfo.lastDecodedAvFramePts) <= pts &&
+        (pts <= lastDecodedAvFrameEnd);
+  };
+
+  while (!shouldStopDecoding()) {
     try {
       AVFrameStream avFrameStream =
           decodeAVFrame([&streamInfo](AVFrame* avFrame) {

test/decoders/test_ops.py

@@ -750,6 +750,7 @@ def get_reference_frames(start_seconds, stop_seconds):
             frames, get_reference_frames(start_seconds, stop_seconds)
         )
 
+        # "seeking" forward is OK
         start_seconds, stop_seconds = 3, 4
         frames = get_frames_by_pts_in_range_audio(
             decoder, start_seconds=start_seconds, stop_seconds=stop_seconds
@@ -771,7 +772,7 @@ def get_reference_frames(start_seconds, stop_seconds):
             frames, get_reference_frames(start_seconds, stop_seconds)
         )
 
-        # But starting immediately on the same frame isn't OK
+        # but starting immediately on the same frame raises
         with pytest.raises(
             RuntimeError,
             match="The previous call's stop_seconds is larger than the current calls's start_seconds",
@@ -788,13 +789,13 @@ def get_reference_frames(start_seconds, stop_seconds):
                 decoder, start_seconds=stop_seconds + 1e-4, stop_seconds=6
             )
 
-        start_seconds, stop_seconds = 0, 2
-        frames = get_frames_by_pts_in_range_audio(
-            decoder, start_seconds=start_seconds, stop_seconds=stop_seconds
-        )
-        with pytest.raises(AssertionError):
-            torch.testing.assert_close(
-                frames, get_reference_frames(start_seconds, stop_seconds)
+        # and seeking backwards doesn't work either
+        with pytest.raises(
+            RuntimeError,
+            match="The previous call's stop_seconds is larger than the current calls's start_seconds",
+        ):
+            frames = get_frames_by_pts_in_range_audio(
+                decoder, start_seconds=0, stop_seconds=2
             )