Remove preAllocatedTensor from getFrameAtIndex (#322)

Author: NicolasHug

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -1008,10 +1008,8 @@ void VideoDecoder::validateFrameIndex(
 
 VideoDecoder::DecodedOutput VideoDecoder::getFrameAtIndex(
     int streamIndex,
-    int64_t frameIndex,
-    std::optional<torch::Tensor> preAllocatedOutputTensor) {
-  auto output = getFrameAtIndexInternal(
-      streamIndex, frameIndex, preAllocatedOutputTensor);
+    int64_t frameIndex) {
+  auto output = getFrameAtIndexInternal(streamIndex, frameIndex);
   output.frame = MaybePermuteHWC2CHW(streamIndex, output.frame);
   return output;
 }

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -225,7 +225,11 @@ class VideoDecoder {
   // seconds=5.999, etc.
   DecodedOutput getFramePlayedAtTimestampNoDemux(double seconds);
 
-  DecodedOutput getFrameAtIndex(
+  DecodedOutput getFrameAtIndex(int streamIndex, int64_t frameIndex);
+  // This is morally private but needs to be exposed for C++ tests. Once
+  // getFrameAtIndex supports the preAllocatedOutputTensor parameter, we can
+  // move it back to private.
+  DecodedOutput getFrameAtIndexInternal(
       int streamIndex,
       int64_t frameIndex,
       std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
@@ -387,10 +391,6 @@ class VideoDecoder {
       DecodedOutput& output,
       std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
 
-  DecodedOutput getFrameAtIndexInternal(
-      int streamIndex,
-      int64_t frameIndex,
-      std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
   DecodedOutput getNextFrameOutputNoDemuxInternal(
       std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
 

test/decoders/VideoDecoderTest.cpp

@@ -400,7 +400,7 @@ TEST_P(VideoDecoderTest, PreAllocatedTensorFilterGraph) {
       bestVideoStreamIndex,
       VideoDecoder::VideoStreamDecoderOptions(
           "color_conversion_library=filtergraph"));
-  auto output = ourDecoder->getFrameAtIndex(
+  auto output = ourDecoder->getFrameAtIndexInternal(
       bestVideoStreamIndex, 0, preAllocatedOutputTensor);
   EXPECT_EQ(output.frame.data_ptr(), preAllocatedOutputTensor.data_ptr());
 }
@@ -418,7 +418,7 @@ TEST_P(VideoDecoderTest, PreAllocatedTensorSwscale) {
       bestVideoStreamIndex,
       VideoDecoder::VideoStreamDecoderOptions(
           "color_conversion_library=swscale"));
-  auto output = ourDecoder->getFrameAtIndex(
+  auto output = ourDecoder->getFrameAtIndexInternal(
       bestVideoStreamIndex, 0, preAllocatedOutputTensor);
   EXPECT_EQ(output.frame.data_ptr(), preAllocatedOutputTensor.data_ptr());
 }