Support multiple files in the Model Explorer Mediapipe Adapter

PiperOrigin-RevId: 705688677
google-ai-edge · Dec 16, 2024 · be0fbc5 · be0fbc5
1 parent 2da380f
commit be0fbc5
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 0 deletions.
diff --git a/src/builtin-adapter/BUILD b/src/builtin-adapter/BUILD
@@ -133,6 +133,7 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
+        "//mediapipe_adapter:mediapipe_to_json",
         "@org_tensorflow//tensorflow/core:lib",
     ],
 )

diff --git a/src/builtin-adapter/python/convert_wrapper/BUILD b/src/builtin-adapter/python/convert_wrapper/BUILD
@@ -23,5 +23,6 @@ pybind_extension(
         "//:direct_saved_model_to_json_graph_convert",
         "//:model_json_graph_convert",
         "//:visualize_config",
+        "//mediapipe_adapter:mediapipe_to_json",
     ],
 )
diff --git a/src/builtin-adapter/python/convert_wrapper/_pywrap_convert_wrapper.pyi b/src/builtin-adapter/python/convert_wrapper/_pywrap_convert_wrapper.pyi
@@ -20,6 +20,7 @@ class VisualizeConfig:
 def ConvertFlatbufferDirectlyToJson(arg0: VisualizeConfig, arg1: str) -> str: ...
 def ConvertFlatbufferToJson(arg0: VisualizeConfig, arg1: str, arg2: bool) -> str: ...
 def ConvertGraphDefDirectlyToJson(arg0: VisualizeConfig, arg1: str) -> str: ...
+def ConvertMediapipeToJson(arg0: VisualizeConfig, arg1: str) -> str: ...
 def ConvertMlirToJson(arg0: VisualizeConfig, arg1: str) -> str: ...
 def ConvertSavedModelDirectlyToJson(arg0: VisualizeConfig, arg1: str) -> str: ...
 def ConvertSavedModelToJson(arg0: VisualizeConfig, arg1: str) -> str: ...
diff --git a/src/builtin-adapter/python/convert_wrapper/convert_wrapper.cc b/src/builtin-adapter/python/convert_wrapper/convert_wrapper.cc
@@ -21,6 +21,7 @@
 #include "pybind11/pybind11.h"
 #include "direct_flatbuffer_to_json_graph_convert.h"
 #include "direct_saved_model_to_json_graph_convert.h"
+#include "mediapipe_adapter/mediapipe_to_json.h"
 #include "model_json_graph_convert.h"
 #include "visualize_config.h"
 
@@ -134,6 +135,22 @@ PYBIND11_MODULE(_pywrap_convert_wrapper, m) {
       Converts a MLIR textual/bytecode file to visualizer JSON string.
       Raises `RuntimeError` exception if failed.
       )pbdoc");
+  m.def(
+      "ConvertMediapipeToJson",
+      [](const VisualizeConfig& config,
+         absl::string_view model_path) -> std::string {
+        const absl::StatusOr<std::string> json_or_status =
+            ::tooling::visualization_client::ConvertMediapipeToJson(config,
+                                                                    model_path);
+        if (!json_or_status.ok()) {
+          throw std::runtime_error(json_or_status.status().ToString());
+        }
+        return json_or_status.value();
+      },
+      R"pbdoc(
+      Converts a Mediapipe pipeline to visualizer JSON string.
+      Raises `RuntimeError` exception if failed.
+      )pbdoc");
 }
 
 }  // namespace pybind11