Merge branch 'use_audioStreamOptions' into encoding_sample_rate_lezzzgo

Author: NicolasHug

src/torchcodec/_core/Encoder.cpp

@@ -101,10 +101,8 @@ AudioEncoder::AudioEncoder(
     const torch::Tensor wf,
     int sampleRate,
     std::string_view fileName,
-    std::optional<int64_t> bitRate,
-    std::optional<int64_t> numChannels,
-    std::optional<int64_t> desiredSampleRate)
-    : wf_(validateWf(wf)), sampleRateInput_(static_cast<int>(sampleRate)) {
+    const AudioStreamOptions& audioStreamOptions)
+    : wf_(validateWf(wf)), sampleRateInput_(sampleRate) {
   setFFmpegLogLevel();
   AVFormatContext* avFormatContext = nullptr;
   int status = avformat_alloc_output_context2(
@@ -127,19 +125,17 @@ AudioEncoder::AudioEncoder(
       ", make sure it's a valid path? ",
       getFFMPEGErrorStringFromErrorCode(status));
 
-  initializeEncoder(bitRate, numChannels, desiredSampleRate);
+  initializeEncoder(audioStreamOptions);
 }
 
 AudioEncoder::AudioEncoder(
     const torch::Tensor wf,
     int sampleRate,
     std::string_view formatName,
     std::unique_ptr<AVIOToTensorContext> avioContextHolder,
-    std::optional<int64_t> bitRate,
-    std::optional<int64_t> numChannels,
-    std::optional<int64_t> desiredSampleRate)
+    const AudioStreamOptions& audioStreamOptions)
     : wf_(validateWf(wf)),
-      sampleRateInput_(static_cast<int>(sampleRate)),
+      sampleRateInput_(sampleRate),
       avioContextHolder_(std::move(avioContextHolder)) {
   setFFmpegLogLevel();
   AVFormatContext* avFormatContext = nullptr;
@@ -157,13 +153,11 @@ AudioEncoder::AudioEncoder(
 
   avFormatContext_->pb = avioContextHolder_->getAVIOContext();
 
-  initializeEncoder(bitRate, numChannels, desiredSampleRate);
+  initializeEncoder(audioStreamOptions);
 }
 
 void AudioEncoder::initializeEncoder(
-    std::optional<int64_t> bitRate,
-    std::optional<int64_t> numChannels,
-    std::optional<int64_t> desiredSampleRate) {
+    const AudioStreamOptions& audioStreamOptions) {
   // We use the AVFormatContext's default codec for that
   // specific format/container.
   const AVCodec* avCodec =
@@ -174,23 +168,26 @@ void AudioEncoder::initializeEncoder(
   TORCH_CHECK(avCodecContext != nullptr, "Couldn't allocate codec context.");
   avCodecContext_.reset(avCodecContext);
 
-  if (bitRate.has_value()) {
-    TORCH_CHECK(*bitRate >= 0, "bit_rate=", *bitRate, " must be >= 0.");
+  auto desiredBitRate = audioStreamOptions.bitRate;
+  if (desiredBitRate.has_value()) {
+    TORCH_CHECK(
+        *desiredBitRate >= 0, "bit_rate=", *desiredBitRate, " must be >= 0.");
   }
-  // bit_rate=None defaults to 0, which is what the FFmpeg CLI seems to use
-  // as well when "-b:a" isn't specified.
-  avCodecContext_->bit_rate = bitRate.value_or(0);
+  // bit_rate=None defaults to 0, which is what the FFmpeg CLI seems to use as
+  // well when "-b:a" isn't specified.
+  avCodecContext_->bit_rate = desiredBitRate.value_or(0);
 
-  numChannelsOutput_ = static_cast<int>(numChannels.value_or(wf_.sizes()[0]));
-  validateNumChannels(*avCodec, numChannelsOutput_);
+  outNumChannels_ =
+      static_cast<int>(audioStreamOptions.numChannels.value_or(wf_.sizes()[0]));
+  validateNumChannels(*avCodec, outNumChannels_);
   // The avCodecContext layout defines the layout of the encoded output, it's
   // not related to the input sampes.
-  setDefaultChannelLayout(avCodecContext_, numChannelsOutput_);
+  setDefaultChannelLayout(avCodecContext_, outNumChannels_);
 
-  sampleRateOutput_ =
-      static_cast<int>(desiredSampleRate.value_or(sampleRateInput_));
-  validateSampleRate(*avCodec, sampleRateOutput_);
-  avCodecContext_->sample_rate = sampleRateOutput_;
+  outSampleRate_ = static_cast<int>(
+      audioStreamOptions.sampleRate.value_or(sampleRateInput_));
+  validateSampleRate(*avCodec, outSampleRate_);
+  avCodecContext_->sample_rate = outSampleRate_;
 
   // Input waveform is expected to be FLTP. Not all encoders support FLTP,
   // so we may need to convert the wf into a supported output sample format,
@@ -310,8 +307,8 @@ void AudioEncoder::encodeInnerLoop(
   bool mustConvert =
       (srcAVFrame != nullptr &&
        (avCodecContext_->sample_fmt != AV_SAMPLE_FMT_FLTP ||
-        getNumChannels(srcAVFrame) != numChannelsOutput_ ||
-        srcAVFrame->sample_rate != sampleRateOutput_));
+        getNumChannels(srcAVFrame) != outNumChannels_ ||
+        srcAVFrame->sample_rate != outSampleRate_));
 
   UniqueAVFrame convertedAVFrame;
   if (mustConvert) {
@@ -320,17 +317,17 @@ void AudioEncoder::encodeInnerLoop(
           AV_SAMPLE_FMT_FLTP,
           avCodecContext_->sample_fmt,
           srcAVFrame->sample_rate,
-          sampleRateOutput_,
+          outSampleRate_,
           srcAVFrame,
-          numChannelsOutput_));
+          outNumChannels_));
     }
     convertedAVFrame = convertAudioAVFrameSamples(
         swrContext_,
         srcAVFrame,
         avCodecContext_->sample_fmt,
-        sampleRateOutput_,
-        numChannelsOutput_);
-    if (sampleRateOutput_ == sampleRateInput_) {
+        outSampleRate_,
+        outNumChannels_);
+    if (outSampleRate_ == sampleRateInput_) {
       TORCH_CHECK(
           convertedAVFrame->nb_samples == srcAVFrame->nb_samples,
           "convertedAVFrame->nb_samples=",

src/torchcodec/_core/Encoder.h

@@ -2,6 +2,7 @@
 #include <torch/types.h>
 #include "src/torchcodec/_core/AVIOBytesContext.h"
 #include "src/torchcodec/_core/FFMPEGCommon.h"
+#include "src/torchcodec/_core/StreamOptions.h"
 
 namespace facebook::torchcodec {
 class AudioEncoder {
@@ -13,36 +14,28 @@ class AudioEncoder {
   // like passing 0, which results in choosing the minimum supported bit rate.
   // Passing 44_100 could result in output being 44000 if only 44000 is
   // supported.
-  //
-  // TODO-ENCODING: bundle the optional params like bitRate, numChannels, etc.
-  // into an AudioStreamOptions struct, or similar.
   AudioEncoder(
       const torch::Tensor wf,
+      // TODO-ENCODING: update this comment when we support an output sample
+      // rate. This will become the input sample rate.
       // The *output* sample rate. We can't really decide for the user what it
       // should be. Particularly, the sample rate of the input waveform should
       // match this, and that's up to the user. If sample rates don't match,
       // encoding will still work but audio will be distorted.
       int sampleRate,
       std::string_view fileName,
-      std::optional<int64_t> bitRate = std::nullopt,
-      std::optional<int64_t> numChannels = std::nullopt,
-      std::optional<int64_t> desiredSampleRate = std::nullopt);
+      const AudioStreamOptions& audioStreamOptions);
   AudioEncoder(
       const torch::Tensor wf,
       int sampleRate,
       std::string_view formatName,
       std::unique_ptr<AVIOToTensorContext> avioContextHolder,
-      std::optional<int64_t> bitRate = std::nullopt,
-      std::optional<int64_t> numChannels = std::nullopt,
-      std::optional<int64_t> desiredSampleRate = std::nullopt);
+      const AudioStreamOptions& audioStreamOptions);
   void encode();
   torch::Tensor encodeToTensor();
 
  private:
-  void initializeEncoder(
-      std::optional<int64_t> bitRate = std::nullopt,
-      std::optional<int64_t> numChannels = std::nullopt,
-      std::optional<int64_t> desiredSampleRate = std::nullopt);
+  void initializeEncoder(const AudioStreamOptions& audioStreamOptions);
   void encodeInnerLoop(
       AutoAVPacket& autoAVPacket,
       const UniqueAVFrame& srcAVFrame);
@@ -52,10 +45,10 @@ class AudioEncoder {
   UniqueAVCodecContext avCodecContext_;
   int streamIndex_;
   UniqueSwrContext swrContext_;
-  // TODO-ENCODING: These should just be part of an options struct,
-  // see other TODO above.
-  int numChannelsOutput_ = -1;
-  int sampleRateOutput_ = -1;
+  AudioStreamOptions audioStreamOptions;
+
+  int outNumChannels_ = -1;
+  int outSampleRate_ = -1;
 
   const torch::Tensor wf_;
   int sampleRateInput_ = -1;

src/torchcodec/_core/FFMPEGCommon.cpp

@@ -159,74 +159,74 @@ namespace {
 #if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
 
 // Returns:
-// - the srcAVFrame's channel layout if srcAVFrame has desiredNumChannels
-// - the default channel layout with desiredNumChannels otherwise.
-AVChannelLayout getDesiredChannelLayout(
-    int desiredNumChannels,
+// - the srcAVFrame's channel layout if srcAVFrame has outNumChannels
+// - the default channel layout with outNumChannels otherwise.
+AVChannelLayout getOutputChannelLayout(
+    int outNumChannels,
     const UniqueAVFrame& srcAVFrame) {
-  AVChannelLayout desiredLayout;
-  if (desiredNumChannels == getNumChannels(srcAVFrame)) {
-    desiredLayout = srcAVFrame->ch_layout;
+  AVChannelLayout outLayout;
+  if (outNumChannels == getNumChannels(srcAVFrame)) {
+    outLayout = srcAVFrame->ch_layout;
   } else {
-    av_channel_layout_default(&desiredLayout, desiredNumChannels);
+    av_channel_layout_default(&outLayout, outNumChannels);
   }
-  return desiredLayout;
+  return outLayout;
 }
 
 #else
 
 // Same as above
-int64_t getDesiredChannelLayout(
-    int desiredNumChannels,
+int64_t getOutputChannelLayout(
+    int outNumChannels,
     const UniqueAVFrame& srcAVFrame) {
-  int64_t desiredLayout;
-  if (desiredNumChannels == getNumChannels(srcAVFrame)) {
-    desiredLayout = srcAVFrame->channel_layout;
+  int64_t outLayout;
+  if (outNumChannels == getNumChannels(srcAVFrame)) {
+    outLayout = srcAVFrame->channel_layout;
   } else {
-    desiredLayout = av_get_default_channel_layout(desiredNumChannels);
+    outLayout = av_get_default_channel_layout(outNumChannels);
   }
-  return desiredLayout;
+  return outLayout;
 }
 #endif
 } // namespace
 
-// Sets dstAVFrame' channel layout to getDesiredChannelLayout(): see doc above
+// Sets dstAVFrame' channel layout to getOutputChannelLayout(): see doc above
 void setChannelLayout(
     UniqueAVFrame& dstAVFrame,
     const UniqueAVFrame& srcAVFrame,
-    int desiredNumChannels) {
+    int outNumChannels) {
 #if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
-  AVChannelLayout desiredLayout =
-      getDesiredChannelLayout(desiredNumChannels, srcAVFrame);
-  auto status = av_channel_layout_copy(&dstAVFrame->ch_layout, &desiredLayout);
+  AVChannelLayout outLayout =
+      getOutputChannelLayout(outNumChannels, srcAVFrame);
+  auto status = av_channel_layout_copy(&dstAVFrame->ch_layout, &outLayout);
   TORCH_CHECK(
       status == AVSUCCESS,
       "Couldn't copy channel layout to avFrame: ",
       getFFMPEGErrorStringFromErrorCode(status));
 #else
   dstAVFrame->channel_layout =
-      getDesiredChannelLayout(desiredNumChannels, srcAVFrame);
-  dstAVFrame->channels = desiredNumChannels;
+      getOutputChannelLayout(outNumChannels, srcAVFrame);
+  dstAVFrame->channels = outNumChannels;
 #endif
 }
 
 SwrContext* createSwrContext(
     AVSampleFormat srcSampleFormat,
-    AVSampleFormat desiredSampleFormat,
+    AVSampleFormat outSampleFormat,
     int srcSampleRate,
-    int desiredSampleRate,
+    int outSampleRate,
     const UniqueAVFrame& srcAVFrame,
-    int desiredNumChannels) {
+    int outNumChannels) {
   SwrContext* swrContext = nullptr;
   int status = AVSUCCESS;
 #if LIBAVFILTER_VERSION_MAJOR > 7 // FFmpeg > 4
-  AVChannelLayout desiredLayout =
-      getDesiredChannelLayout(desiredNumChannels, srcAVFrame);
+  AVChannelLayout outLayout =
+      getOutputChannelLayout(outNumChannels, srcAVFrame);
   status = swr_alloc_set_opts2(
       &swrContext,
-      &desiredLayout,
-      desiredSampleFormat,
-      desiredSampleRate,
+      &outLayout,
+      outSampleFormat,
+      outSampleRate,
       &srcAVFrame->ch_layout,
       srcSampleFormat,
       srcSampleRate,
@@ -238,13 +238,12 @@ SwrContext* createSwrContext(
       "Couldn't create SwrContext: ",
       getFFMPEGErrorStringFromErrorCode(status));
 #else
-  int64_t desiredLayout =
-      getDesiredChannelLayout(desiredNumChannels, srcAVFrame);
+  int64_t outLayout = getOutputChannelLayout(outNumChannels, srcAVFrame);
   swrContext = swr_alloc_set_opts(
       nullptr,
-      desiredLayout,
-      desiredSampleFormat,
-      desiredSampleRate,
+      outLayout,
+      outSampleFormat,
+      outSampleRate,
       srcAVFrame->channel_layout,
       srcSampleFormat,
       srcSampleRate,
@@ -267,19 +266,19 @@ SwrContext* createSwrContext(
 UniqueAVFrame convertAudioAVFrameSamples(
     const UniqueSwrContext& swrContext,
     const UniqueAVFrame& srcAVFrame,
-    AVSampleFormat desiredSampleFormat,
-    int desiredSampleRate,
-    int desiredNumChannels) {
+    AVSampleFormat outSampleFormat,
+    int outSampleRate,
+    int outNumChannels) {
   UniqueAVFrame convertedAVFrame(av_frame_alloc());
   TORCH_CHECK(
       convertedAVFrame,
       "Could not allocate frame for sample format conversion.");
 
-  convertedAVFrame->format = static_cast<int>(desiredSampleFormat);
+  convertedAVFrame->format = static_cast<int>(outSampleFormat);
 
-  convertedAVFrame->sample_rate = desiredSampleRate;
+  convertedAVFrame->sample_rate = outSampleRate;
   int srcSampleRate = srcAVFrame->sample_rate;
-  if (srcSampleRate != desiredSampleRate) {
+  if (srcSampleRate != outSampleRate) {
     // 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
@@ -290,14 +289,14 @@ UniqueAVFrame convertAudioAVFrameSamples(
     // tighter bound.
     convertedAVFrame->nb_samples = av_rescale_rnd(
         swr_get_delay(swrContext.get(), srcSampleRate) + srcAVFrame->nb_samples,
-        desiredSampleRate,
+        outSampleRate,
         srcSampleRate,
         AV_ROUND_UP);
   } else {
     convertedAVFrame->nb_samples = srcAVFrame->nb_samples;
   }
 
-  setChannelLayout(convertedAVFrame, srcAVFrame, desiredNumChannels);
+  setChannelLayout(convertedAVFrame, srcAVFrame, outNumChannels);
 
   auto status = av_frame_get_buffer(convertedAVFrame.get(), 0);
   TORCH_CHECK(