From be0fbc5a06292b749d100a4ffe06cf6cc59b4075 Mon Sep 17 00:00:00 2001
From: Google AI Edge <noreply@google.com>
Date: Thu, 12 Dec 2024 17:49:50 -0800
Subject: [PATCH] Support multiple files in the Model Explorer Mediapipe
 Adapter

PiperOrigin-RevId: 705688677
---
 src/builtin-adapter/BUILD                       |  1 +
 .../python/convert_wrapper/BUILD                |  1 +
 .../convert_wrapper/_pywrap_convert_wrapper.pyi |  1 +
 .../python/convert_wrapper/convert_wrapper.cc   | 17 +++++++++++++++++
 4 files changed, 20 insertions(+)

diff --git a/src/builtin-adapter/BUILD b/src/builtin-adapter/BUILD
index b5177731..972dcf1d 100644
--- 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
index 64e1eff4..45be96c3 100644
--- 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
index b0878ad1..c49f033a 100644
--- 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
index 1dfe4a90..b22ec963 100644
--- 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