Add AudioConverter to the Blocks API (#1667)

Author: NicolasHug

examples/decoding/blocks.py

@@ -5,9 +5,9 @@
 # LICENSE file in the root directory of this source tree.
 
 """
-===============================================
+========================================
 Blocks: build your own decoding pipeline
-===============================================
+========================================
 
 .. warning::
 
@@ -22,13 +22,16 @@
 .. code-block::
 
    VideoDemuxer  ->  VideoPacketDecoder  ->  ColorConverter
-    Packet         RawFrame          RGB Frame
+     Packet            RawFrame               RGB Frame
 
 The blocks are passive: they never create threads, and they release the GIL.
 You decide how they are composed, on which threads, and where to stop. Below
 we illustrate a few things this enables: overlapping stages on multiple
 threads, accessing raw (YUV) frames, and decoding streams of unknown -
 possibly infinite - length.
+
+Audio works the same way, through ``AudioDemuxer``, ``AudioPacketDecoder`` and
+``AudioConverter``; we come back to it at the end.
 """
 
 # %%
@@ -333,7 +336,7 @@ def upsample(plane):
 
 # %%
 # Raw HDR frames
-# ~~~~~~~~~~~~~~
+# --------------
 #
 # Raw planes come at the source's own precision, so a 10-bit HDR video gives
 # ``uint16`` planes with all 10 bits intact - no clipping to 8 bits, and no
@@ -427,3 +430,76 @@ def start_live_stream():
 
 ffmpeg.kill()
 ffmpeg.wait()
+
+# %%
+# Audio
+# -----
+#
+# Audio has the same three stages:
+#
+# .. code-block::
+#
+#    AudioDemuxer  ->  AudioPacketDecoder  ->  AudioConverter
+#      Packet           RawAudioSamples          AudioSamples
+#
+# Decoding is the same operation either way, so the two packet decoders are a
+# single class in C++; they are separate in Python because what they hand out
+# isn't. Audio comes out as ``RawAudioSamples``: the codec's own samples, in
+# the codec's own sample type, as a ``[num_channels, num_samples]`` tensor.
+from torchcodec.decoders._blocks import (
+    AudioConverter,
+    AudioDemuxer,
+    AudioPacketDecoder,
+)
+
+audio_path = temp_dir / "audio.wav"
+subprocess.run(
+    [
+        "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
+        "-f", "lavfi", "-i", "sine=frequency=440:sample_rate=44100:duration=5",
+        "-c:a", "pcm_s16le", str(audio_path),
+    ],
+    check=True,
+)
+
+demuxer = AudioDemuxer(audio_path)
+packet_decoder = AudioPacketDecoder(demuxer)
+raw = next(iter(packet_decoder.decode(next(iter(demuxer)))))
+print(f"{raw.sample_format = }, {raw.data.dtype = }, {raw.data.shape = }, "
+      f"{raw.sample_rate = }")
+
+# %%
+# Those are the true source samples - 16-bit integers here, not floats in
+# ``[-1, 1]``. ``AudioConverter`` is what normalizes them, and it can resample
+# and change the channel count on the way.
+#
+# It differs from ``ColorConverter`` in one important way: resampling is an
+# interpolation filter, so the sample it emits at a given instant depends on
+# input samples on *both* sides of it. The converter therefore holds the tail
+# of each frame back until the next one arrives - which is why ``convert()``
+# can return fewer samples than it was given, and why the pipeline ends with
+# ``drain()``. Leave that call out and you lose the end of the stream.
+demuxer = AudioDemuxer(audio_path)
+packet_decoder = AudioPacketDecoder(demuxer)
+audio_converter = AudioConverter(sample_rate=16_000, num_channels=1)
+
+chunks = []
+for packet in demuxer:
+    chunks += [audio_converter.convert(raw) for raw in packet_decoder.decode(packet)]
+chunks += [audio_converter.convert(raw) for raw in packet_decoder.drain()]
+chunks.append(audio_converter.drain())  # don't forget me
+
+samples = torch.cat([chunk.data for chunk in chunks], dim=1)
+print(f"{samples.shape = }, {samples.dtype = }, "
+      f"{chunks[-1].data.shape[1]} samples came out of drain()")
+
+# %%
+# .. warning::
+#
+#    These blocks do no pre-roll. A lossy codec's first frames after a seek are
+#    subtly wrong until it re-primes, and a resampler started mid-stream emits
+#    samples on a grid of its own, so samples decoded after a
+#    ``demuxer.seek()`` do not line up bit-for-bit with the same region of a
+#    whole-file decode. Decoding a margin before your target and discarding it
+#    is up to you. :class:`~torchcodec.decoders.AudioDecoder` does all of this
+#    for you.

src/torchcodec/_core/AudioCommon.cpp

@@ -6,6 +6,8 @@
 
 #include "AudioCommon.h"
 
+#include <vector>
+
 namespace facebook::torchcodec {
 
 torch::headeronly::ScalarType sample_format_dtype(
@@ -35,4 +37,65 @@ torch::headeronly::ScalarType sample_format_dtype(
   return kStableUInt8;
 }
 
+AVSampleFormat planar_sample_format(torch::headeronly::ScalarType dtype) {
+  switch (dtype) {
+    case kStableUInt8:
+      return AV_SAMPLE_FMT_U8P;
+    case kStableInt16:
+      return AV_SAMPLE_FMT_S16P;
+    case kStableInt32:
+      return AV_SAMPLE_FMT_S32P;
+    case kStableInt64:
+      return AV_SAMPLE_FMT_S64P;
+    case kStableFloat32:
+      return AV_SAMPLE_FMT_FLTP;
+    case kStableFloat64:
+      return AV_SAMPLE_FMT_DBLP;
+    default:
+      break;
+  }
+  STD_TORCH_CHECK(
+      false,
+      "Unsupported dtype for audio samples. Expected one of uint8, int16, "
+      "int32, int64, float32 or float64.");
+  return AV_SAMPLE_FMT_NONE;
+}
+
+torch::stable::Tensor swr_convert_to_tensor(
+    const UniqueSwrContext& swr_context,
+    const uint8_t** src_planes,
+    int num_src_samples,
+    int num_out_channels,
+    int64_t num_out_samples_bound) {
+  torch::stable::Tensor out = torch::stable::empty(
+      {num_out_channels, num_out_samples_bound}, kStableFloat32);
+  if (num_out_samples_bound == 0) {
+    return out;
+  }
+
+  // A contiguous [num_out_channels, N] float32 tensor is exactly FLTP: one
+  // plane per row.
+  int64_t bytes_per_channel =
+      num_out_samples_bound * av_get_bytes_per_sample(kAudioOutSampleFormat);
+  auto* base = static_cast<uint8_t*>(out.mutable_data_ptr());
+  std::vector<uint8_t*> out_planes(num_out_channels);
+  for (int channel = 0; channel < num_out_channels; ++channel) {
+    out_planes[channel] = base + channel * bytes_per_channel;
+  }
+
+  int num_out_samples = swr_convert(
+      swr_context.get(),
+      out_planes.data(),
+      static_cast<int>(num_out_samples_bound),
+      src_planes,
+      num_src_samples);
+  STD_TORCH_CHECK(
+      num_out_samples >= 0,
+      "Error in swr_convert: ",
+      get_ffmpeg_error_string_from_error_code(num_out_samples));
+
+  return torch::stable::narrow(
+      out, /*dim=*/1, /*start=*/0, /*length=*/num_out_samples);
+}
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/AudioCommon.h

@@ -15,9 +15,34 @@
 
 namespace facebook::torchcodec {
 
+// Decoded audio always leaves torchcodec as normalized float32 samples,
+// one plane per channel.
+constexpr AVSampleFormat kAudioOutSampleFormat = AV_SAMPLE_FMT_FLTP;
+
 // The dtype that holds `sample_format`'s samples exactly. Planar and packed
 // variants of a format share a sample type, which is why this doesn't care
 // which one it is given.
 torch::headeronly::ScalarType sample_format_dtype(AVSampleFormat sample_format);
 
+// The inverse: the sample format that a contiguous [num_channels, num_samples]
+// tensor of `dtype` already is. One row per channel is precisely what planar
+// means, which is what lets swresample read such a tensor's rows directly.
+AVSampleFormat planar_sample_format(torch::headeronly::ScalarType dtype);
+
+// Runs swr_convert() straight into a fresh float32 [num_out_channels, N]
+// tensor, narrowed to the number of samples it actually produced - which is
+// at most `num_out_samples_bound` and is often fewer, since swresample holds
+// back the samples that need future input. Pass a null `src_planes` to flush
+// those out.
+//
+// `src_planes` is deliberately `const uint8_t**` rather than
+// `const uint8_t* const*`: swr_convert() takes the former up to FFmpeg 6 and
+// the latter from 7 on, and only the former converts to both.
+torch::stable::Tensor swr_convert_to_tensor(
+    const UniqueSwrContext& swr_context,
+    const uint8_t** src_planes,
+    int num_src_samples,
+    int num_out_channels,
+    int64_t num_out_samples_bound);
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/AudioConverter.cpp

@@ -0,0 +1,130 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "AudioConverter.h"
+
+#include <vector>
+
+#include "AudioCommon.h"
+
+namespace facebook::torchcodec {
+
+AudioConverter::AudioConverter(
+    std::optional<int> sample_rate,
+    std::optional<int> num_channels)
+    : requested_sample_rate_(sample_rate),
+      requested_num_channels_(num_channels) {
+  STD_TORCH_CHECK(
+      !sample_rate.has_value() || *sample_rate > 0,
+      "sample_rate must be > 0. Got: ",
+      sample_rate.value_or(0));
+  STD_TORCH_CHECK(
+      !num_channels.has_value() || *num_channels > 0,
+      "num_channels must be > 0. Got: ",
+      num_channels.value_or(0));
+}
+
+void AudioConverter::reset() {
+  swr_context_.reset();
+  src_sample_format_ = AV_SAMPLE_FMT_NONE;
+  src_sample_rate_ = 0;
+  src_num_channels_ = 0;
+  out_sample_rate_ = 0;
+  out_num_channels_ = 0;
+}
+
+torch::stable::Tensor AudioConverter::convert(
+    const torch::stable::Tensor& samples,
+    int sample_rate) {
+  STD_TORCH_CHECK(
+      samples.dim() == 2,
+      "Expected a 2D [num_channels, num_samples] tensor, got a ",
+      samples.dim(),
+      "D one.");
+  STD_TORCH_CHECK(samples.is_contiguous(), "The samples must be contiguous.");
+  STD_TORCH_CHECK(sample_rate > 0, "sample_rate must be > 0.");
+
+  AVSampleFormat src_sample_format =
+      planar_sample_format(samples.scalar_type());
+  int num_channels = static_cast<int>(samples.sizes()[0]);
+  int num_samples = static_cast<int>(samples.sizes()[1]);
+  STD_TORCH_CHECK(
+      num_channels > 0, "The samples must have at least 1 channel.");
+
+  if (swr_context_ == nullptr) {
+    src_sample_format_ = src_sample_format;
+    src_sample_rate_ = sample_rate;
+    src_num_channels_ = num_channels;
+    out_sample_rate_ = requested_sample_rate_.value_or(sample_rate);
+    out_num_channels_ = requested_num_channels_.value_or(num_channels);
+    swr_context_.reset(create_swr_context(
+        src_sample_format_,
+        kAudioOutSampleFormat,
+        src_sample_rate_,
+        out_sample_rate_,
+        src_num_channels_,
+        out_num_channels_));
+  } else {
+    // swresample is configured once, from the first samples we see, and its
+    // buffered state is tied to that configuration. Rather than silently
+    // reconfiguring - which would discard whatever it still holds - we make the
+    // caller decide, by reset()ing.
+    STD_TORCH_CHECK(
+        src_sample_format == src_sample_format_ &&
+            sample_rate == src_sample_rate_ &&
+            num_channels == src_num_channels_,
+        "This AudioConverter was set up for ",
+        src_num_channels_,
+        " channels of ",
+        av_get_sample_fmt_name(src_sample_format_),
+        " at ",
+        src_sample_rate_,
+        " Hz, but got ",
+        num_channels,
+        " channels of ",
+        av_get_sample_fmt_name(src_sample_format),
+        " at ",
+        sample_rate,
+        " Hz. Call reset() to convert a different stream.");
+  }
+
+  const auto* base = static_cast<const uint8_t*>(samples.const_data_ptr());
+  int64_t bytes_per_channel =
+      num_samples * av_get_bytes_per_sample(src_sample_format);
+  std::vector<const uint8_t*> src_planes(num_channels);
+  for (int channel = 0; channel < num_channels; ++channel) {
+    src_planes[channel] = base + channel * bytes_per_channel;
+  }
+
+  // TODO_API_BREAKDOWN CC P2: the CPUDeviceInterface relies on
+  // swr_convert_to_tensor but only for flushing. We effectively have two
+  // code-paths doing libswresample conversion. Worth aligning.
+  return swr_convert_to_tensor(
+      swr_context_,
+      src_planes.data(),
+      num_samples,
+      out_num_channels_,
+      get_swr_output_num_samples_bound(
+          swr_context_, num_samples, src_sample_rate_, out_sample_rate_));
+}
+
+torch::stable::Tensor AudioConverter::drain() {
+  STD_TORCH_CHECK(
+      swr_context_ != nullptr,
+      "This AudioConverter hasn't converted any samples, so there is nothing "
+      "to drain and no way to know what shape the result should have.");
+  // A null input is what tells swr_convert() to flush. Unlike the convert()
+  // path we ask swresample how much it is holding rather than deriving a bound
+  // from an input size, since here there is no input.
+  return swr_convert_to_tensor(
+      swr_context_,
+      /*src_planes=*/nullptr,
+      /*num_src_samples=*/0,
+      out_num_channels_,
+      swr_get_out_samples(swr_context_.get(), 0));
+}
+
+} // namespace facebook::torchcodec