WIP

Author: NicolasHug

src/torchcodec/decoders/_core/CMakeLists.txt

@@ -4,7 +4,8 @@ set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
 find_package(Torch REQUIRED)
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -pedantic -Werror ${TORCH_CXX_FLAGS}")
+# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -pedantic -Werror ${TORCH_CXX_FLAGS}")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra ${TORCH_CXX_FLAGS}")
 find_package(Python3 ${PYTHON_VERSION} EXACT COMPONENTS Development)
 
 function(make_torchcodec_library library_name ffmpeg_target)
@@ -97,6 +98,7 @@ else()
         libavformat
         libavcodec
         libavutil
+        libswresample
         libswscale
     )
 

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -23,6 +23,7 @@ extern "C" {
 #include <libavutil/imgutils.h>
 #include <libavutil/log.h>
 #include <libavutil/pixdesc.h>
+#include <libswresample/swresample.h>
 #include <libswscale/swscale.h>
 }
 
@@ -541,14 +542,18 @@ void VideoDecoder::addVideoStream(
       videoStreamOptions.colorConversionLibrary.value_or(defaultLibrary);
 }
 
-void VideoDecoder::addAudioStream(int streamIndex) {
+void VideoDecoder::addAudioStream(
+    int streamIndex,
+    const AudioStreamOptions& audioStreamOptions) {
   TORCH_CHECK(
       seekMode_ == SeekMode::approximate,
       "seek_mode must be 'approximate' for audio streams.");
 
   addStream(streamIndex, AVMEDIA_TYPE_AUDIO);
 
   auto& streamInfo = streamInfos_[activeStreamIndex_];
+  streamInfo.audioStreamOptions = audioStreamOptions;
+
   auto& streamMetadata =
       containerMetadata_.allStreamMetadata[activeStreamIndex_];
   streamMetadata.sampleRate =
@@ -1332,6 +1337,82 @@ void VideoDecoder::convertAudioAVFrameToFrameOutputOnCPU(
       "pre-allocated audio tensor not supported yet.");
 
   const AVFrame* avFrame = avFrameStream.avFrame.get();
+  AVFrame* output_frame = nullptr;
+  SwrContext* swr_ctx = NULL; // TODO RAII
+
+  const auto sampleRate =
+      streamInfos_[activeStreamIndex_].audioStreamOptions.sampleRate;
+  if (sampleRate.has_value()) {
+    int outRate = static_cast<int>(*sampleRate);
+    auto& streamMetadata =
+        containerMetadata_.allStreamMetadata[activeStreamIndex_];
+    int inRate = static_cast<int>(streamMetadata.sampleRate.value());
+
+    printf("RESAMPLEING FROM %d to %d\n", outRate, inRate);
+    AVSampleFormat sampleFormat = AV_SAMPLE_FMT_FLTP;
+
+    AVChannelLayout stereoLayout = AV_CHANNEL_LAYOUT_STEREO;
+    const AVChannelLayout* chl = &stereoLayout;
+
+    int status = swr_alloc_set_opts2(
+        &swr_ctx,
+        chl,
+        sampleFormat,
+        outRate,
+        chl,
+        sampleFormat,
+        inRate,
+        0,
+        NULL);
+
+    TORCH_CHECK(status == 0, "IS NULL");
+
+    if (swr_init(swr_ctx) < 0) {
+      swr_free(&swr_ctx);
+      TORCH_CHECK(false, "Failed to initialize the resampling context\n");
+    }
+
+    // Allocate output frame
+    output_frame = av_frame_alloc();
+    if (!output_frame) {
+      swr_free(&swr_ctx);
+      TORCH_CHECK(false, "Could not allocate output frame\n");
+    }
+    output_frame->ch_layout = stereoLayout;
+    output_frame->sample_rate = outRate;
+    output_frame->format = sampleFormat;
+
+    output_frame->nb_samples = av_rescale_rnd(
+        swr_get_delay(swr_ctx, inRate) + avFrame->nb_samples,
+        outRate,
+        inRate,
+        AV_ROUND_UP);
+
+    if (av_frame_get_buffer(output_frame, 0) < 0) {
+      av_frame_free(&output_frame);
+      swr_free(&swr_ctx);
+      TORCH_CHECK(false, "Could not allocate output frame samples");
+    }
+
+    int ret = swr_convert(
+        swr_ctx,
+        output_frame->data,
+        output_frame->nb_samples,
+        (const uint8_t**)avFrame->data,
+        avFrame->nb_samples);
+    if (ret < 0) {
+      av_frame_free(&output_frame);
+      swr_free(&swr_ctx);
+      TORCH_CHECK(false, "Error while converting\n");
+    }
+
+    printf(
+        "nb_samples: %d %d\n", avFrame->nb_samples, output_frame->nb_samples);
+
+    avFrame = output_frame; // lmao
+  } else {
+    printf("NO RESAMPLING\n");
+  }
 
   auto numSamples = avFrame->nb_samples; // per channel
   auto numChannels = getNumChannels(avFrame);
@@ -1360,6 +1441,10 @@ void VideoDecoder::convertAudioAVFrameToFrameOutputOnCPU(
           av_get_sample_fmt_name(format));
   }
   frameOutput.data = outputData;
+
+  // TODO
+  av_frame_free(&output_frame);
+  swr_free(&swr_ctx);
 }
 
 // --------------------------------------------------------------------------

src/torchcodec/decoders/_core/VideoDecoder.h

@@ -138,10 +138,19 @@ class VideoDecoder {
     torch::Device device = torch::kCPU;
   };
 
+  struct AudioStreamOptions {
+    AudioStreamOptions() {}
+
+    // explicit AudioStreamOptions(const std::string& optionsString);
+    std::optional<int> sampleRate;
+  };
+
   void addVideoStream(
       int streamIndex,
       const VideoStreamOptions& videoStreamOptions = VideoStreamOptions());
-  void addAudioStream(int streamIndex);
+  void addAudioStream(
+      int streamIndex,
+      const AudioStreamOptions& audioStreamOptions = AudioStreamOptions());
 
   // --------------------------------------------------------------------------
   // DECODING AND SEEKING APIs
@@ -221,7 +230,6 @@ class VideoDecoder {
       double startSeconds,
       double stopSeconds);
 
-  // TODO-AUDIO: Should accept sampleRate
   torch::Tensor getFramesPlayedInRangeAudio(
       double startSeconds,
       std::optional<double> stopSecondsOptional = std::nullopt);
@@ -343,6 +351,7 @@ class VideoDecoder {
     int64_t lastDecodedAvFramePts = 0;
     int64_t lastDecodedAvFrameDuration = 0;
     VideoStreamOptions videoStreamOptions;
+    AudioStreamOptions audioStreamOptions;
 
     // color-conversion fields. Only one of FilterGraphContext and
     // UniqueSwsContext should be non-null.

src/torchcodec/decoders/_core/VideoDecoderOps.cpp

@@ -34,7 +34,7 @@ TORCH_LIBRARY(torchcodec_ns, m) {
   m.def(
       "add_video_stream(Tensor(a!) decoder, *, int? width=None, int? height=None, int? num_threads=None, str? dimension_order=None, int? stream_index=None, str? device=None) -> ()");
   m.def(
-      "add_audio_stream(Tensor(a!) decoder, *, int? stream_index=None) -> ()");
+      "add_audio_stream(Tensor(a!) decoder, *, int? stream_index=None, int? sample_rate=None) -> ()");
   m.def("seek_to_pts(Tensor(a!) decoder, float seconds) -> ()");
   m.def("get_next_frame(Tensor(a!) decoder) -> (Tensor, Tensor, Tensor)");
   m.def(
@@ -213,9 +213,13 @@ void _add_video_stream(
 
 void add_audio_stream(
     at::Tensor& decoder,
-    std::optional<int64_t> stream_index) {
+    std::optional<int64_t> stream_index,
+    std::optional<int64_t> sample_rate) {
+  VideoDecoder::AudioStreamOptions audioStreamOptions;
+  audioStreamOptions.sampleRate = sample_rate;
+
   auto videoDecoder = unwrapTensorToGetDecoder(decoder);
-  videoDecoder->addAudioStream(stream_index.value_or(-1));
+  videoDecoder->addAudioStream(stream_index.value_or(-1), audioStreamOptions);
 }
 
 void seek_to_pts(at::Tensor& decoder, double seconds) {

src/torchcodec/decoders/_core/VideoDecoderOps.h

@@ -50,7 +50,8 @@ void _add_video_stream(
 
 void add_audio_stream(
     at::Tensor& decoder,
-    std::optional<int64_t> stream_index = std::nullopt);
+    std::optional<int64_t> stream_index = std::nullopt,
+    std::optional<int64_t> sample_rate = std::nullopt);
 
 // Seek to a particular presentation timestamp in the video in seconds.
 void seek_to_pts(at::Tensor& decoder, double seconds);