Change return type to a 1D tensor of int64

Author: scotts

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -553,15 +553,17 @@ VideoDecoder::ContainerMetadata VideoDecoder::getContainerMetadata() const {
   return containerMetadata_;
 }
 
-std::vector<int64_t> VideoDecoder::getKeyFrameIndices(int streamIndex) {
+torch::Tensor 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));
+  const std::vector<FrameInfo>& keyFrames = streamInfos_[streamIndex].keyFrames;
+  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);
   }
 
   return keyFrameIndices;

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -100,7 +100,9 @@ class VideoDecoder {
   // Returns the metadata for the container.
   ContainerMetadata getContainerMetadata() const;
 
-  std::vector<int64_t> getKeyFrameIndices(int streamIndex);
+  // Returns the key frame indices as a tensor. The tensor is 1D and contains
+  // int64 values, where each value is the frame index for a key frame.
+  torch::Tensor getKeyFrameIndices(int streamIndex);
 
   // --------------------------------------------------------------------------
   // ADDING STREAMS API

src/torchcodec/decoders/_core/VideoDecoderOps.cpp

@@ -48,7 +48,8 @@ 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_key_frame_indices(Tensor(a!) decoder, int stream_index) -> Tensor");
   m.def("get_json_metadata(Tensor(a!) decoder) -> str");
   m.def("get_container_json_metadata(Tensor(a!) decoder) -> str");
   m.def(
@@ -335,7 +336,7 @@ bool _test_frame_pts_equality(
       videoDecoder->getPtsSecondsForFrame(stream_index, frame_index);
 }
 
-std::vector<int64_t> _get_key_frame_indices(
+torch::Tensor _get_key_frame_indices(
     at::Tensor& decoder,
     int64_t stream_index) {
   auto videoDecoder = unwrapTensorToGetDecoder(decoder);

src/torchcodec/decoders/_core/VideoDecoderOps.h

@@ -137,9 +137,7 @@ 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);
+torch::Tensor _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/_video_decoder.py

@@ -186,7 +186,9 @@ def __getitem__(self, key: Union[numbers.Integral, slice]) -> Tensor:
         )
 
     def _get_key_frame_indices(self) -> list[int]:
-        return core._get_key_frame_indices(self._decoder, stream_index=self.stream_index)
+        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

@@ -835,7 +835,9 @@ 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()
-        assert len(key_frame_indices) > 0
+        size = key_frame_indices.size()
+        assert size[0] > 0
+        assert len(size) == 1
 
 
 if __name__ == "__main__":