Rough implementation of getting key frame indices

Author: scotts

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -553,6 +553,20 @@ VideoDecoder::ContainerMetadata VideoDecoder::getContainerMetadata() const {
   return containerMetadata_;
 }
 
+std::vector<int64_t> VideoDecoder::getKeyFrameIndices(int streamIndex) {
+  validateUserProvidedStreamIndex(streamIndex);
+  validateScannedAllStreams("getKeyFrameIndices");
+
+  std::vector<int64_t> keyFrameIndices;
+  const StreamInfo& streamInfo = streamInfos_[streamIndex];
+  for (const FrameInfo& frameInfo : streamInfo.keyFrames) {
+    keyFrameIndices.push_back(
+        getKeyFrameIndexForPtsUsingScannedIndex(streamInfo.keyFrames, frameInfo.pts));
+  }
+
+  return keyFrameIndices;
+}
+
 int VideoDecoder::getKeyFrameIndexForPtsUsingEncoderIndex(
     AVStream* stream,
     int64_t pts) const {

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -100,6 +100,8 @@ class VideoDecoder {
   // Returns the metadata for the container.
   ContainerMetadata getContainerMetadata() const;
 
+  std::vector<int64_t> getKeyFrameIndices(int streamIndex);
+
   // --------------------------------------------------------------------------
   // ADDING STREAMS API
   // --------------------------------------------------------------------------

src/torchcodec/decoders/_core/VideoDecoderOps.cpp

@@ -48,6 +48,7 @@ TORCH_LIBRARY(torchcodec_ns, m) {
       "get_frames_by_pts_in_range(Tensor(a!) decoder, *, int stream_index, float start_seconds, float stop_seconds) -> (Tensor, Tensor, Tensor)");
   m.def(
       "get_frames_by_pts(Tensor(a!) decoder, *, int stream_index, float[] timestamps) -> (Tensor, Tensor, Tensor)");
+  m.def("get_key_frame_indices(Tensor(a!) decoder, int stream_index) -> int[]");
   m.def("get_json_metadata(Tensor(a!) decoder) -> str");
   m.def("get_container_json_metadata(Tensor(a!) decoder) -> str");
   m.def(
@@ -334,6 +335,13 @@ bool _test_frame_pts_equality(
       videoDecoder->getPtsSecondsForFrame(stream_index, frame_index);
 }
 
+std::vector<int64_t> get_key_frame_indices(
+    at::Tensor& decoder,
+    int64_t stream_index) {
+  auto videoDecoder = unwrapTensorToGetDecoder(decoder);
+  return videoDecoder->getKeyFrameIndices(stream_index);
+}
+
 std::string get_json_metadata(at::Tensor& decoder) {
   auto videoDecoder = unwrapTensorToGetDecoder(decoder);
 
@@ -526,6 +534,7 @@ TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) {
   m.impl("add_video_stream", &add_video_stream);
   m.impl("_add_video_stream", &_add_video_stream);
   m.impl("get_next_frame", &get_next_frame);
+  m.impl("get_key_frame_indices", &get_key_frame_indices);
   m.impl("get_json_metadata", &get_json_metadata);
   m.impl("get_container_json_metadata", &get_container_json_metadata);
   m.impl("get_stream_json_metadata", &get_stream_json_metadata);

src/torchcodec/decoders/_core/VideoDecoderOps.h

@@ -137,6 +137,10 @@ bool _test_frame_pts_equality(
     int64_t frame_index,
     double pts_seconds_to_test);
 
+std::vector<int64_t> get_key_frame_indices(
+    at::Tensor& decoder,
+    int64_t stream_index);
+
 // Get the metadata from the video as a string.
 std::string get_json_metadata(at::Tensor& decoder);
 

src/torchcodec/decoders/_core/__init__.py

@@ -26,6 +26,7 @@
     get_frames_by_pts_in_range,
     get_frames_in_range,
     get_json_metadata,
+    get_key_frame_indices,
     get_next_frame,
     scan_all_streams_to_update_metadata,
     seek_to_pts,

src/torchcodec/decoders/_core/video_decoder_ops.py

@@ -77,6 +77,7 @@ def load_torchcodec_extension():
 get_frames_by_pts = torch.ops.torchcodec_ns.get_frames_by_pts.default
 get_frames_in_range = torch.ops.torchcodec_ns.get_frames_in_range.default
 get_frames_by_pts_in_range = torch.ops.torchcodec_ns.get_frames_by_pts_in_range.default
+get_key_frame_indices = torch.ops.torchcodec_ns.get_key_frame_indices.default
 get_json_metadata = torch.ops.torchcodec_ns.get_json_metadata.default
 _test_frame_pts_equality = torch.ops.torchcodec_ns._test_frame_pts_equality.default
 _get_container_json_metadata = (
@@ -255,6 +256,11 @@ def get_frames_by_pts_in_range_abstract(
     )
 
 
+@register_fake("torchcodec_ns::get_key_frame_indices")
+def get_key_frame_indices_abstract(decoder: torch.Tensor, *, stream_index: int) -> List[int]:
+    return []
+
+
 @register_fake("torchcodec_ns::get_json_metadata")
 def get_json_metadata_abstract(decoder: torch.Tensor) -> str:
     return ""

src/torchcodec/decoders/_video_decoder.py

@@ -185,6 +185,9 @@ def __getitem__(self, key: Union[numbers.Integral, slice]) -> Tensor:
             f"Unsupported key type: {type(key)}. Supported types are int and slice."
         )
 
+    def _get_key_frame_indices(self) -> list[int]:
+        return core.get_key_frame_indices(self._decoder, stream_index=self.stream_index)
+
     def get_frame_at(self, index: int) -> Frame:
         """Return a single frame at the given index.
 

test/decoders/test_video_decoder.py

@@ -831,6 +831,12 @@ def test_get_frames_by_pts_in_range_fails(self, device, seek_mode):
         with pytest.raises(ValueError, match="Invalid stop seconds"):
             frame = decoder.get_frames_played_in_range(0, 23)  # noqa
 
+    @pytest.mark.parametrize("device", cpu_and_cuda())
+    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()
+        assert len(key_frame_indices) > 0
+
 
 if __name__ == "__main__":
     pytest.main()