Enable decoding audio with >8 channels (#1166)

Author: Dan-Flores

src/torchcodec/_core/Encoder.cpp

@@ -361,6 +361,12 @@ UniqueAVFrame AudioEncoder::maybeConvertAVFrame(const UniqueAVFrame& avFrame) {
         avFrame,
         outNumChannels_));
   }
+  // convertAudioAVFrameSamples uses avFrame's extended_data field, so we ensure
+  // it's the same as data. This should always be the case since we validated
+  // earlier that we have less than AV_NUM_DATA_POINTERS channels.
+  TORCH_CHECK(
+      avFrame->data == avFrame->extended_data,
+      "Codec context data and extended_data pointers differ, this is unexpected.")
   UniqueAVFrame convertedAVFrame = convertAudioAVFrameSamples(
       swrContext_,
       avFrame,

src/torchcodec/_core/FFMPEGCommon.cpp

@@ -512,12 +512,16 @@ UniqueAVFrame convertAudioAVFrameSamples(
       "Could not allocate frame buffers for sample format conversion: ",
       getFFMPEGErrorStringFromErrorCode(status));
 
+  // Below we use AVFrame->extended_data instead of AVFrame->data to support
+  // decoding audio with >8 audio channels. extended_data contains pointers
+  // for all channels, while data only contains AV_NUM_DATA_POINTERS (8).
+  // https://ffmpeg.org/doxygen/trunk/structAVFrame.html#afca04d808393822625e09b5ba91c6756
   auto numConvertedSamples = swr_convert(
       swrContext.get(),
-      convertedAVFrame->data,
+      convertedAVFrame->extended_data,
       convertedAVFrame->nb_samples,
       static_cast<const uint8_t**>(
-          const_cast<const uint8_t**>(srcAVFrame->data)),
+          const_cast<const uint8_t**>(srcAVFrame->extended_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

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -571,9 +571,8 @@ void SingleStreamDecoder::addAudioStream(
       "seek_mode must be 'approximate' for audio streams.");
   if (audioStreamOptions.numChannels.has_value()) {
     TORCH_CHECK(
-        *audioStreamOptions.numChannels > 0 &&
-            *audioStreamOptions.numChannels <= AV_NUM_DATA_POINTERS,
-        "num_channels must be > 0 and <= AV_NUM_DATA_POINTERS (usually 8). Got: ",
+        *audioStreamOptions.numChannels > 0,
+        "num_channels must be > 0. Got: ",
         *audioStreamOptions.numChannels);
   }
 

test/resources/sine_16ch_s16.wav

@@ -0,0 +1,66 @@
+[
+  {
+    "pts_time": "0.000000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.064000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.128000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.192000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.256000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.320000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.384000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.448000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.512000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.576000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.640000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.704000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.768000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.832000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.896000",
+    "duration_time": "0.064000"
+  },
+  {
+    "pts_time": "0.960000",
+    "duration_time": "0.040000"
+  }
+]

test/resources/sine_16ch_s16.wav.stream0.all_frames_info.json

@@ -42,6 +42,7 @@
     needs_cuda,
     needs_ffmpeg_cli,
     psnr,
+    SINE_16_CHANNEL_S16,
     SINE_MONO_S16,
     SINE_MONO_S32,
     SINE_MONO_S32_44100,
@@ -1947,7 +1948,9 @@ def test_cpu_fallback_no_fallback_on_supported_video(self, device):
 
 
 class TestAudioDecoder:
-    @pytest.mark.parametrize("asset", (NASA_AUDIO, NASA_AUDIO_MP3, SINE_MONO_S32))
+    @pytest.mark.parametrize(
+        "asset", (NASA_AUDIO, NASA_AUDIO_MP3, SINE_MONO_S32, SINE_16_CHANNEL_S16)
+    )
     def test_metadata(self, asset):
         decoder = AudioDecoder(asset.path)
         assert isinstance(decoder.metadata, AudioStreamMetadata)
@@ -2005,7 +2008,7 @@ def test_get_all_samples_with_range(self, asset, stop_seconds):
         assert samples.sample_rate == asset.sample_rate
         assert samples.pts_seconds == asset.get_frame_info(idx=0).pts_seconds
 
-    @pytest.mark.parametrize("asset", (NASA_AUDIO, NASA_AUDIO_MP3))
+    @pytest.mark.parametrize("asset", (NASA_AUDIO, NASA_AUDIO_MP3, SINE_16_CHANNEL_S16))
     def test_get_all_samples(self, asset):
         decoder = AudioDecoder(asset.path)
         torch.testing.assert_close(
@@ -2274,8 +2277,7 @@ def test_samples_duration(self, asset, sample_rate):
     # that the extra tensor allocation that happens within
     # maybeFlushSwrBuffers() is correct.
     @pytest.mark.parametrize("sample_rate", (None, 16_000))
-    # FFmpeg can handle up to AV_NUM_DATA_POINTERS=8 channels
-    @pytest.mark.parametrize("num_channels", (1, 2, 8, None))
+    @pytest.mark.parametrize("num_channels", (1, 2, 8, 16, 24, None))
     def test_num_channels(self, asset, sample_rate, num_channels):
         decoder = AudioDecoder(
             asset.path, sample_rate=sample_rate, num_channels=num_channels
@@ -2289,12 +2291,12 @@ def test_num_channels(self, asset, sample_rate, num_channels):
 
     @pytest.mark.parametrize("asset", (SINE_MONO_S32, NASA_AUDIO_MP3))
     def test_num_channels_errors(self, asset):
-        with pytest.raises(
-            RuntimeError, match="num_channels must be > 0 and <= AV_NUM_DATA_POINTERS"
-        ):
+        with pytest.raises(RuntimeError, match="num_channels must be > 0"):
             AudioDecoder(asset.path, num_channels=0)
-        with pytest.raises(
-            RuntimeError, match="num_channels must be > 0 and <= AV_NUM_DATA_POINTERS"
-        ):
-            # FFmpeg can handle up to AV_NUM_DATA_POINTERS=8 channels
-            AudioDecoder(asset.path, num_channels=9)
+        for num_channels in (15, 23):
+            with pytest.raises(RuntimeError, match="Couldn't initialize SwrContext:"):
+                decoder = AudioDecoder(asset.path, num_channels=num_channels)
+                # Call get_all_samples to trigger num_channels conversion.
+                # FFmpeg fails to find a default layout for certain channel counts,
+                # which causes SwrContext to fail to initialize.
+                decoder.get_all_samples()

test/test_decoders.py

@@ -834,3 +834,20 @@ def sample_format(self) -> str:
         )
     },
 )
+
+# 16-channel audio for testing support for >8 channels. Generated with:
+# ffmpeg -i test/resources/sine_mono_s32.wav -t 1 -filter_complex "[0]asplit=16[s0][s1][s2][s3][s4][s5][s6][s7][s8][s9][s10][s11][s12][s13][s14][s15];[s0][s1][s2][s3][s4][s5][s6][s7][s8][s9][s10][s11][s12][s13][s14][s15]amerge=inputs=16" -c:a pcm_s16le test/resources/sine_16ch_s16.wav
+SINE_16_CHANNEL_S16 = TestAudio(
+    filename="sine_16ch_s16.wav",
+    default_stream_index=0,
+    frames={},  # Automatically loaded from json file
+    stream_infos={
+        0: TestAudioStreamInfo(
+            sample_rate=16_000,
+            num_channels=16,
+            duration_seconds=1,
+            num_frames=16,
+            sample_format="s16",
+        )
+    },
+)