Migrate the binding of FFmpeg utils to PyBind11 (#3228)

Summary:
Utilities functions are only available to Python, so no need to use TorchBind for them.
This should allow us to remove link-whole flag when linking `libtorchaudio_ffmpeg` part.

Pull Request resolved: #3228

Reviewed By: nateanl

Differential Revision: D44639560

Pulled By: mthrok

fbshipit-source-id: 5116073ee8c5ab572c63ad123942c4826bfe1100

torchaudio/_extension/utils.py

@@ -88,9 +88,9 @@ def _init_ffmpeg():
 
     import torchaudio.lib._torchaudio_ffmpeg  # noqa
 
-    torch.ops.torchaudio.ffmpeg_init()
-    if torch.ops.torchaudio.ffmpeg_get_log_level() > 8:
-        torch.ops.torchaudio.ffmpeg_set_log_level(8)
+    torchaudio.lib._torchaudio_ffmpeg.init()
+    if torchaudio.lib._torchaudio_ffmpeg.get_log_level() > 8:
+        torchaudio.lib._torchaudio_ffmpeg.set_log_level(8)
 
 
 def _init_dll_path():

torchaudio/csrc/ffmpeg/CMakeLists.txt

@@ -21,7 +21,6 @@ set(
   stream_writer/stream_writer.cpp
   stream_writer/tensor_converter.cpp
   compat.cpp
-  utils.cpp
   )
 
 if (USE_CUDA)

torchaudio/csrc/ffmpeg/pybind/pybind.cpp

@@ -8,6 +8,97 @@ namespace torchaudio {
 namespace io {
 namespace {
 
+std::map<std::string, std::tuple<int64_t, int64_t, int64_t>> get_versions() {
+  std::map<std::string, std::tuple<int64_t, int64_t, int64_t>> ret;
+
+#define add_version(NAME)            \
+  {                                  \
+    int ver = NAME##_version();      \
+    ret.emplace(                     \
+        "lib" #NAME,                 \
+        std::make_tuple<>(           \
+            AV_VERSION_MAJOR(ver),   \
+            AV_VERSION_MINOR(ver),   \
+            AV_VERSION_MICRO(ver))); \
+  }
+
+  add_version(avutil);
+  add_version(avcodec);
+  add_version(avformat);
+  add_version(avfilter);
+  add_version(avdevice);
+  return ret;
+
+#undef add_version
+}
+
+std::map<std::string, std::string> get_demuxers(bool req_device) {
+  std::map<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.emplace(fmt->name, fmt->long_name);
+    }
+  }
+  return ret;
+}
+
+std::map<std::string, std::string> get_muxers(bool req_device) {
+  std::map<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.emplace(fmt->name, fmt->long_name);
+    }
+  }
+  return ret;
+}
+
+std::map<std::string, std::string> get_codecs(
+    AVMediaType type,
+    bool req_encoder) {
+  const AVCodec* c = nullptr;
+  void* i = nullptr;
+  std::map<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.emplace(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;
+}
+
+std::string get_build_config() {
+  return avcodec_configuration();
+}
+
 // The reason we inherit FileObj instead of making it an attribute
 // is so that FileObj is instantiated first.
 // AVIOContext must be initialized before AVFormat, and outlive AVFormat.
@@ -31,7 +122,31 @@ struct StreamWriterFileObj : private FileObj, public StreamWriter {
 };
 
 PYBIND11_MODULE(_torchaudio_ffmpeg, m) {
+  m.def("init", []() { avdevice_register_all(); });
+  m.def("get_log_level", []() { return av_log_get_level(); });
+  m.def("set_log_level", [](int level) { av_log_set_level(level); });
+  m.def("get_versions", &get_versions);
+  m.def("get_muxers", []() { return get_muxers(false); });
+  m.def("get_demuxers", []() { return get_demuxers(false); });
+  m.def("get_input_devices", []() { return get_demuxers(true); });
+  m.def("get_build_config", &get_build_config);
+  m.def("get_output_devices", []() { return get_muxers(true); });
+  m.def("get_audio_decoders", []() {
+    return get_codecs(AVMEDIA_TYPE_AUDIO, false);
+  });
+  m.def("get_audio_encoders", []() {
+    return get_codecs(AVMEDIA_TYPE_AUDIO, true);
+  });
+  m.def("get_video_decoders", []() {
+    return get_codecs(AVMEDIA_TYPE_VIDEO, false);
+  });
+  m.def("get_video_encoders", []() {
+    return get_codecs(AVMEDIA_TYPE_VIDEO, true);
+  });
+  m.def("get_input_protocols", []() { return get_protocols(false); });
+  m.def("get_output_protocols", []() { return get_protocols(true); });
   m.def("clear_cuda_context_cache", &clear_cuda_context_cache);
+
   py::class_<Chunk>(m, "Chunk", py::module_local())
       .def_readwrite("frames", &Chunk::frames)
       .def_readwrite("pts", &Chunk::pts);