Move convertAudioAVFrameSampleFormatAndSampleRate in ffmpeg file

Author: NicolasHug

src/torchcodec/_core/FFMPEGCommon.cpp

@@ -167,6 +167,66 @@ SwrContext* createSwrContext(
   return swrContext;
 }
 
+UniqueAVFrame convertAudioAVFrameSampleFormatAndSampleRate(
+    const UniqueSwrContext& swrContext,
+    const UniqueAVFrame& srcAVFrame,
+    AVSampleFormat desiredSampleFormat,
+    int sourceSampleRate,
+    int desiredSampleRate) {
+  UniqueAVFrame convertedAVFrame(av_frame_alloc());
+  TORCH_CHECK(
+      convertedAVFrame,
+      "Could not allocate frame for sample format conversion.");
+
+  setChannelLayout(convertedAVFrame, srcAVFrame);
+  convertedAVFrame->format = static_cast<int>(desiredSampleFormat);
+  convertedAVFrame->sample_rate = desiredSampleRate;
+  if (sourceSampleRate != desiredSampleRate) {
+    // Note that this is an upper bound on the number of output samples.
+    // `swr_convert()` will likely not fill convertedAVFrame with that many
+    // samples if sample rate conversion is needed. It will buffer the last few
+    // ones because those require future samples. That's also why we reset
+    // nb_samples after the call to `swr_convert()`.
+    // We could also use `swr_get_out_samples()` to determine the number of
+    // output samples, but empirically `av_rescale_rnd()` seems to provide a
+    // tighter bound.
+    convertedAVFrame->nb_samples = av_rescale_rnd(
+        swr_get_delay(swrContext.get(), sourceSampleRate) +
+            srcAVFrame->nb_samples,
+        desiredSampleRate,
+        sourceSampleRate,
+        AV_ROUND_UP);
+  } else {
+    convertedAVFrame->nb_samples = srcAVFrame->nb_samples;
+  }
+
+  auto status = av_frame_get_buffer(convertedAVFrame.get(), 0);
+  TORCH_CHECK(
+      status == AVSUCCESS,
+      "Could not allocate frame buffers for sample format conversion: ",
+      getFFMPEGErrorStringFromErrorCode(status));
+
+  auto numConvertedSamples = swr_convert(
+      swrContext.get(),
+      convertedAVFrame->data,
+      convertedAVFrame->nb_samples,
+      static_cast<const uint8_t**>(
+          const_cast<const uint8_t**>(srcAVFrame->data)),
+      srcAVFrame->nb_samples);
+  // numConvertedSamples can be 0 if we're downsampling by a great factor and
+  // the first frame doesn't contain a lot of samples. It should be handled
+  // properly by the caller.
+  TORCH_CHECK(
+      numConvertedSamples >= 0,
+      "Error in swr_convert: ",
+      getFFMPEGErrorStringFromErrorCode(numConvertedSamples));
+
+  // See comment above about nb_samples
+  convertedAVFrame->nb_samples = numConvertedSamples;
+
+  return convertedAVFrame;
+}
+
 void setFFmpegLogLevel() {
   auto logLevel = AV_LOG_QUIET;
   const char* logLevelEnvPtr = std::getenv("TORCHCODEC_FFMPEG_LOG_LEVEL");

src/torchcodec/_core/FFMPEGCommon.h

@@ -165,6 +165,13 @@ SwrContext* createSwrContext(
     int sourceSampleRate,
     int desiredSampleRate);
 
+UniqueAVFrame convertAudioAVFrameSampleFormatAndSampleRate(
+    const UniqueSwrContext& swrContext,
+    const UniqueAVFrame& srcAVFrame,
+    AVSampleFormat desiredSampleFormat,
+    int sourceSampleRate,
+    int desiredSampleRate);
+
 // Returns true if sws_scale can handle unaligned data.
 bool canSwsScaleHandleUnalignedData();
 

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -1402,66 +1402,6 @@ void SingleStreamDecoder::convertAudioAVFrameToFrameOutputOnCPU(
   }
 }
 
-UniqueAVFrame SingleStreamDecoder::convertAudioAVFrameSampleFormatAndSampleRate(
-    const UniqueSwrContext& swrContext,
-    const UniqueAVFrame& srcAVFrame,
-    AVSampleFormat desiredSampleFormat,
-    int sourceSampleRate,
-    int desiredSampleRate) {
-  UniqueAVFrame convertedAVFrame(av_frame_alloc());
-  TORCH_CHECK(
-      convertedAVFrame,
-      "Could not allocate frame for sample format conversion.");
-
-  setChannelLayout(convertedAVFrame, srcAVFrame);
-  convertedAVFrame->format = static_cast<int>(desiredSampleFormat);
-  convertedAVFrame->sample_rate = desiredSampleRate;
-  if (sourceSampleRate != desiredSampleRate) {
-    // Note that this is an upper bound on the number of output samples.
-    // `swr_convert()` will likely not fill convertedAVFrame with that many
-    // samples if sample rate conversion is needed. It will buffer the last few
-    // ones because those require future samples. That's also why we reset
-    // nb_samples after the call to `swr_convert()`.
-    // We could also use `swr_get_out_samples()` to determine the number of
-    // output samples, but empirically `av_rescale_rnd()` seems to provide a
-    // tighter bound.
-    convertedAVFrame->nb_samples = av_rescale_rnd(
-        swr_get_delay(swrContext.get(), sourceSampleRate) +
-            srcAVFrame->nb_samples,
-        desiredSampleRate,
-        sourceSampleRate,
-        AV_ROUND_UP);
-  } else {
-    convertedAVFrame->nb_samples = srcAVFrame->nb_samples;
-  }
-
-  auto status = av_frame_get_buffer(convertedAVFrame.get(), 0);
-  TORCH_CHECK(
-      status == AVSUCCESS,
-      "Could not allocate frame buffers for sample format conversion: ",
-      getFFMPEGErrorStringFromErrorCode(status));
-
-  auto numConvertedSamples = swr_convert(
-      swrContext.get(),
-      convertedAVFrame->data,
-      convertedAVFrame->nb_samples,
-      static_cast<const uint8_t**>(
-          const_cast<const uint8_t**>(srcAVFrame->data)),
-      srcAVFrame->nb_samples);
-  // numConvertedSamples can be 0 if we're downsampling by a great factor and
-  // the first frame doesn't contain a lot of samples. It should be handled
-  // properly by the caller.
-  TORCH_CHECK(
-      numConvertedSamples >= 0,
-      "Error in swr_convert: ",
-      getFFMPEGErrorStringFromErrorCode(numConvertedSamples));
-
-  // See comment above about nb_samples
-  convertedAVFrame->nb_samples = numConvertedSamples;
-
-  return convertedAVFrame;
-}
-
 std::optional<torch::Tensor> SingleStreamDecoder::maybeFlushSwrBuffers() {
   // When sample rate conversion is involved, swresample buffers some of the
   // samples in-between calls to swr_convert (see the libswresample docs).

src/torchcodec/_core/SingleStreamDecoder.h

@@ -287,13 +287,6 @@ class SingleStreamDecoder {
       const UniqueAVFrame& avFrame,
       torch::Tensor& outputTensor);
 
-  UniqueAVFrame convertAudioAVFrameSampleFormatAndSampleRate(
-      const UniqueSwrContext& swrContext,
-      const UniqueAVFrame& srcAVFrame,
-      AVSampleFormat desiredSampleFormat,
-      int sourceSampleRate,
-      int desiredSampleRate);
-
   std::optional<torch::Tensor> maybeFlushSwrBuffers();
 
   // --------------------------------------------------------------------------