From dbd6d72696d195d3e5358f7a9a081fc2fc9d1a72 Mon Sep 17 00:00:00 2001 From: Jiuqiang Tang Date: Tue, 7 Mar 2023 00:41:41 -0800 Subject: [PATCH] Add GetInputImageTensorSpecs into BaseVisionTaskApi for tasks api users to get input image tensor specifications. PiperOrigin-RevId: 514650593 --- mediapipe/tasks/cc/vision/core/BUILD | 2 + .../cc/vision/core/base_vision_task_api.h | 38 +++++++++++++++++++ .../vision/object_detector/object_detector.h | 2 +- .../object_detector/object_detector_test.cc | 31 +++++++++++++++ .../cc/vision/utils/image_tensor_specs.h | 2 +- 5 files changed, 73 insertions(+), 2 deletions(-) diff --git a/mediapipe/tasks/cc/vision/core/BUILD b/mediapipe/tasks/cc/vision/core/BUILD index 1f5ab5fa..5791c25f 100644 --- a/mediapipe/tasks/cc/vision/core/BUILD +++ b/mediapipe/tasks/cc/vision/core/BUILD @@ -38,10 +38,12 @@ cc_library( ":image_processing_options", ":running_mode", "//mediapipe/calculators/core:flow_limiter_calculator", + "//mediapipe/calculators/tensor:image_to_tensor_calculator_cc_proto", "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/tasks/cc/components/containers:rect", "//mediapipe/tasks/cc/core:base_task_api", "//mediapipe/tasks/cc/core:task_runner", + "//mediapipe/tasks/cc/vision/utils:image_tensor_specs", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", diff --git a/mediapipe/tasks/cc/vision/core/base_vision_task_api.h b/mediapipe/tasks/cc/vision/core/base_vision_task_api.h index a86b2cca..8b27ca19 100644 --- a/mediapipe/tasks/cc/vision/core/base_vision_task_api.h +++ b/mediapipe/tasks/cc/vision/core/base_vision_task_api.h @@ -25,12 +25,14 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" +#include "mediapipe/calculators/tensor/image_to_tensor_calculator.pb.h" #include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/tasks/cc/components/containers/rect.h" #include "mediapipe/tasks/cc/core/base_task_api.h" #include "mediapipe/tasks/cc/core/task_runner.h" #include "mediapipe/tasks/cc/vision/core/image_processing_options.h" #include "mediapipe/tasks/cc/vision/core/running_mode.h" +#include "mediapipe/tasks/cc/vision/utils/image_tensor_specs.h" namespace mediapipe { namespace tasks { @@ -44,6 +46,42 @@ class BaseVisionTaskApi : public tasks::core::BaseTaskApi { explicit BaseVisionTaskApi(std::unique_ptr runner, RunningMode running_mode) : BaseTaskApi(std::move(runner)), running_mode_(running_mode) {} + virtual ~BaseVisionTaskApi() {} + + virtual absl::StatusOr GetInputImageTensorSpecs() { + ImageTensorSpecs image_tensor_specs; + bool found_image_to_tensor_calculator = false; + for (auto& node : runner_->GetGraphConfig().node()) { + if (node.calculator() == "ImageToTensorCalculator") { + if (!found_image_to_tensor_calculator) { + found_image_to_tensor_calculator = true; + } else { + return absl::Status(CreateStatusWithPayload( + absl::StatusCode::kFailedPrecondition, + absl::StrCat( + "The graph has more than one ImageToTensorCalculator."))); + } + mediapipe::ImageToTensorCalculatorOptions options = + node.options().GetExtension( + mediapipe::ImageToTensorCalculatorOptions::ext); + image_tensor_specs.image_width = options.output_tensor_width(); + image_tensor_specs.image_height = options.output_tensor_height(); + image_tensor_specs.color_space = + tflite::ColorSpaceType::ColorSpaceType_RGB; + if (options.has_output_tensor_uint_range()) { + image_tensor_specs.tensor_type = tflite::TensorType_UINT8; + } else if (options.has_output_tensor_float_range()) { + image_tensor_specs.tensor_type = tflite::TensorType_FLOAT32; + } + } + } + if (!found_image_to_tensor_calculator) { + return absl::Status(CreateStatusWithPayload( + absl::StatusCode::kNotFound, + absl::StrCat("The graph doesn't contain ImageToTensorCalculator."))); + } + return image_tensor_specs; + } protected: // A synchronous method to process single image inputs. diff --git a/mediapipe/tasks/cc/vision/object_detector/object_detector.h b/mediapipe/tasks/cc/vision/object_detector/object_detector.h index 249a2ebf..6113ed10 100644 --- a/mediapipe/tasks/cc/vision/object_detector/object_detector.h +++ b/mediapipe/tasks/cc/vision/object_detector/object_detector.h @@ -130,7 +130,7 @@ struct ObjectDetectorOptions { // // [1]: // https://github.com/google/mediapipe/blob/6cdc6443b6a7ed662744e2a2ce2d58d9c83e6d6f/mediapipe/tasks/metadata/metadata_schema.fbs#L456 -class ObjectDetector : tasks::vision::core::BaseVisionTaskApi { +class ObjectDetector : public tasks::vision::core::BaseVisionTaskApi { public: using BaseVisionTaskApi::BaseVisionTaskApi; diff --git a/mediapipe/tasks/cc/vision/object_detector/object_detector_test.cc b/mediapipe/tasks/cc/vision/object_detector/object_detector_test.cc index 798e3f23..1b0dcf2f 100644 --- a/mediapipe/tasks/cc/vision/object_detector/object_detector_test.cc +++ b/mediapipe/tasks/cc/vision/object_detector/object_detector_test.cc @@ -39,6 +39,7 @@ limitations under the License. #include "mediapipe/tasks/cc/components/containers/rect.h" #include "mediapipe/tasks/cc/vision/core/image_processing_options.h" #include "mediapipe/tasks/cc/vision/core/running_mode.h" +#include "mediapipe/tasks/cc/vision/utils/image_tensor_specs.h" #include "mediapipe/tasks/cc/vision/utils/image_utils.h" #include "tensorflow/lite/c/common.h" #include "tensorflow/lite/core/api/op_resolver.h" @@ -298,6 +299,36 @@ TEST_F(CreateFromOptionsTest, FailsWithMissingCallbackInLiveStreamMode) { MediaPipeTasksStatus::kInvalidTaskGraphConfigError)))); } +TEST_F(CreateFromOptionsTest, InputTensorSpecsForMobileSsdModel) { + auto options = std::make_unique(); + options->base_options.model_asset_path = + JoinPath("./", kTestDataDirectory, kMobileSsdWithMetadata); + MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr object_detector, + ObjectDetector::Create(std::move(options))); + MP_ASSERT_OK_AND_ASSIGN(auto image_tensor_specs, + object_detector->GetInputImageTensorSpecs()); + EXPECT_EQ(image_tensor_specs.image_width, 300); + EXPECT_EQ(image_tensor_specs.image_height, 300); + EXPECT_EQ(image_tensor_specs.color_space, + tflite::ColorSpaceType::ColorSpaceType_RGB); + EXPECT_EQ(image_tensor_specs.tensor_type, tflite::TensorType_UINT8); +} + +TEST_F(CreateFromOptionsTest, InputTensorSpecsForEfficientDetModel) { + auto options = std::make_unique(); + options->base_options.model_asset_path = + JoinPath("./", kTestDataDirectory, kEfficientDetWithMetadata); + MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr object_detector, + ObjectDetector::Create(std::move(options))); + MP_ASSERT_OK_AND_ASSIGN(auto image_tensor_specs, + object_detector->GetInputImageTensorSpecs()); + EXPECT_EQ(image_tensor_specs.image_width, 320); + EXPECT_EQ(image_tensor_specs.image_height, 320); + EXPECT_EQ(image_tensor_specs.color_space, + tflite::ColorSpaceType::ColorSpaceType_RGB); + EXPECT_EQ(image_tensor_specs.tensor_type, tflite::TensorType_UINT8); +} + // TODO: Add NumThreadsTest back after having an // "acceleration configuration" field in the ObjectDetectorOptions. diff --git a/mediapipe/tasks/cc/vision/utils/image_tensor_specs.h b/mediapipe/tasks/cc/vision/utils/image_tensor_specs.h index 0c0d4c81..f1ad2a80 100644 --- a/mediapipe/tasks/cc/vision/utils/image_tensor_specs.h +++ b/mediapipe/tasks/cc/vision/utils/image_tensor_specs.h @@ -94,7 +94,7 @@ absl::StatusOr BuildInputImageTensorSpecs( // Build ImageTensorSpec from model resources. The tflite model must contain // single subgraph with single input tensor. absl::StatusOr BuildInputImageTensorSpecs( - const core::ModelResources& model_resources); + const tasks::core::ModelResources& model_resources); } // namespace vision } // namespace tasks