Integrates the kMemoryManagerService into ImageToTensorCalculator and…

… InferenceCalculatorDarwinn. PiperOrigin-RevId: 604743255
google-ai-edge · Feb 6, 2024 · 9c54f1a · 9c54f1a
1 parent a18ed2e
commit 9c54f1a
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/mediapipe/calculators/tensor/BUILD b/mediapipe/calculators/tensor/BUILD
@@ -1175,6 +1175,8 @@ cc_library(
         ":image_to_tensor_utils",
         ":loose_headers",
         "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:memory_manager",
+        "//mediapipe/framework:memory_manager_service",
         "//mediapipe/framework:port",
         "//mediapipe/framework/api2:node",
         "//mediapipe/framework/formats:image",

diff --git a/mediapipe/calculators/tensor/image_to_tensor_calculator.cc b/mediapipe/calculators/tensor/image_to_tensor_calculator.cc
@@ -26,6 +26,8 @@
 #include "mediapipe/framework/formats/image_frame.h"
 #include "mediapipe/framework/formats/rect.pb.h"
 #include "mediapipe/framework/formats/tensor.h"
+#include "mediapipe/framework/memory_manager.h"
+#include "mediapipe/framework/memory_manager_service.h"
 #include "mediapipe/framework/port.h"
 #include "mediapipe/framework/port/canonical_errors.h"
 #include "mediapipe/framework/port/ret_check.h"
@@ -156,10 +158,14 @@ class ImageToTensorCalculator : public Node {
 #endif  // MEDIAPIPE_METAL_ENABLED
 #endif  // MEDIAPIPE_DISABLE_GPU
 
+    cc->UseService(kMemoryManagerService).Optional();
     return absl::OkStatus();
   }
 
   absl::Status Open(CalculatorContext* cc) {
+    if (cc->Service(kMemoryManagerService).IsAvailable()) {
+      memory_manager_ = &cc->Service(kMemoryManagerService).GetObject();
+    }
     options_ = cc->Options<mediapipe::ImageToTensorCalculatorOptions>();
     params_ = GetOutputTensorParams(options_);
     return absl::OkStatus();
@@ -221,8 +227,10 @@ class ImageToTensorCalculator : public Node {
 
     Tensor::ElementType output_tensor_type =
         GetOutputTensorType(image->UsesGpu(), params_);
-    Tensor tensor(output_tensor_type, {1, tensor_height, tensor_width,
-                                       GetNumOutputChannels(*image)});
+    Tensor tensor(
+        output_tensor_type,
+        {1, tensor_height, tensor_width, GetNumOutputChannels(*image)},
+        memory_manager_);
     MP_RETURN_IF_ERROR((image->UsesGpu() ? gpu_converter_ : cpu_converter_)
                            ->Convert(*image, roi, params_.range_min,
                                      params_.range_max,
@@ -301,6 +309,7 @@ class ImageToTensorCalculator : public Node {
   std::unique_ptr<ImageToTensorConverter> cpu_converter_;
   mediapipe::ImageToTensorCalculatorOptions options_;
   OutputTensorParams params_;
+  MemoryManager* memory_manager_ = nullptr;
 };
 
 MEDIAPIPE_REGISTER_NODE(ImageToTensorCalculator);

diff --git a/mediapipe/framework/BUILD b/mediapipe/framework/BUILD
@@ -1814,6 +1814,5 @@ cc_library(
     deps = [
         ":graph_service",
         ":memory_manager",
-        "@com_google_absl//absl/base:core_headers",
     ],
 )