Rename 'wf' to 'samples' in AudioEncoder

Author: danielflores3

src/torchcodec/_core/Encoder.cpp

@@ -8,16 +8,16 @@ namespace facebook::torchcodec {
 
 namespace {
 
-torch::Tensor validateWf(torch::Tensor wf) {
+torch::Tensor validateSamples(torch::Tensor samples) {
   TORCH_CHECK(
-      wf.dtype() == torch::kFloat32,
-      "waveform must have float32 dtype, got ",
-      wf.dtype());
-  TORCH_CHECK(wf.dim() == 2, "waveform must have 2 dimensions, got ", wf.dim());
+      samples.dtype() == torch::kFloat32,
+      "samples must have float32 dtype, got ",
+      samples.dtype());
+  TORCH_CHECK(samples.dim() == 2, "samples must have 2 dimensions, got ", samples.dim());
 
   // We enforce this, but if we get user reports we should investigate whether
   // that's actually needed.
-  int numChannels = static_cast<int>(wf.sizes()[0]);
+  int numChannels = static_cast<int>(samples.sizes()[0]);
   TORCH_CHECK(
       numChannels <= AV_NUM_DATA_POINTERS,
       "Trying to encode ",
@@ -26,7 +26,7 @@ torch::Tensor validateWf(torch::Tensor wf) {
       AV_NUM_DATA_POINTERS,
       " channels per frame.");
 
-  return wf.contiguous();
+  return samples.contiguous();
 }
 
 void validateSampleRate(const AVCodec& avCodec, int sampleRate) {
@@ -71,7 +71,7 @@ static const std::vector<AVSampleFormat> preferredFormatsOrder = {
 
 AVSampleFormat findBestOutputSampleFormat(const AVCodec& avCodec) {
   // Find a sample format that the encoder supports. We prefer using FLT[P],
-  // since this is the format of the input waveform. If FLTP isn't supported
+  // since this is the format of the input samples. If FLTP isn't supported
   // then we'll need to convert the AVFrame's format. Our heuristic is to encode
   // into the format with the highest resolution.
   if (avCodec.sample_fmts == nullptr) {
@@ -98,12 +98,12 @@ AVSampleFormat findBestOutputSampleFormat(const AVCodec& avCodec) {
 AudioEncoder::~AudioEncoder() {}
 
 AudioEncoder::AudioEncoder(
-    const torch::Tensor wf,
+    const torch::Tensor samples,
     int sampleRate,
     std::string_view fileName,
     std::optional<int64_t> bitRate,
     std::optional<int64_t> numChannels)
-    : wf_(validateWf(wf)) {
+    : samples_(validateSamples(samples)) {
   setFFmpegLogLevel();
   AVFormatContext* avFormatContext = nullptr;
   int status = avformat_alloc_output_context2(
@@ -130,13 +130,13 @@ AudioEncoder::AudioEncoder(
 }
 
 AudioEncoder::AudioEncoder(
-    const torch::Tensor wf,
+    const torch::Tensor samples,
     int sampleRate,
     std::string_view formatName,
     std::unique_ptr<AVIOToTensorContext> avioContextHolder,
     std::optional<int64_t> bitRate,
     std::optional<int64_t> numChannels)
-    : wf_(validateWf(wf)), avioContextHolder_(std::move(avioContextHolder)) {
+    : samples_(validateSamples(samples)), avioContextHolder_(std::move(avioContextHolder)) {
   setFFmpegLogLevel();
   AVFormatContext* avFormatContext = nullptr;
   int status = avformat_alloc_output_context2(
@@ -177,7 +177,7 @@ void AudioEncoder::initializeEncoder(
   // well when "-b:a" isn't specified.
   avCodecContext_->bit_rate = bitRate.value_or(0);
 
-  desiredNumChannels_ = static_cast<int>(numChannels.value_or(wf_.sizes()[0]));
+  desiredNumChannels_ = static_cast<int>(numChannels.value_or(samples_.sizes()[0]));
   validateNumChannels(*avCodec, desiredNumChannels_);
   // The avCodecContext layout defines the layout of the encoded output, it's
   // not related to the input sampes.
@@ -186,11 +186,13 @@ void AudioEncoder::initializeEncoder(
   validateSampleRate(*avCodec, sampleRate);
   avCodecContext_->sample_rate = sampleRate;
 
-  // 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, which is
+  // Input samples are expected to be FLTP. Not all encoders support FLTP, so we
+  // may need to convert the samples into a supported output sample format, which is
   // what the `.sample_fmt` defines.
   avCodecContext_->sample_fmt = findBestOutputSampleFormat(*avCodec);
 
+  setDefaultChannelLayout(avCodecContext_, static_cast<int>(samples_.sizes()[0]));
+
   int status = avcodec_open2(avCodecContext_.get(), avCodec, nullptr);
   TORCH_CHECK(
       status == AVSUCCESS,
@@ -237,7 +239,7 @@ void AudioEncoder::encode() {
   avFrame->pts = 0;
   // We set the channel layout of the frame to the default layout corresponding
   // to the input samples' number of channels
-  setDefaultChannelLayout(avFrame, static_cast<int>(wf_.sizes()[0]));
+  setDefaultChannelLayout(avFrame, static_cast<int>(samples_.sizes()[0]));
 
   auto status = av_frame_get_buffer(avFrame.get(), 0);
   TORCH_CHECK(
@@ -247,10 +249,10 @@ void AudioEncoder::encode() {
 
   AutoAVPacket autoAVPacket;
 
-  uint8_t* pwf = static_cast<uint8_t*>(wf_.data_ptr());
-  int numSamples = static_cast<int>(wf_.sizes()[1]); // per channel
+  uint8_t* psamples = static_cast<uint8_t*>(samples_.data_ptr());
+  int numSamples = static_cast<int>(samples_.sizes()[1]); // per channel
   int numEncodedSamples = 0; // per channel
-  int numBytesPerSample = static_cast<int>(wf_.element_size());
+  int numBytesPerSample = static_cast<int>(samples_.element_size());
   int numBytesPerChannel = numSamples * numBytesPerSample;
 
   status = avformat_write_header(avFormatContext_.get(), nullptr);
@@ -270,11 +272,11 @@ void AudioEncoder::encode() {
         std::min(numSamplesAllocatedPerFrame, numSamples - numEncodedSamples);
     int numBytesToEncode = numSamplesToEncode * numBytesPerSample;
 
-    for (int ch = 0; ch < wf_.sizes()[0]; ch++) {
+    for (int ch = 0; ch < samples_.sizes()[0]; ch++) {
       std::memcpy(
-          avFrame->data[ch], pwf + ch * numBytesPerChannel, numBytesToEncode);
+          avFrame->data[ch], psamples + ch * numBytesPerChannel, numBytesToEncode);
     }
-    pwf += numBytesToEncode;
+    psamples += numBytesToEncode;
 
     // Above, we set the AVFrame's .nb_samples to AVCodecContext.frame_size so
     // that the frame buffers are allocated to a big enough size. Here, we reset

src/torchcodec/_core/Encoder.h

@@ -17,17 +17,17 @@ class AudioEncoder {
   // TODO-ENCODING: bundle the optional params like bitRate, numChannels, etc.
   // into an AudioStreamOptions struct, or similar.
   AudioEncoder(
-      const torch::Tensor wf,
+      const torch::Tensor samples,
       // 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
+      // should be. Particularly, the sample rate of the input samples 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);
   AudioEncoder(
-      const torch::Tensor wf,
+      const torch::Tensor samples,
       int sampleRate,
       std::string_view formatName,
       std::unique_ptr<AVIOToTensorContext> avioContextHolder,
@@ -54,7 +54,7 @@ class AudioEncoder {
   // see other TODO above.
   int desiredNumChannels_ = -1;
 
-  const torch::Tensor wf_;
+  const torch::Tensor samples_;
 
   // Stores the AVIOContext for the output tensor buffer.
   std::unique_ptr<AVIOToTensorContext> avioContextHolder_;

src/torchcodec/_core/custom_ops.cpp

@@ -29,9 +29,9 @@ TORCH_LIBRARY(torchcodec_ns, m) {
       "torchcodec._core.ops", "//pytorch/torchcodec:torchcodec");
   m.def("create_from_file(str filename, str? seek_mode=None) -> Tensor");
   m.def(
-      "encode_audio_to_file(Tensor wf, int sample_rate, str filename, int? bit_rate=None, int? num_channels=None) -> ()");
+      "encode_audio_to_file(Tensor samples, int sample_rate, str filename, int? bit_rate=None, int? num_channels=None) -> ()");
   m.def(
-      "encode_audio_to_tensor(Tensor wf, int sample_rate, str format, int? bit_rate=None, int? num_channels=None) -> Tensor");
+      "encode_audio_to_tensor(Tensor samples, int sample_rate, str format, int? bit_rate=None, int? num_channels=None) -> Tensor");
   m.def(
       "create_from_tensor(Tensor video_tensor, str? seek_mode=None) -> Tensor");
   m.def("_convert_to_tensor(int decoder_ptr) -> Tensor");
@@ -388,25 +388,25 @@ OpsAudioFramesOutput get_frames_by_pts_in_range_audio(
 }
 
 void encode_audio_to_file(
-    const at::Tensor wf,
+    const at::Tensor samples,
     int64_t sample_rate,
     std::string_view file_name,
     std::optional<int64_t> bit_rate = std::nullopt,
     std::optional<int64_t> num_channels = std::nullopt) {
   AudioEncoder(
-      wf, validateSampleRate(sample_rate), file_name, bit_rate, num_channels)
+      samples, validateSampleRate(sample_rate), file_name, bit_rate, num_channels)
       .encode();
 }
 
 at::Tensor encode_audio_to_tensor(
-    const at::Tensor wf,
+    const at::Tensor samples,
     int64_t sample_rate,
     std::string_view format,
     std::optional<int64_t> bit_rate = std::nullopt,
     std::optional<int64_t> num_channels = std::nullopt) {
   auto avioContextHolder = std::make_unique<AVIOToTensorContext>();
   return AudioEncoder(
-             wf,
+             samples,
              validateSampleRate(sample_rate),
              format,
              std::move(avioContextHolder),