Add utility functions to fetch available formats/devices/codecs/protocols. (#2958)

Summary:
This commit adds utility functions that fetch the available/supported formats/devices/codecs.

These functions are mostly same with commands like `ffmpeg -decoders`. But the use of `ffmpeg` CLI can report different resutls if there are multiple installation of FFmpegs. Or, the CLI might not be available.

Pull Request resolved: https://github.com/pytorch/audio/pull/2958

Reviewed By: hwangjeff

Differential Revision: D42371640

Pulled By: mthrok

fbshipit-source-id: 96a96183815a126cb1adc97ab7754aef216fff6f

test/torchaudio_unittest/utils/ffmpeg_utils_test.py

@@ -25,3 +25,17 @@ class TestFFmpegUtils(PytorchTestCase):
         """`get_versions` does not crash"""
         versions = ffmpeg_utils.get_versions()
         assert set(versions.keys()) == {"libavutil", "libavcodec", "libavformat", "libavfilter", "libavdevice"}
+
+    def test_available_stuff(self):
+        """get_encoders|decoders|muxers|demuxers|devices function does not segfault"""
+
+        ffmpeg_utils.get_demuxers()
+        ffmpeg_utils.get_muxers()
+        ffmpeg_utils.get_audio_decoders()
+        ffmpeg_utils.get_audio_encoders()
+        ffmpeg_utils.get_video_decoders()
+        ffmpeg_utils.get_video_encoders()
+        ffmpeg_utils.get_input_devices()
+        ffmpeg_utils.get_output_devices()
+        ffmpeg_utils.get_input_protocols()
+        ffmpeg_utils.get_output_protocols()

torchaudio/csrc/ffmpeg/utils.cpp

@@ -29,8 +29,98 @@ c10::Dict<std::string, std::tuple<int64_t, int64_t, int64_t>> get_versions() {
 #undef add_version
 }
 
+c10::Dict<std::string, std::string> get_demuxers(bool req_device) {
+  c10::Dict<std::string, std::string> ret;
+  const AVInputFormat* fmt = nullptr;
+  void* i = nullptr;
+  while ((fmt = av_demuxer_iterate(&i))) {
+    assert(fmt);
+    bool is_device = [&]() {
+      const AVClass* avclass = fmt->priv_class;
+      return avclass && AV_IS_INPUT_DEVICE(avclass->category);
+    }();
+    if (req_device == is_device) {
+      ret.insert(fmt->name, fmt->long_name);
+    }
+  }
+  return ret;
+}
+
+c10::Dict<std::string, std::string> get_muxers(bool req_device) {
+  c10::Dict<std::string, std::string> ret;
+  const AVOutputFormat* fmt = nullptr;
+  void* i = nullptr;
+  while ((fmt = av_muxer_iterate(&i))) {
+    assert(fmt);
+    bool is_device = [&]() {
+      const AVClass* avclass = fmt->priv_class;
+      return avclass && AV_IS_OUTPUT_DEVICE(avclass->category);
+    }();
+    if (req_device == is_device) {
+      ret.insert(fmt->name, fmt->long_name);
+    }
+  }
+  return ret;
+}
+
+c10::Dict<std::string, std::string> get_codecs(
+    AVMediaType type,
+    bool req_encoder) {
+  const AVCodec* c = nullptr;
+  void* i = nullptr;
+  c10::Dict<std::string, std::string> ret;
+  while ((c = av_codec_iterate(&i))) {
+    assert(c);
+    if ((req_encoder && av_codec_is_encoder(c)) ||
+        (!req_encoder && av_codec_is_decoder(c))) {
+      if (c->type == type && c->name) {
+        ret.insert(c->name, c->long_name ? c->long_name : "");
+      }
+    }
+  }
+  return ret;
+}
+
+std::vector<std::string> get_protocols(bool output) {
+  void* opaque = nullptr;
+  const char* name = nullptr;
+  std::vector<std::string> ret;
+  while ((name = avio_enum_protocols(&opaque, output))) {
+    assert(name);
+    ret.emplace_back(name);
+  }
+  return ret;
+}
+
 TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   m.def("torchaudio::ffmpeg_get_versions", &get_versions);
+  m.def("torchaudio::ffmpeg_get_muxers", []() { return get_muxers(false); });
+  m.def(
+      "torchaudio::ffmpeg_get_demuxers", []() { return get_demuxers(false); });
+  m.def("torchaudio::ffmpeg_get_input_devices", []() {
+    return get_demuxers(true);
+  });
+  m.def("torchaudio::ffmpeg_get_output_devices", []() {
+    return get_muxers(true);
+  });
+  m.def("torchaudio::ffmpeg_get_audio_decoders", []() {
+    return get_codecs(AVMEDIA_TYPE_AUDIO, false);
+  });
+  m.def("torchaudio::ffmpeg_get_audio_encoders", []() {
+    return get_codecs(AVMEDIA_TYPE_AUDIO, true);
+  });
+  m.def("torchaudio::ffmpeg_get_video_decoders", []() {
+    return get_codecs(AVMEDIA_TYPE_VIDEO, false);
+  });
+  m.def("torchaudio::ffmpeg_get_video_encoders", []() {
+    return get_codecs(AVMEDIA_TYPE_VIDEO, true);
+  });
+  m.def("torchaudio::ffmpeg_get_input_protocols", []() {
+    return get_protocols(false);
+  });
+  m.def("torchaudio::ffmpeg_get_output_protocols", []() {
+    return get_protocols(true);
+  });
 }
 
 } // namespace

torchaudio/io/_stream_reader.py

@@ -234,7 +234,9 @@ _video_stream_index = """The source video stream index.
 _decoder = """The name of the decoder to be used.
                 When provided, use the specified decoder instead of the default one.
 
-                To list the available decoders, you can use `ffmpeg -decoders` command.
+                To list the available decoders, please use
+                :py:func:`~torchaudio.utils.ffmpeg_utils.get_audio_decoders` for audio, and
+                :py:func:`~torchaudio.utils.ffmpeg_utils.get_video_decoders` for video.
 
                 Default: ``None``."""
 
@@ -340,14 +342,16 @@ class StreamReader:
 
                https://ffmpeg.org/ffmpeg-formats.html#Demuxers
 
-               Use `ffmpeg -demuxers` to list the values available in the current environment.
+               Please use :py:func:`~torchaudio.utils.ffmpeg_utils.get_demuxers` to list the
+               demultiplexers available in the current environment.
 
                For device access, the available values vary based on hardware (AV device) and
                software configuration (ffmpeg build).
 
                https://ffmpeg.org/ffmpeg-devices.html#Input-Devices
 
-               Use `ffmpeg -devices` to list the values available in the current environment.
+               Please use :py:func:`~torchaudio.utils.ffmpeg_utils.get_input_devices` to list
+               the input devices available in the current environment.
 
         option (dict of str to str, optional):
             Custom option passed when initializing format context (opening source).

torchaudio/io/_stream_writer.py

@@ -15,7 +15,9 @@ def _format_doc(**kwargs):
 _encoder = """The name of the encoder to be used.
                 When provided, use the specified encoder instead of the default one.
 
-                To list the available encoders, you can use ``ffmpeg -encoders`` command.
+                To list the available encoders, please use
+                :py:func:`~torchaudio.utils.ffmpeg_utils.get_audio_encoders` for audio, and
+                :py:func:`~torchaudio.utils.ffmpeg_utils.get_video_encoders` for video.
 
                 Default: ``None``."""
 
@@ -82,7 +84,8 @@ class StreamWriter:
 
                https://ffmpeg.org/ffmpeg-formats.html#Muxers
 
-               Use `ffmpeg -muxers` to list the values available in the current environment.
+               Please use :py:func:`~torchaudio.utils.ffmpeg_utils.get_muxers` to list the
+               multiplexers available in the current environment.
 
                For device access, the available values vary based on hardware (AV device) and
                software configuration (ffmpeg build).
@@ -90,7 +93,8 @@ class StreamWriter:
 
                https://ffmpeg.org/ffmpeg-devices.html#Output-Devices
 
-               Use `ffmpeg -devices` to list the values available in the current environment.
+               Please use :py:func:`~torchaudio.utils.ffmpeg_utils.get_output_devices` to list
+               the output devices available in the current environment.
 
         buffer_size (int):
             The internal buffer size in byte. Used only when `dst` is a file-like object.

torchaudio/utils/ffmpeg_utils.py

@@ -2,7 +2,7 @@
 
 It affects functionalities in :py:mod:`torchaudio.io` (and indirectly :py:func:`torchaudio.load`).
 """
-from typing import Dict, Tuple
+from typing import Dict, List, Tuple
 
 import torch
 
@@ -59,3 +59,169 @@ def set_log_level(level: int):
 
     """
     torch.ops.torchaudio.ffmpeg_set_log_level(level)
+
+
+def get_demuxers() -> Dict[str, str]:
+    """Get the available demuxers.
+
+    Returns:
+        Dict[str, str]: Mapping from demuxer (format) short name to long name.
+
+    Example
+        >>> for k, v in get_demuxers().items():
+        >>>     print(f"{k}: {v}")
+        ... aa: Audible AA format files
+        ... aac: raw ADTS AAC (Advanced Audio Coding)
+        ... aax: CRI AAX
+        ... ac3: raw AC-3
+    """
+    return torch.ops.torchaudio.ffmpeg_get_demuxers()
+
+
+def get_muxers() -> Dict[str, str]:
+    """Get the available muxers.
+
+    Returns:
+        Dict[str, str]: Mapping from muxer (format) short name to long name.
+
+    Example
+        >>> for k, v in get_muxers().items():
+        >>>     print(f"{k}: {v}")
+        ... a64: a64 - video for Commodore 64
+        ... ac3: raw AC-3
+        ... adts: ADTS AAC (Advanced Audio Coding)
+        ... adx: CRI ADX
+        ... aiff: Audio IFF
+    """
+    return torch.ops.torchaudio.ffmpeg_get_muxers()
+
+
+def get_audio_decoders() -> Dict[str, str]:
+    """Get the available audio decoders.
+
+    Returns:
+        Dict[str, str]: Mapping from decoder short name to long name.
+
+    Example
+        >>> for k, v in get_audio_decoders().items():
+        >>>     print(f"{k}: {v}")
+        ... a64: a64 - video for Commodore 64
+        ... ac3: raw AC-3
+        ... adts: ADTS AAC (Advanced Audio Coding)
+        ... adx: CRI ADX
+        ... aiff: Audio IFF
+    """
+    return torch.ops.torchaudio.ffmpeg_get_audio_decoders()
+
+
+def get_audio_encoders() -> Dict[str, str]:
+    """Get the available audio encoders.
+
+    Returns:
+        Dict[str, str]: Mapping from encoder short name to long name.
+
+    Example
+        >>> for k, v in get_audio_encoders().items():
+        >>>     print(f"{k}: {v}")
+        ... comfortnoise: RFC 3389 comfort noise generator
+        ... s302m: SMPTE 302M
+        ... aac: AAC (Advanced Audio Coding)
+        ... ac3: ATSC A/52A (AC-3)
+        ... ac3_fixed: ATSC A/52A (AC-3)
+        ... alac: ALAC (Apple Lossless Audio Codec)
+    """
+    return torch.ops.torchaudio.ffmpeg_get_audio_encoders()
+
+
+def get_video_decoders() -> Dict[str, str]:
+    """Get the available video decoders.
+
+    Returns:
+        Dict[str, str]: Mapping from decoder short name to long name.
+
+    Example
+        >>> for k, v in get_video_decoders().items():
+        >>>     print(f"{k}: {v}")
+        ... aasc: Autodesk RLE
+        ... aic: Apple Intermediate Codec
+        ... alias_pix: Alias/Wavefront PIX image
+        ... agm: Amuse Graphics Movie
+        ... amv: AMV Video
+        ... anm: Deluxe Paint Animation
+    """
+    return torch.ops.torchaudio.ffmpeg_get_video_decoders()
+
+
+def get_video_encoders() -> Dict[str, str]:
+    """Get the available video encoders.
+
+    Returns:
+        Dict[str, str]: Mapping from encoder short name to long name.
+
+    Example
+        >>> for k, v in get_audio_encoders().items():
+        >>>     print(f"{k}: {v}")
+        ... a64multi: Multicolor charset for Commodore 64
+        ... a64multi5: Multicolor charset for Commodore 64, extended with 5th color (colram)
+        ... alias_pix: Alias/Wavefront PIX image
+        ... amv: AMV Video
+        ... apng: APNG (Animated Portable Network Graphics) image
+        ... asv1: ASUS V1
+        ... asv2: ASUS V2
+    """
+    return torch.ops.torchaudio.ffmpeg_get_video_encoders()
+
+
+def get_input_devices() -> Dict[str, str]:
+    """Get the available input devices.
+
+    Returns:
+        Dict[str, str]: Mapping from device short name to long name.
+
+    Example
+        >>> for k, v in get_input_devices().items():
+        >>>     print(f"{k}: {v}")
+        ... avfoundation: AVFoundation input device
+        ... lavfi: Libavfilter virtual input device
+    """
+    return torch.ops.torchaudio.ffmpeg_get_input_devices()
+
+
+def get_output_devices() -> Dict[str, str]:
+    """Get the available output devices.
+
+    Returns:
+        Dict[str, str]: Mapping from device short name to long name.
+
+    Example
+        >>> for k, v in get_output_devices().items():
+        >>>     print(f"{k}: {v}")
+        ... audiotoolbox: AudioToolbox output device
+    """
+    return torch.ops.torchaudio.ffmpeg_get_output_devices()
+
+
+def get_input_protocols() -> List[str]:
+    """Get the supported input protocols.
+
+    Returns:
+        List[str]: The names of supported input protocols
+
+    Example
+        >>> print(get_input_protocols())
+        ... ['file', 'ftp', 'hls', 'http','https', 'pipe', 'rtmp', 'tcp', 'tls', 'udp', 'unix']
+    """
+    return torch.ops.torchaudio.ffmpeg_get_input_protocols()
+
+
+def get_output_protocols() -> List[str]:
+    """Get the supported output protocols.
+
+    Returns:
+        list of str: The names of supported output protocols
+
+    Example
+        >>> print(get_output_protocols())
+        ... ['file', 'ftp', 'http', 'https', 'md5', 'pipe', 'prompeg', 'rtmp', 'tee', 'tcp', 'tls', 'udp', 'unix']
+    """
+    return torch.ops.torchaudio.ffmpeg_get_output_protocols()