From 7d8d3ab1968c6146106805566f3163189c2b6409 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Thu, 6 Apr 2023 08:40:12 -0700 Subject: [PATCH] Add EDGETPU_NNAPI delegate option in MediaPipe tasks API PiperOrigin-RevId: 522344828 --- .../processors/image_preprocessing_graph.cc | 4 +++- mediapipe/tasks/cc/core/BUILD | 2 ++ mediapipe/tasks/cc/core/base_options.cc | 5 +++++ mediapipe/tasks/cc/core/base_options.h | 2 ++ mediapipe/tasks/cc/core/base_options_test.cc | 18 +++++++++++++++++- mediapipe/tasks/cc/core/model_task_graph.cc | 3 +++ .../tasks/cc/core/proto/acceleration.proto | 1 + mediapipe/tasks/web/core/task_runner_test.ts | 2 ++ 8 files changed, 35 insertions(+), 2 deletions(-) diff --git a/mediapipe/tasks/cc/components/processors/image_preprocessing_graph.cc b/mediapipe/tasks/cc/components/processors/image_preprocessing_graph.cc index f07b2da8b..7a45bb148 100644 --- a/mediapipe/tasks/cc/components/processors/image_preprocessing_graph.cc +++ b/mediapipe/tasks/cc/components/processors/image_preprocessing_graph.cc @@ -117,7 +117,9 @@ absl::Status ConfigureImageToTensorCalculator( bool DetermineImagePreprocessingGpuBackend( const core::proto::Acceleration& acceleration) { - return acceleration.has_gpu(); + return acceleration.has_gpu() || + (acceleration.has_nnapi() && + acceleration.nnapi().accelerator_name() == "google-edgetpu"); } absl::Status ConfigureImagePreprocessingGraph( diff --git a/mediapipe/tasks/cc/core/BUILD b/mediapipe/tasks/cc/core/BUILD index 8b79aeaf5..1ba1fde33 100644 --- a/mediapipe/tasks/cc/core/BUILD +++ b/mediapipe/tasks/cc/core/BUILD @@ -44,7 +44,9 @@ cc_test( deps = [ ":base_options", ":utils", + "//mediapipe/calculators/tensor:inference_calculator_cc_proto", "//mediapipe/framework/port:gtest", + "//mediapipe/tasks/cc/core/proto:acceleration_cc_proto", "//mediapipe/tasks/cc/core/proto:external_file_cc_proto", "@com_google_googletest//:gtest_main", ], diff --git a/mediapipe/tasks/cc/core/base_options.cc b/mediapipe/tasks/cc/core/base_options.cc index 02a39081a..8bee6b469 100644 --- a/mediapipe/tasks/cc/core/base_options.cc +++ b/mediapipe/tasks/cc/core/base_options.cc @@ -57,6 +57,11 @@ proto::BaseOptions ConvertBaseOptionsToProto(BaseOptions* base_options) { case BaseOptions::Delegate::GPU: base_options_proto.mutable_acceleration()->mutable_gpu(); break; + case BaseOptions::Delegate::EDGETPU_NNAPI: + base_options_proto.mutable_acceleration() + ->mutable_nnapi() + ->set_accelerator_name("google-edgetpu"); + break; } return base_options_proto; diff --git a/mediapipe/tasks/cc/core/base_options.h b/mediapipe/tasks/cc/core/base_options.h index cdb3998c8..b6a0f0556 100644 --- a/mediapipe/tasks/cc/core/base_options.h +++ b/mediapipe/tasks/cc/core/base_options.h @@ -42,6 +42,8 @@ struct BaseOptions { enum Delegate { CPU = 0, GPU = 1, + // Edge TPU acceleration using NNAPI delegate. + EDGETPU_NNAPI = 2, }; Delegate delegate = CPU; diff --git a/mediapipe/tasks/cc/core/base_options_test.cc b/mediapipe/tasks/cc/core/base_options_test.cc index 23eb83d4a..dce95050d 100644 --- a/mediapipe/tasks/cc/core/base_options_test.cc +++ b/mediapipe/tasks/cc/core/base_options_test.cc @@ -2,8 +2,10 @@ #include +#include "mediapipe/calculators/tensor/inference_calculator.pb.h" #include "mediapipe/framework/port/gmock.h" #include "mediapipe/framework/port/gtest.h" +#include "mediapipe/tasks/cc/core/proto/acceleration.pb.h" #include "mediapipe/tasks/cc/core/proto/external_file.pb.h" #include "mediapipe/tasks/cc/core/utils.h" @@ -15,7 +17,7 @@ namespace tasks { namespace core { namespace { -TEST(BaseOptionsTest, ConverBaseOptionsToProtoWithFile) { +TEST(BaseOptionsTest, ConvertBaseOptionsToProtoWithFile) { BaseOptions base_options; base_options.model_asset_buffer = std::make_unique(LoadBinaryContent(kTestModelBundlePath)); @@ -24,6 +26,20 @@ TEST(BaseOptionsTest, ConverBaseOptionsToProtoWithFile) { EXPECT_TRUE(proto.model_asset().has_file_content()); } +TEST(BaseOptionsTest, ConvertBaseOptionsToProtoWithAcceleration) { + BaseOptions base_options; + proto::BaseOptions proto = ConvertBaseOptionsToProto(&base_options); + EXPECT_TRUE(proto.acceleration().has_tflite()); + + base_options.delegate = BaseOptions::Delegate::GPU; + proto = ConvertBaseOptionsToProto(&base_options); + EXPECT_TRUE(proto.acceleration().has_gpu()); + + base_options.delegate = BaseOptions::Delegate::EDGETPU_NNAPI; + proto = ConvertBaseOptionsToProto(&base_options); + EXPECT_EQ(proto.acceleration().nnapi().accelerator_name(), "google-edgetpu"); +} + } // namespace } // namespace core } // namespace tasks diff --git a/mediapipe/tasks/cc/core/model_task_graph.cc b/mediapipe/tasks/cc/core/model_task_graph.cc index 653c6b9ff..8767fb48b 100644 --- a/mediapipe/tasks/cc/core/model_task_graph.cc +++ b/mediapipe/tasks/cc/core/model_task_graph.cc @@ -134,6 +134,9 @@ class InferenceSubgraph : public Subgraph { case Acceleration::kGpu: delegate.mutable_gpu()->CopyFrom(acceleration.gpu()); break; + case Acceleration::kNnapi: + delegate.mutable_nnapi()->CopyFrom(acceleration.nnapi()); + break; case Acceleration::kTflite: delegate.mutable_tflite()->CopyFrom(acceleration.tflite()); break; diff --git a/mediapipe/tasks/cc/core/proto/acceleration.proto b/mediapipe/tasks/cc/core/proto/acceleration.proto index bdfaff4d2..c7215604a 100644 --- a/mediapipe/tasks/cc/core/proto/acceleration.proto +++ b/mediapipe/tasks/cc/core/proto/acceleration.proto @@ -33,5 +33,6 @@ message Acceleration { mediapipe.InferenceCalculatorOptions.Delegate.Xnnpack xnnpack = 1; mediapipe.InferenceCalculatorOptions.Delegate.Gpu gpu = 2; mediapipe.InferenceCalculatorOptions.Delegate.TfLite tflite = 4; + mediapipe.InferenceCalculatorOptions.Delegate.Nnapi nnapi = 5; } } diff --git a/mediapipe/tasks/web/core/task_runner_test.ts b/mediapipe/tasks/web/core/task_runner_test.ts index 52df0550a..8c9a1a5db 100644 --- a/mediapipe/tasks/web/core/task_runner_test.ts +++ b/mediapipe/tasks/web/core/task_runner_test.ts @@ -96,6 +96,7 @@ describe('TaskRunner', () => { xnnpack: undefined, gpu: undefined, tflite: {}, + nnapi: undefined, }, }; const mockBytesResultWithGpuDelegate = { @@ -113,6 +114,7 @@ describe('TaskRunner', () => { .SUSTAINED_SPEED, }, tflite: undefined, + nnapi: undefined, }, };