From af2ad1abbe38c07961410a9fe7d000fd15cfa78e Mon Sep 17 00:00:00 2001
From: Jiuqiang Tang <jqtang@google.com>
Date: Fri, 30 Sep 2022 08:32:42 +0000
Subject: [PATCH] Switch MediaPipe Tasks Python and Java base layer to use
 MediaPipeBuiltinOpResolver by default.

PiperOrigin-RevId: 477927852
---
 .../tasks/java/com/google/mediapipe/tasks/core/jni/BUILD      | 2 +-
 .../java/com/google/mediapipe/tasks/core/jni/BUILD.bazel      | 2 +-
 .../mediapipe/tasks/core/jni/model_resources_cache_jni.cc     | 4 +++-
 mediapipe/tasks/python/core/pybind/BUILD                      | 1 +
 mediapipe/tasks/python/core/pybind/task_runner.cc             | 3 ++-
 5 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/BUILD b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/BUILD
index a1ec67517..cb3ef9656 100644
--- a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/BUILD
+++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/BUILD
@@ -28,10 +28,10 @@ cc_library_with_tflite(
     ],
     tflite_deps = [
         "//mediapipe/tasks/cc/core:model_resources_cache",
-        "@org_tensorflow//tensorflow/lite/core/shims:builtin_ops",
     ],
     deps = [
         "//mediapipe/java/com/google/mediapipe/framework/jni:mediapipe_framework_jni",
+        "//mediapipe/tasks/cc/core:mediapipe_builtin_op_resolver",
     ] + select({
         "//conditions:default": ["//third_party/java/jdk:jni"],
         "//mediapipe:android": [],
diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/BUILD.bazel b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/BUILD.bazel
index fba314b28..0eb74e7ff 100644
--- a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/BUILD.bazel
+++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/BUILD.bazel
@@ -34,10 +34,10 @@ cc_library_with_tflite(
     }),
     tflite_deps = [
         "//mediapipe/tasks/cc/core:model_resources_cache",
-        "@org_tensorflow//tensorflow/lite/core/shims:builtin_ops",
     ],
     deps = [
         "//mediapipe/java/com/google/mediapipe/framework/jni:mediapipe_framework_jni",
+        "//mediapipe/tasks/cc/core:mediapipe_builtin_op_resolver",
     ] + select({
         "//conditions:default": [],
         "//mediapipe:android": [],
diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/model_resources_cache_jni.cc b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/model_resources_cache_jni.cc
index 74ff4a689..aab022dec 100644
--- a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/model_resources_cache_jni.cc
+++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni/model_resources_cache_jni.cc
@@ -17,11 +17,13 @@
 #include <utility>
 
 #include "mediapipe/java/com/google/mediapipe/framework/jni/graph_service_jni.h"
+#include "mediapipe/tasks/cc/core/mediapipe_builtin_op_resolver.h"
 #include "mediapipe/tasks/cc/core/model_resources_cache.h"
 #include "tensorflow/lite/core/shims/cc/kernels/register.h"
 
 namespace {
 using ::mediapipe::tasks::core::kModelResourcesCacheService;
+using ::mediapipe::tasks::core::MediaPipeBuiltinOpResolver;
 using ::mediapipe::tasks::core::ModelResourcesCache;
 using HandleType = std::shared_ptr<ModelResourcesCache>*;
 }  // namespace
@@ -29,7 +31,7 @@ using HandleType = std::shared_ptr<ModelResourcesCache>*;
 JNIEXPORT jlong JNICALL MODEL_RESOURCES_CACHE_METHOD(
     nativeCreateModelResourcesCache)(JNIEnv* env, jobject thiz) {
   auto ptr = std::make_shared<ModelResourcesCache>(
-      absl::make_unique<tflite_shims::ops::builtin::BuiltinOpResolver>());
+      absl::make_unique<MediaPipeBuiltinOpResolver>());
   HandleType handle = new std::shared_ptr<ModelResourcesCache>(std::move(ptr));
   return reinterpret_cast<jlong>(handle);
 }
diff --git a/mediapipe/tasks/python/core/pybind/BUILD b/mediapipe/tasks/python/core/pybind/BUILD
index fab878135..b59635dc3 100644
--- a/mediapipe/tasks/python/core/pybind/BUILD
+++ b/mediapipe/tasks/python/core/pybind/BUILD
@@ -27,6 +27,7 @@ pybind_library(
         "//mediapipe/framework/api2:builder",
         "//mediapipe/framework/port:parse_text_proto",
         "//mediapipe/python/pybind:util",
+        "//mediapipe/tasks/cc/core:mediapipe_builtin_op_resolver",
         "//mediapipe/tasks/cc/core:task_runner",
         "@org_tensorflow//tensorflow/lite/core/api:op_resolver",
         "@pybind11_protobuf//pybind11_protobuf:native_proto_caster",
diff --git a/mediapipe/tasks/python/core/pybind/task_runner.cc b/mediapipe/tasks/python/core/pybind/task_runner.cc
index 52834bab2..cb13787c3 100644
--- a/mediapipe/tasks/python/core/pybind/task_runner.cc
+++ b/mediapipe/tasks/python/core/pybind/task_runner.cc
@@ -16,6 +16,7 @@
 
 #include "mediapipe/framework/calculator.pb.h"
 #include "mediapipe/python/pybind/util.h"
+#include "mediapipe/tasks/cc/core/mediapipe_builtin_op_resolver.h"
 #include "mediapipe/tasks/cc/core/task_runner.h"
 #include "pybind11/stl.h"
 #include "pybind11_protobuf/native_proto_caster.h"
@@ -75,7 +76,7 @@ mode) or not (synchronous mode).)doc");
         }
         auto task_runner = TaskRunner::Create(
             std::move(graph_config),
-            absl::make_unique<tflite::ops::builtin::BuiltinOpResolver>(),
+            absl::make_unique<core::MediaPipeBuiltinOpResolver>(),
             std::move(callback));
         RaisePyErrorIfNotOk(task_runner.status());
         return std::move(*task_runner);