Create key frame index manually

Author: scotts

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -561,9 +561,7 @@ torch::Tensor VideoDecoder::getKeyFrameIndices(int streamIndex) {
   torch::Tensor keyFrameIndices =
       torch::empty({static_cast<int64_t>(keyFrames.size())}, {torch::kInt64});
   for (size_t i = 0; i < keyFrames.size(); ++i) {
-    int64_t pts = keyFrames[i].pts;
-    keyFrameIndices[i] =
-        getKeyFrameIndexForPtsUsingScannedIndex(keyFrames, pts);
+    keyFrameIndices[i] = keyFrames[i].frameIndex;
   }
 
   return keyFrameIndices;
@@ -685,7 +683,13 @@ void VideoDecoder::scanFileAndUpdateMetadataAndIndex() {
           return frameInfo1.pts < frameInfo2.pts;
         });
 
+    size_t keyIndex = 0;
     for (size_t i = 0; i < streamInfo.allFrames.size(); ++i) {
+      streamInfo.allFrames[i].frameIndex = i;
+      if (streamInfo.keyFrames[keyIndex].pts == streamInfo.allFrames[i].pts) {
+        streamInfo.keyFrames[keyIndex].frameIndex = i;
+        ++keyIndex;
+      }
       if (i + 1 < streamInfo.allFrames.size()) {
         streamInfo.allFrames[i].nextPts = streamInfo.allFrames[i + 1].pts;
       }

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -291,12 +291,19 @@ class VideoDecoder {
 
   struct FrameInfo {
     int64_t pts = 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.)
+
+    // The value of the nextPts 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;
+
+    // Note that frameIndex is ALWAYS the index into all of the frames in that
+    // stream, even when the FrameInfo is part of the key frame index. Given a
+    // FrameInfo for a key frame, the frameIndex allows us to know which frame
+    // that is in the stream.
+    int64_t frameIndex = 0;
   };
 
   struct FilterGraphContext {

test/decoders/test_video_decoder.py

@@ -835,9 +835,23 @@ def test_get_frames_by_pts_in_range_fails(self, device, seek_mode):
     def test_get_key_frame_indices(self, device):
         decoder = VideoDecoder(NASA_VIDEO.path, device=device, seek_mode="exact")
         key_frame_indices = decoder._get_key_frame_indices()
-        size = key_frame_indices.size()
-        assert size[0] > 0
-        assert len(size) == 1
+
+        # The key frame indices were generated from the following command:
+        #   $ ffprobe -v error -hide_banner -select_streams v:1 -show_frames -of csv test/resources/nasa_13013.mp4 | grep -n ",I," | cut -d ':' -f 1 > key_frames.txt
+        # What it's doing:
+        #   1. Calling ffprobe on the second video stream, which is absolute stream index 3.
+        #   2. Showing all frames for that stream.
+        #   3. Using grep to find and count the "I" frames, which are the key frames.
+        #   4. Using cut to extract just the count for the frame.
+        # Finally, because the above produces a count, which is index + 1, we subtract
+        # one from all values manually to arrive at the values below.
+        # TODO: decide if/how we want to incorporate key frame indices into the utils
+        # framework.
+        reference_key_frame_indices = torch.tensor([0, 240])
+
+        torch.testing.assert_close(
+            key_frame_indices, reference_key_frame_indices, atol=0, rtol=0
+        )
 
 
 if __name__ == "__main__":