Clean up some logic

Author: NicolasHug

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -658,20 +658,16 @@ VideoDecoder::getFramesInRange(int64_t start, int64_t stop, int64_t step) {
   return frameBatchOutput;
 }
 
-VideoDecoder::FrameOutput VideoDecoder::getFramePlayedAt(
-    double seconds) {
-  for (auto& [streamIndex, streamInfo] : streamInfos_) {
-    double frameStartTime =
-        ptsToSeconds(streamInfo.currentPts, streamInfo.timeBase);
-    double frameEndTime = ptsToSeconds(
-        streamInfo.currentPts + streamInfo.currentDuration,
-        streamInfo.timeBase);
-    if (seconds >= frameStartTime && seconds < frameEndTime) {
-      // We are in the same frame as the one we just returned. However, since we
-      // don't cache it locally, we have to rewind back.
-      seconds = frameStartTime;
-      break;
-    }
+VideoDecoder::FrameOutput VideoDecoder::getFramePlayedAt(double seconds) {
+  StreamInfo& streamInfo = streamInfos_[activeStreamIndex_];
+  double frameStartTime =
+      ptsToSeconds(streamInfo.currentPts, streamInfo.timeBase);
+  double frameEndTime = ptsToSeconds(
+      streamInfo.currentPts + streamInfo.currentDuration, streamInfo.timeBase);
+  if (seconds >= frameStartTime && seconds < frameEndTime) {
+    // We are in the same frame as the one we just returned. However, since we
+    // don't cache it locally, we have to rewind back.
+    seconds = frameStartTime;
   }
 
   setCursorPtsInSeconds(seconds);

test/decoders/VideoDecoderTest.cpp

@@ -285,8 +285,7 @@ TEST_P(VideoDecoderTest, GetsFramePlayedAtTimestamp) {
       kPtsOfLastFrameInVideoStream + kDurationOfLastFrameInVideoStream;
   // Sanity check: make sure duration is strictly positive.
   EXPECT_GT(kPtsPlusDurationOfLastFrame, kPtsOfLastFrameInVideoStream);
-  output =
-      ourDecoder->getFramePlayedAt(kPtsPlusDurationOfLastFrame - 1e-6);
+  output = ourDecoder->getFramePlayedAt(kPtsPlusDurationOfLastFrame - 1e-6);
   EXPECT_EQ(output.ptsSeconds, kPtsOfLastFrameInVideoStream);
 }