diff --git a/Dockerfile b/Dockerfile index 4d6c68e7..3df22cc0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -30,6 +30,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ git \ wget \ unzip \ + nodejs \ + npm \ python3-dev \ python3-opencv \ python3-pip \ diff --git a/WORKSPACE b/WORKSPACE index 146916c5..5a47cf6b 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -172,6 +172,10 @@ http_archive( urls = [ "https://github.com/google/sentencepiece/archive/1.0.0.zip", ], + patches = [ + "//third_party:com_google_sentencepiece_no_gflag_no_gtest.diff", + ], + patch_args = ["-p1"], repo_mapping = {"@com_google_glog" : "@com_github_glog_glog"}, ) diff --git a/docs/BUILD b/docs/BUILD new file mode 100644 index 00000000..cb8794da --- /dev/null +++ b/docs/BUILD @@ -0,0 +1,14 @@ +# Placeholder for internal Python strict binary compatibility macro. + +py_binary( + name = "build_py_api_docs", + srcs = ["build_py_api_docs.py"], + deps = [ + "//mediapipe", + "//third_party/py/absl:app", + "//third_party/py/absl/flags", + "//third_party/py/tensorflow_docs", + "//third_party/py/tensorflow_docs/api_generator:generate_lib", + "//third_party/py/tensorflow_docs/api_generator:public_api", + ], +) diff --git a/docs/build_py_api_docs.py b/docs/build_py_api_docs.py new file mode 100644 index 00000000..9911d073 --- /dev/null +++ b/docs/build_py_api_docs.py @@ -0,0 +1,85 @@ +# Copyright 2022 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +r"""MediaPipe reference docs generation script. + +This script generates API reference docs for the `mediapipe` PIP package. + +$> pip install -U git+https://github.com/tensorflow/docs mediapipe +$> python build_py_api_docs.py +""" + +import os + +from absl import app +from absl import flags + +from tensorflow_docs.api_generator import generate_lib +from tensorflow_docs.api_generator import public_api + +try: + # mediapipe has not been set up to work with bazel yet, so catch & report. + import mediapipe # pytype: disable=import-error +except ImportError as e: + raise ImportError('Please `pip install mediapipe`.') from e + + +PROJECT_SHORT_NAME = 'mp' +PROJECT_FULL_NAME = 'MediaPipe' + +_OUTPUT_DIR = flags.DEFINE_string( + 'output_dir', + default='/tmp/generated_docs', + help='Where to write the resulting docs.') + +_URL_PREFIX = flags.DEFINE_string( + 'code_url_prefix', + 'https://github.com/google/mediapipe/tree/master/mediapipe', + 'The url prefix for links to code.') + +_SEARCH_HINTS = flags.DEFINE_bool( + 'search_hints', True, + 'Include metadata search hints in the generated files') + +_SITE_PATH = flags.DEFINE_string('site_path', '/mediapipe/api_docs/python', + 'Path prefix in the _toc.yaml') + + +def gen_api_docs(): + """Generates API docs for the mediapipe package.""" + + doc_generator = generate_lib.DocGenerator( + root_title=PROJECT_FULL_NAME, + py_modules=[(PROJECT_SHORT_NAME, mediapipe)], + base_dir=os.path.dirname(mediapipe.__file__), + code_url_prefix=_URL_PREFIX.value, + search_hints=_SEARCH_HINTS.value, + site_path=_SITE_PATH.value, + # This callback ensures that docs are only generated for objects that + # are explicitly imported in your __init__.py files. There are other + # options but this is a good starting point. + callbacks=[public_api.explicit_package_contents_filter], + ) + + doc_generator.build(_OUTPUT_DIR.value) + + print('Docs output to:', _OUTPUT_DIR.value) + + +def main(_): + gen_api_docs() + + +if __name__ == '__main__': + app.run(main) diff --git a/mediapipe/calculators/core/BUILD b/mediapipe/calculators/core/BUILD index 8fbc9982..74398be4 100644 --- a/mediapipe/calculators/core/BUILD +++ b/mediapipe/calculators/core/BUILD @@ -222,10 +222,10 @@ cc_library( "//mediapipe/framework:calculator_contract", "//mediapipe/framework:calculator_framework", "//mediapipe/framework:collection_item_id", - "//mediapipe/framework:packet", "//mediapipe/framework/formats:classification_cc_proto", "//mediapipe/framework/formats:detection_cc_proto", "//mediapipe/framework/formats:landmark_cc_proto", + "//mediapipe/framework/formats:matrix", "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/framework/port:integral_types", "//mediapipe/framework/port:ret_check", @@ -328,6 +328,7 @@ cc_library( ":concatenate_vector_calculator_cc_proto", "//mediapipe/framework:calculator_framework", "//mediapipe/framework/api2:node", + "//mediapipe/framework/formats:classification_cc_proto", "//mediapipe/framework/formats:landmark_cc_proto", "//mediapipe/framework/port:ret_check", "//mediapipe/framework/port:status", @@ -344,6 +345,7 @@ cc_test( "//mediapipe/framework:calculator_framework", "//mediapipe/framework:calculator_runner", "//mediapipe/framework:timestamp", + "//mediapipe/framework/formats:classification_cc_proto", "//mediapipe/framework/formats:landmark_cc_proto", "//mediapipe/framework/port:gtest_main", "//mediapipe/framework/port:parse_text_proto", diff --git a/mediapipe/calculators/core/bypass_calculator_test.cc b/mediapipe/calculators/core/bypass_calculator_test.cc index 4d1cd8f7..224742a1 100644 --- a/mediapipe/calculators/core/bypass_calculator_test.cc +++ b/mediapipe/calculators/core/bypass_calculator_test.cc @@ -75,6 +75,7 @@ constexpr char kTestGraphConfig2[] = R"pb( output_stream: "FEDERATED_GAZE_OUTPUT:federated_gaze_output" options { [mediapipe.SwitchContainerOptions.ext] { + async_selection: true contained_node: { calculator: "AppearancesPassThroughSubgraph" } } } @@ -101,6 +102,7 @@ constexpr char kTestGraphConfig3[] = R"pb( output_stream: "FEDERATED_GAZE_OUTPUT:federated_gaze_output" options { [mediapipe.SwitchContainerOptions.ext] { + async_selection: true contained_node: { calculator: "BypassCalculator" node_options: { diff --git a/mediapipe/calculators/core/concatenate_proto_list_calculator.cc b/mediapipe/calculators/core/concatenate_proto_list_calculator.cc index 9dd0dfd9..6c58e111 100644 --- a/mediapipe/calculators/core/concatenate_proto_list_calculator.cc +++ b/mediapipe/calculators/core/concatenate_proto_list_calculator.cc @@ -18,6 +18,7 @@ #include "mediapipe/calculators/core/concatenate_vector_calculator.pb.h" #include "mediapipe/framework/api2/node.h" #include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/classification.pb.h" #include "mediapipe/framework/formats/landmark.pb.h" #include "mediapipe/framework/port/canonical_errors.h" #include "mediapipe/framework/port/ret_check.h" @@ -111,6 +112,22 @@ class ConcatenateLandmarkListCalculator }; MEDIAPIPE_REGISTER_NODE(ConcatenateLandmarkListCalculator); +class ConcatenateClassificationListCalculator + : public ConcatenateListsCalculator { + protected: + int ListSize(const ClassificationList& list) const override { + return list.classification_size(); + } + const Classification GetItem(const ClassificationList& list, + int idx) const override { + return list.classification(idx); + } + Classification* AddItem(ClassificationList& list) const override { + return list.add_classification(); + } +}; +MEDIAPIPE_REGISTER_NODE(ConcatenateClassificationListCalculator); + } // namespace api2 } // namespace mediapipe diff --git a/mediapipe/calculators/core/concatenate_proto_list_calculator_test.cc b/mediapipe/calculators/core/concatenate_proto_list_calculator_test.cc index fd116ece..2167cd9d 100644 --- a/mediapipe/calculators/core/concatenate_proto_list_calculator_test.cc +++ b/mediapipe/calculators/core/concatenate_proto_list_calculator_test.cc @@ -18,6 +18,7 @@ #include "mediapipe/framework/calculator_framework.h" #include "mediapipe/framework/calculator_runner.h" +#include "mediapipe/framework/formats/classification.pb.h" #include "mediapipe/framework/formats/landmark.pb.h" #include "mediapipe/framework/port/gmock.h" #include "mediapipe/framework/port/gtest.h" @@ -70,6 +71,16 @@ void AddInputLandmarkLists( } } +void AddInputClassificationLists( + const std::vector& input_classifications_vec, + int64 timestamp, CalculatorRunner* runner) { + for (int i = 0; i < input_classifications_vec.size(); ++i) { + runner->MutableInputs()->Index(i).packets.push_back( + MakePacket(input_classifications_vec[i]) + .At(Timestamp(timestamp))); + } +} + TEST(ConcatenateNormalizedLandmarkListCalculatorTest, EmptyVectorInputs) { CalculatorRunner runner("ConcatenateNormalizedLandmarkListCalculator", /*options_string=*/"", /*num_inputs=*/3, @@ -181,4 +192,39 @@ TEST(ConcatenateNormalizedLandmarkListCalculatorTest, OneEmptyStreamNoOutput) { EXPECT_EQ(0, outputs.size()); } +TEST(ConcatenateClassificationListCalculatorTest, OneTimestamp) { + CalculatorRunner runner("ConcatenateClassificationListCalculator", + /*options_string=*/ + "[mediapipe.ConcatenateVectorCalculatorOptions.ext]: " + "{only_emit_if_all_present: true}", + /*num_inputs=*/2, + /*num_outputs=*/1, /*num_side_packets=*/0); + + auto input_0 = ParseTextProtoOrDie(R"pb( + classification: { index: 0 score: 0.2 label: "test_0" } + classification: { index: 1 score: 0.3 label: "test_1" } + classification: { index: 2 score: 0.4 label: "test_2" } + )pb"); + auto input_1 = ParseTextProtoOrDie(R"pb( + classification: { index: 3 score: 0.2 label: "test_3" } + classification: { index: 4 score: 0.3 label: "test_4" } + )pb"); + std::vector inputs = {input_0, input_1}; + AddInputClassificationLists(inputs, /*timestamp=*/1, &runner); + MP_ASSERT_OK(runner.Run()); + + const std::vector& outputs = runner.Outputs().Index(0).packets; + EXPECT_EQ(1, outputs.size()); + EXPECT_EQ(Timestamp(1), outputs[0].Timestamp()); + auto result = outputs[0].Get(); + EXPECT_THAT(ParseTextProtoOrDie(R"pb( + classification: { index: 0 score: 0.2 label: "test_0" } + classification: { index: 1 score: 0.3 label: "test_1" } + classification: { index: 2 score: 0.4 label: "test_2" } + classification: { index: 3 score: 0.2 label: "test_3" } + classification: { index: 4 score: 0.3 label: "test_4" } + )pb"), + EqualsProto(result)); +} + } // namespace mediapipe diff --git a/mediapipe/calculators/core/end_loop_calculator.cc b/mediapipe/calculators/core/end_loop_calculator.cc index b321f427..45cd8a9f 100644 --- a/mediapipe/calculators/core/end_loop_calculator.cc +++ b/mediapipe/calculators/core/end_loop_calculator.cc @@ -19,6 +19,7 @@ #include "mediapipe/framework/formats/classification.pb.h" #include "mediapipe/framework/formats/detection.pb.h" #include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/formats/matrix.h" #include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/util/render_data.pb.h" #include "tensorflow/lite/interpreter.h" @@ -58,4 +59,7 @@ typedef EndLoopCalculator> EndLoopDetectionCalculator; REGISTER_CALCULATOR(EndLoopDetectionCalculator); +typedef EndLoopCalculator> EndLoopMatrixCalculator; +REGISTER_CALCULATOR(EndLoopMatrixCalculator); + } // namespace mediapipe diff --git a/mediapipe/calculators/core/end_loop_calculator.h b/mediapipe/calculators/core/end_loop_calculator.h index e40301e8..9f56657d 100644 --- a/mediapipe/calculators/core/end_loop_calculator.h +++ b/mediapipe/calculators/core/end_loop_calculator.h @@ -50,7 +50,7 @@ namespace mediapipe { // calculator: "EndLoopWithOutputCalculator" // input_stream: "ITEM:output_of_loop_body" # ItemU @loop_internal_ts // input_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts -// output_stream: "OUTPUT:aggregated_result" # IterableU @ext_ts +// output_stream: "ITERABLE:aggregated_result" # IterableU @ext_ts // } template class EndLoopCalculator : public CalculatorBase { diff --git a/mediapipe/calculators/tensor/BUILD b/mediapipe/calculators/tensor/BUILD index 654b7135..92e786b6 100644 --- a/mediapipe/calculators/tensor/BUILD +++ b/mediapipe/calculators/tensor/BUILD @@ -109,6 +109,56 @@ cc_test( ], ) +mediapipe_proto_library( + name = "tensors_to_audio_calculator_proto", + srcs = ["tensors_to_audio_calculator.proto"], + visibility = [ + "//mediapipe/framework:mediapipe_internal", + ], + deps = [ + "//mediapipe/framework:calculator_options_proto", + "//mediapipe/framework:calculator_proto", + ], +) + +cc_library( + name = "tensors_to_audio_calculator", + srcs = ["tensors_to_audio_calculator.cc"], + visibility = [ + "//mediapipe/framework:mediapipe_internal", + ], + deps = [ + ":tensors_to_audio_calculator_cc_proto", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework/api2:node", + "//mediapipe/framework/formats:matrix", + "//mediapipe/framework/formats:tensor", + "//mediapipe/framework/port:ret_check", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/status", + "@com_google_audio_tools//audio/dsp:window_functions", + "@pffft", + ], + alwayslink = 1, +) + +cc_test( + name = "tensors_to_audio_calculator_test", + srcs = ["tensors_to_audio_calculator_test.cc"], + deps = [ + ":audio_to_tensor_calculator", + ":audio_to_tensor_calculator_cc_proto", + ":tensors_to_audio_calculator", + ":tensors_to_audio_calculator_cc_proto", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework/formats:matrix", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:parse_text_proto", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + ], +) + mediapipe_proto_library( name = "feedback_tensors_calculator_proto", srcs = ["feedback_tensors_calculator.proto"], @@ -253,6 +303,26 @@ cc_library( alwayslink = 1, ) +cc_test( + name = "regex_preprocessor_calculator_test", + srcs = ["regex_preprocessor_calculator_test.cc"], + data = ["//mediapipe/tasks/testdata/text:text_classifier_models"], + linkopts = ["-ldl"], + deps = [ + ":regex_preprocessor_calculator", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework/formats:tensor", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:parse_text_proto", + "//mediapipe/framework/tool:sink", + "//mediapipe/tasks/cc/core:utils", + "//mediapipe/tasks/cc/metadata:metadata_extractor", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + ], +) + cc_library( name = "text_to_tensor_calculator", srcs = ["text_to_tensor_calculator.cc"], @@ -304,6 +374,28 @@ cc_library( "@com_google_absl//absl/status", "@com_google_absl//absl/strings", ], + alwayslink = 1, +) + +cc_test( + name = "universal_sentence_encoder_preprocessor_calculator_test", + srcs = ["universal_sentence_encoder_preprocessor_calculator_test.cc"], + data = ["//mediapipe/tasks/testdata/text:universal_sentence_encoder_qa"], + deps = [ + ":universal_sentence_encoder_preprocessor_calculator", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework:packet", + "//mediapipe/framework/formats:tensor", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:parse_text_proto", + "//mediapipe/framework/port:status", + "//mediapipe/framework/tool:options_map", + "//mediapipe/tasks/cc/core:utils", + "//mediapipe/tasks/cc/metadata:metadata_extractor", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + ], ) mediapipe_proto_library( @@ -438,6 +530,7 @@ cc_library( }), visibility = ["//visibility:public"], deps = [ + "//mediapipe/framework:calculator_context", "//mediapipe/framework/formats:tensor", "@com_google_absl//absl/status:statusor", ], @@ -458,6 +551,7 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":inference_runner", + "//mediapipe/framework:mediapipe_profiling", "//mediapipe/framework/api2:packet", "//mediapipe/framework/formats:tensor", "//mediapipe/framework/port:ret_check", @@ -1200,13 +1294,30 @@ cc_library( name = "image_to_tensor_utils", srcs = ["image_to_tensor_utils.cc"], hdrs = ["image_to_tensor_utils.h"], + copts = select({ + "//mediapipe:apple": [ + "-x objective-c++", + "-fobjc-arc", # enable reference-counting + ], + "//conditions:default": [], + }), visibility = ["//visibility:public"], deps = [ + ":image_to_tensor_calculator_cc_proto", + "@com_google_absl//absl/status", + "@com_google_absl//absl/types:optional", + "//mediapipe/framework/api2:packet", + "//mediapipe/framework/api2:port", + "//mediapipe/framework/formats:image", "//mediapipe/framework/formats:rect_cc_proto", + "//mediapipe/framework/formats:tensor", "//mediapipe/framework/port:ret_check", "//mediapipe/framework/port:statusor", - "@com_google_absl//absl/types:optional", - ], + "//mediapipe/gpu:gpu_origin_cc_proto", + ] + select({ + "//mediapipe/gpu:disable_gpu": [], + "//conditions:default": ["//mediapipe/gpu:gpu_buffer"], + }), ) cc_test( @@ -1216,6 +1327,8 @@ cc_test( ":image_to_tensor_utils", "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:parse_text_proto", + "//mediapipe/framework/port:status", ], ) diff --git a/mediapipe/calculators/tensor/audio_to_tensor_calculator.cc b/mediapipe/calculators/tensor/audio_to_tensor_calculator.cc index 59c12919..d0513518 100644 --- a/mediapipe/calculators/tensor/audio_to_tensor_calculator.cc +++ b/mediapipe/calculators/tensor/audio_to_tensor_calculator.cc @@ -133,7 +133,7 @@ bool IsValidFftSize(int size) { // invocation. In the non-streaming mode, the vector contains all of the // output timestamps for an input audio buffer. // DC_AND_NYQUIST - std::pair @Optional. -// A pair of dc component and nyquest component. Only can be connected when +// A pair of dc component and nyquist component. Only can be connected when // the calculator performs fft (the fft_size is set in the calculator // options). // diff --git a/mediapipe/calculators/tensor/image_to_tensor_calculator.cc b/mediapipe/calculators/tensor/image_to_tensor_calculator.cc index ec7d4afa..5af4cdb6 100644 --- a/mediapipe/calculators/tensor/image_to_tensor_calculator.cc +++ b/mediapipe/calculators/tensor/image_to_tensor_calculator.cc @@ -54,13 +54,6 @@ namespace mediapipe { namespace api2 { -#if MEDIAPIPE_DISABLE_GPU -// Just a placeholder to not have to depend on mediapipe::GpuBuffer. -using GpuBuffer = AnyType; -#else -using GpuBuffer = mediapipe::GpuBuffer; -#endif // MEDIAPIPE_DISABLE_GPU - // Converts image into Tensor, possibly with cropping, resizing and // normalization, according to specified inputs and options. // @@ -141,42 +134,7 @@ class ImageToTensorCalculator : public Node { const auto& options = cc->Options(); - RET_CHECK(options.has_output_tensor_float_range() || - options.has_output_tensor_int_range() || - options.has_output_tensor_uint_range()) - << "Output tensor range is required."; - if (options.has_output_tensor_float_range()) { - RET_CHECK_LT(options.output_tensor_float_range().min(), - options.output_tensor_float_range().max()) - << "Valid output float tensor range is required."; - } - if (options.has_output_tensor_uint_range()) { - RET_CHECK_LT(options.output_tensor_uint_range().min(), - options.output_tensor_uint_range().max()) - << "Valid output uint tensor range is required."; - RET_CHECK_GE(options.output_tensor_uint_range().min(), 0) - << "The minimum of the output uint tensor range must be " - "non-negative."; - RET_CHECK_LE(options.output_tensor_uint_range().max(), 255) - << "The maximum of the output uint tensor range must be less than or " - "equal to 255."; - } - if (options.has_output_tensor_int_range()) { - RET_CHECK_LT(options.output_tensor_int_range().min(), - options.output_tensor_int_range().max()) - << "Valid output int tensor range is required."; - RET_CHECK_GE(options.output_tensor_int_range().min(), -128) - << "The minimum of the output int tensor range must be greater than " - "or equal to -128."; - RET_CHECK_LE(options.output_tensor_int_range().max(), 127) - << "The maximum of the output int tensor range must be less than or " - "equal to 127."; - } - RET_CHECK_GT(options.output_tensor_width(), 0) - << "Valid output tensor width is required."; - RET_CHECK_GT(options.output_tensor_height(), 0) - << "Valid output tensor height is required."; - + RET_CHECK_OK(ValidateOptionOutputDims(options)); RET_CHECK(kIn(cc).IsConnected() ^ kInGpu(cc).IsConnected()) << "One and only one of IMAGE and IMAGE_GPU input is expected."; @@ -198,21 +156,7 @@ class ImageToTensorCalculator : public Node { absl::Status Open(CalculatorContext* cc) { options_ = cc->Options(); - output_width_ = options_.output_tensor_width(); - output_height_ = options_.output_tensor_height(); - is_float_output_ = options_.has_output_tensor_float_range(); - if (options_.has_output_tensor_uint_range()) { - range_min_ = - static_cast(options_.output_tensor_uint_range().min()); - range_max_ = - static_cast(options_.output_tensor_uint_range().max()); - } else if (options_.has_output_tensor_int_range()) { - range_min_ = static_cast(options_.output_tensor_int_range().min()); - range_max_ = static_cast(options_.output_tensor_int_range().max()); - } else { - range_min_ = options_.output_tensor_float_range().min(); - range_max_ = options_.output_tensor_float_range().max(); - } + params_ = GetOutputTensorParams(options_); return absl::OkStatus(); } @@ -242,7 +186,13 @@ class ImageToTensorCalculator : public Node { } } - ASSIGN_OR_RETURN(auto image, GetInputImage(cc)); +#if MEDIAPIPE_DISABLE_GPU + ASSIGN_OR_RETURN(auto image, GetInputImage(kIn(cc))); +#else + const bool is_input_gpu = kInGpu(cc).IsConnected(); + ASSIGN_OR_RETURN(auto image, is_input_gpu ? GetInputImage(kInGpu(cc)) + : GetInputImage(kIn(cc))); +#endif // MEDIAPIPE_DISABLE_GPU RotatedRect roi = GetRoi(image->width(), image->height(), norm_rect); ASSIGN_OR_RETURN(auto padding, PadRoi(options_.output_tensor_width(), @@ -263,11 +213,13 @@ class ImageToTensorCalculator : public Node { MP_RETURN_IF_ERROR(InitConverterIfNecessary(cc, *image.get())); Tensor::ElementType output_tensor_type = - GetOutputTensorType(image->UsesGpu()); - Tensor tensor(output_tensor_type, {1, output_height_, output_width_, - GetNumOutputChannels(*image)}); + GetOutputTensorType(image->UsesGpu(), params_); + Tensor tensor(output_tensor_type, + {1, params_.output_height, params_.output_width, + GetNumOutputChannels(*image)}); MP_RETURN_IF_ERROR((image->UsesGpu() ? gpu_converter_ : cpu_converter_) - ->Convert(*image, roi, range_min_, range_max_, + ->Convert(*image, roi, params_.range_min, + params_.range_max, /*tensor_buffer_offset=*/0, tensor)); auto result = std::make_unique>(); @@ -278,81 +230,11 @@ class ImageToTensorCalculator : public Node { } private: - bool DoesGpuInputStartAtBottom() { - return options_.gpu_origin() != mediapipe::GpuOrigin_Mode_TOP_LEFT; - } - - BorderMode GetBorderMode() { - switch (options_.border_mode()) { - case mediapipe:: - ImageToTensorCalculatorOptions_BorderMode_BORDER_UNSPECIFIED: - return BorderMode::kReplicate; - case mediapipe::ImageToTensorCalculatorOptions_BorderMode_BORDER_ZERO: - return BorderMode::kZero; - case mediapipe:: - ImageToTensorCalculatorOptions_BorderMode_BORDER_REPLICATE: - return BorderMode::kReplicate; - } - } - - Tensor::ElementType GetOutputTensorType(bool uses_gpu) { - if (!uses_gpu) { - if (is_float_output_) { - return Tensor::ElementType::kFloat32; - } - if (range_min_ < 0) { - return Tensor::ElementType::kInt8; - } else { - return Tensor::ElementType::kUInt8; - } - } - // Always use float32 when GPU is enabled. - return Tensor::ElementType::kFloat32; - } - - int GetNumOutputChannels(const Image& image) { -#if !MEDIAPIPE_DISABLE_GPU -#if MEDIAPIPE_METAL_ENABLED - if (image.UsesGpu()) { - return 4; - } -#endif // MEDIAPIPE_METAL_ENABLED -#endif // !MEDIAPIPE_DISABLE_GPU - // All of the processors except for Metal expect 3 channels. - return 3; - } - - absl::StatusOr> GetInputImage( - CalculatorContext* cc) { - if (kIn(cc).IsConnected()) { - const auto& packet = kIn(cc).packet(); - return kIn(cc).Visit( - [&packet](const mediapipe::Image&) { - return SharedPtrWithPacket(packet); - }, - [&packet](const mediapipe::ImageFrame&) { - return std::make_shared( - std::const_pointer_cast( - SharedPtrWithPacket(packet))); - }); - } else { // if (kInGpu(cc).IsConnected()) -#if !MEDIAPIPE_DISABLE_GPU - const GpuBuffer& input = *kInGpu(cc); - // A shallow copy is okay since the resulting 'image' object is local in - // Process(), and thus never outlives 'input'. - return std::make_shared(input); -#else - return absl::UnimplementedError( - "GPU processing is disabled in build flags"); -#endif // !MEDIAPIPE_DISABLE_GPU - } - } - absl::Status InitConverterIfNecessary(CalculatorContext* cc, const Image& image) { // Lazy initialization of the GPU or CPU converter. if (image.UsesGpu()) { - if (!is_float_output_) { + if (!params_.is_float_output) { return absl::UnimplementedError( "ImageToTensorConverter for the input GPU image currently doesn't " "support quantization."); @@ -360,18 +242,20 @@ class ImageToTensorCalculator : public Node { if (!gpu_converter_) { #if !MEDIAPIPE_DISABLE_GPU #if MEDIAPIPE_METAL_ENABLED - ASSIGN_OR_RETURN(gpu_converter_, - CreateMetalConverter(cc, GetBorderMode())); + ASSIGN_OR_RETURN( + gpu_converter_, + CreateMetalConverter(cc, GetBorderMode(options_.border_mode()))); #elif MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31 ASSIGN_OR_RETURN(gpu_converter_, CreateImageToGlBufferTensorConverter( - cc, DoesGpuInputStartAtBottom(), GetBorderMode())); + cc, DoesGpuInputStartAtBottom(options_), + GetBorderMode(options_.border_mode()))); #else if (!gpu_converter_) { - ASSIGN_OR_RETURN( - gpu_converter_, - CreateImageToGlTextureTensorConverter( - cc, DoesGpuInputStartAtBottom(), GetBorderMode())); + ASSIGN_OR_RETURN(gpu_converter_, + CreateImageToGlTextureTensorConverter( + cc, DoesGpuInputStartAtBottom(options_), + GetBorderMode(options_.border_mode()))); } if (!gpu_converter_) { return absl::UnimplementedError( @@ -383,10 +267,10 @@ class ImageToTensorCalculator : public Node { } else { if (!cpu_converter_) { #if !MEDIAPIPE_DISABLE_OPENCV - ASSIGN_OR_RETURN( - cpu_converter_, - CreateOpenCvConverter(cc, GetBorderMode(), - GetOutputTensorType(/*uses_gpu=*/false))); + ASSIGN_OR_RETURN(cpu_converter_, + CreateOpenCvConverter( + cc, GetBorderMode(options_.border_mode()), + GetOutputTensorType(/*uses_gpu=*/false, params_))); #else LOG(FATAL) << "Cannot create image to tensor opencv converter since " "MEDIAPIPE_DISABLE_OPENCV is defined."; @@ -399,11 +283,7 @@ class ImageToTensorCalculator : public Node { std::unique_ptr gpu_converter_; std::unique_ptr cpu_converter_; mediapipe::ImageToTensorCalculatorOptions options_; - int output_width_ = 0; - int output_height_ = 0; - bool is_float_output_ = false; - float range_min_ = 0.0f; - float range_max_ = 1.0f; + OutputTensorParams params_; }; MEDIAPIPE_REGISTER_NODE(ImageToTensorCalculator); diff --git a/mediapipe/calculators/tensor/image_to_tensor_converter.h b/mediapipe/calculators/tensor/image_to_tensor_converter.h index 870ebc30..94d47c09 100644 --- a/mediapipe/calculators/tensor/image_to_tensor_converter.h +++ b/mediapipe/calculators/tensor/image_to_tensor_converter.h @@ -27,12 +27,6 @@ struct Size { int height; }; -// Pixel extrapolation method. -// When converting image to tensor it may happen that tensor needs to read -// pixels outside image boundaries. Border mode helps to specify how such pixels -// will be calculated. -enum class BorderMode { kZero, kReplicate }; - // Converts image to tensor. class ImageToTensorConverter { public: diff --git a/mediapipe/calculators/tensor/image_to_tensor_converter_gl_buffer.cc b/mediapipe/calculators/tensor/image_to_tensor_converter_gl_buffer.cc index 5671d7b4..14de410f 100644 --- a/mediapipe/calculators/tensor/image_to_tensor_converter_gl_buffer.cc +++ b/mediapipe/calculators/tensor/image_to_tensor_converter_gl_buffer.cc @@ -270,10 +270,10 @@ class GlProcessor : public ImageToTensorConverter { Tensor& output_tensor) override { if (input.format() != mediapipe::GpuBufferFormat::kBGRA32 && input.format() != mediapipe::GpuBufferFormat::kRGBAHalf64 && - input.format() != mediapipe::GpuBufferFormat::kRGBAFloat128) { + input.format() != mediapipe::GpuBufferFormat::kRGBAFloat128 && + input.format() != mediapipe::GpuBufferFormat::kRGB24) { return InvalidArgumentError(absl::StrCat( - "Only 4-channel texture input formats are supported, passed format: ", - static_cast(input.format()))); + "Unsupported format: ", static_cast(input.format()))); } const auto& output_shape = output_tensor.shape(); MP_RETURN_IF_ERROR(ValidateTensorShape(output_shape)); @@ -281,12 +281,13 @@ class GlProcessor : public ImageToTensorConverter { MP_RETURN_IF_ERROR(gl_helper_.RunInGlContext( [this, &output_tensor, &input, &roi, &output_shape, range_min, range_max, tensor_buffer_offset]() -> absl::Status { - constexpr int kRgbaNumChannels = 4; + const int input_num_channels = input.channels(); auto source_texture = gl_helper_.CreateSourceTexture(input); tflite::gpu::gl::GlTexture input_texture( - GL_TEXTURE_2D, source_texture.name(), GL_RGBA, + GL_TEXTURE_2D, source_texture.name(), + input_num_channels == 4 ? GL_RGB : GL_RGBA, source_texture.width() * source_texture.height() * - kRgbaNumChannels * sizeof(uint8_t), + input_num_channels * sizeof(uint8_t), /*layer=*/0, /*owned=*/false); diff --git a/mediapipe/calculators/tensor/image_to_tensor_converter_gl_texture.cc b/mediapipe/calculators/tensor/image_to_tensor_converter_gl_texture.cc index 06dfd578..5efd3404 100644 --- a/mediapipe/calculators/tensor/image_to_tensor_converter_gl_texture.cc +++ b/mediapipe/calculators/tensor/image_to_tensor_converter_gl_texture.cc @@ -174,10 +174,10 @@ class GlProcessor : public ImageToTensorConverter { Tensor& output_tensor) override { if (input.format() != mediapipe::GpuBufferFormat::kBGRA32 && input.format() != mediapipe::GpuBufferFormat::kRGBAHalf64 && - input.format() != mediapipe::GpuBufferFormat::kRGBAFloat128) { + input.format() != mediapipe::GpuBufferFormat::kRGBAFloat128 && + input.format() != mediapipe::GpuBufferFormat::kRGB24) { return InvalidArgumentError(absl::StrCat( - "Only 4-channel texture input formats are supported, passed format: ", - static_cast(input.format()))); + "Unsupported format: ", static_cast(input.format()))); } // TODO: support tensor_buffer_offset > 0 scenario. RET_CHECK_EQ(tensor_buffer_offset, 0) diff --git a/mediapipe/calculators/tensor/image_to_tensor_utils.cc b/mediapipe/calculators/tensor/image_to_tensor_utils.cc index 6b3bf08c..3f4c05d4 100644 --- a/mediapipe/calculators/tensor/image_to_tensor_utils.cc +++ b/mediapipe/calculators/tensor/image_to_tensor_utils.cc @@ -16,7 +16,9 @@ #include +#include "absl/status/status.h" #include "absl/types/optional.h" +#include "mediapipe/framework/api2/packet.h" #include "mediapipe/framework/port/ret_check.h" #include "mediapipe/framework/port/statusor.h" @@ -214,4 +216,68 @@ void GetTransposedRotatedSubRectToRectTransformMatrix( matrix[15] = 1.0f; } +BorderMode GetBorderMode( + const mediapipe::ImageToTensorCalculatorOptions::BorderMode& mode) { + switch (mode) { + case mediapipe:: + ImageToTensorCalculatorOptions_BorderMode_BORDER_UNSPECIFIED: + return BorderMode::kReplicate; + case mediapipe::ImageToTensorCalculatorOptions_BorderMode_BORDER_ZERO: + return BorderMode::kZero; + case mediapipe::ImageToTensorCalculatorOptions_BorderMode_BORDER_REPLICATE: + return BorderMode::kReplicate; + } +} + +Tensor::ElementType GetOutputTensorType(bool uses_gpu, + const OutputTensorParams& params) { + if (!uses_gpu) { + if (params.is_float_output) { + return Tensor::ElementType::kFloat32; + } + if (params.range_min < 0) { + return Tensor::ElementType::kInt8; + } else { + return Tensor::ElementType::kUInt8; + } + } + // Always use float32 when GPU is enabled. + return Tensor::ElementType::kFloat32; +} + +int GetNumOutputChannels(const mediapipe::Image& image) { +#if !MEDIAPIPE_DISABLE_GPU +#if MEDIAPIPE_METAL_ENABLED + if (image.UsesGpu()) { + return 4; + } +#endif // MEDIAPIPE_METAL_ENABLED +#endif // !MEDIAPIPE_DISABLE_GPU + // All of the processors except for Metal expect 3 channels. + return 3; +} + +absl::StatusOr> GetInputImage( + const api2::Packet>& + image_packet) { + return image_packet.Visit( + [&image_packet](const mediapipe::Image&) { + return SharedPtrWithPacket(image_packet); + }, + [&image_packet](const mediapipe::ImageFrame&) { + return std::make_shared( + std::const_pointer_cast( + SharedPtrWithPacket(image_packet))); + }); +} + +#if !MEDIAPIPE_DISABLE_GPU +absl::StatusOr> GetInputImage( + const api2::Packet& image_gpu_packet) { + // A shallow copy is okay since the resulting 'image' object is local in + // Process(), and thus never outlives 'input'. + return std::make_shared(image_gpu_packet.Get()); +} +#endif // !MEDIAPIPE_DISABLE_GPU + } // namespace mediapipe diff --git a/mediapipe/calculators/tensor/image_to_tensor_utils.h b/mediapipe/calculators/tensor/image_to_tensor_utils.h index f913875e..dc38ac7b 100644 --- a/mediapipe/calculators/tensor/image_to_tensor_utils.h +++ b/mediapipe/calculators/tensor/image_to_tensor_utils.h @@ -18,8 +18,18 @@ #include #include "absl/types/optional.h" +#include "mediapipe/calculators/tensor/image_to_tensor_calculator.pb.h" +#include "mediapipe/framework/api2/packet.h" +#include "mediapipe/framework/api2/port.h" +#include "mediapipe/framework/formats/image.h" #include "mediapipe/framework/formats/rect.pb.h" +#include "mediapipe/framework/formats/tensor.h" +#include "mediapipe/framework/port/ret_check.h" #include "mediapipe/framework/port/statusor.h" +#if !MEDIAPIPE_DISABLE_GPU +#include "mediapipe/gpu/gpu_buffer.h" +#endif // !MEDIAPIPE_DISABLE_GPU +#include "mediapipe/gpu/gpu_origin.pb.h" namespace mediapipe { @@ -31,6 +41,24 @@ struct RotatedRect { float rotation; }; +// Pixel extrapolation method. +// When converting image to tensor it may happen that tensor needs to read +// pixels outside image boundaries. Border mode helps to specify how such pixels +// will be calculated. +// TODO: Consider moving this to a separate border_mode.h file. +enum class BorderMode { kZero, kReplicate }; + +// Struct that host commonly accessed parameters used in the +// ImageTo[Batch]TensorCalculator. +struct OutputTensorParams { + int output_height; + int output_width; + int output_batch; + bool is_float_output; + float range_min; + float range_max; +}; + // Generates a new ROI or converts it from normalized rect. RotatedRect GetRoi(int input_width, int input_height, absl::optional norm_rect); @@ -95,6 +123,103 @@ void GetTransposedRotatedSubRectToRectTransformMatrix( const RotatedRect& sub_rect, int rect_width, int rect_height, bool flip_horizontaly, std::array* matrix); +// Validates the output dimensions set in the option proto. The input option +// proto is expected to have to following fields: +// output_tensor_float_range, output_tensor_int_range, output_tensor_uint_range +// output_tensor_width, output_tensor_height. +// See ImageToTensorCalculatorOptions for the description of each field. +template +absl::Status ValidateOptionOutputDims(const T& options) { + RET_CHECK(options.has_output_tensor_float_range() || + options.has_output_tensor_int_range() || + options.has_output_tensor_uint_range()) + << "Output tensor range is required."; + if (options.has_output_tensor_float_range()) { + RET_CHECK_LT(options.output_tensor_float_range().min(), + options.output_tensor_float_range().max()) + << "Valid output float tensor range is required."; + } + if (options.has_output_tensor_uint_range()) { + RET_CHECK_LT(options.output_tensor_uint_range().min(), + options.output_tensor_uint_range().max()) + << "Valid output uint tensor range is required."; + RET_CHECK_GE(options.output_tensor_uint_range().min(), 0) + << "The minimum of the output uint tensor range must be " + "non-negative."; + RET_CHECK_LE(options.output_tensor_uint_range().max(), 255) + << "The maximum of the output uint tensor range must be less than or " + "equal to 255."; + } + if (options.has_output_tensor_int_range()) { + RET_CHECK_LT(options.output_tensor_int_range().min(), + options.output_tensor_int_range().max()) + << "Valid output int tensor range is required."; + RET_CHECK_GE(options.output_tensor_int_range().min(), -128) + << "The minimum of the output int tensor range must be greater than " + "or equal to -128."; + RET_CHECK_LE(options.output_tensor_int_range().max(), 127) + << "The maximum of the output int tensor range must be less than or " + "equal to 127."; + } + RET_CHECK_GT(options.output_tensor_width(), 0) + << "Valid output tensor width is required."; + RET_CHECK_GT(options.output_tensor_height(), 0) + << "Valid output tensor height is required."; + return absl::OkStatus(); +} + +template +OutputTensorParams GetOutputTensorParams(const T& options) { + OutputTensorParams params; + if (options.has_output_tensor_uint_range()) { + params.range_min = + static_cast(options.output_tensor_uint_range().min()); + params.range_max = + static_cast(options.output_tensor_uint_range().max()); + } else if (options.has_output_tensor_int_range()) { + params.range_min = + static_cast(options.output_tensor_int_range().min()); + params.range_max = + static_cast(options.output_tensor_int_range().max()); + } else { + params.range_min = options.output_tensor_float_range().min(); + params.range_max = options.output_tensor_float_range().max(); + } + params.output_width = options.output_tensor_width(); + params.output_height = options.output_tensor_height(); + params.is_float_output = options.has_output_tensor_float_range(); + params.output_batch = 1; + return params; +} + +// Returns whether the GPU input format starts at the bottom. +template +bool DoesGpuInputStartAtBottom(const T& options) { + return options.gpu_origin() != mediapipe::GpuOrigin_Mode_TOP_LEFT; +} + +// Converts the BorderMode proto into struct. +BorderMode GetBorderMode( + const mediapipe::ImageToTensorCalculatorOptions::BorderMode& mode); + +// Gets the output tensor type. +Tensor::ElementType GetOutputTensorType(bool uses_gpu, + const OutputTensorParams& params); + +// Gets the number of output channels from the input Image format. +int GetNumOutputChannels(const mediapipe::Image& image); + +// Converts the packet that hosts different format (Image, ImageFrame, +// GpuBuffer) into the mediapipe::Image format. +absl::StatusOr> GetInputImage( + const api2::Packet>& + image_packet); + +#if !MEDIAPIPE_DISABLE_GPU +absl::StatusOr> GetInputImage( + const api2::Packet& image_gpu_packet); +#endif // !MEDIAPIPE_DISABLE_GPU + } // namespace mediapipe #endif // MEDIAPIPE_CALCULATORS_TENSOR_IMAGE_TO_TENSOR_UTILS_H_ diff --git a/mediapipe/calculators/tensor/image_to_tensor_utils_test.cc b/mediapipe/calculators/tensor/image_to_tensor_utils_test.cc index 814b4c34..9ba7d013 100644 --- a/mediapipe/calculators/tensor/image_to_tensor_utils_test.cc +++ b/mediapipe/calculators/tensor/image_to_tensor_utils_test.cc @@ -16,6 +16,8 @@ #include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/port/gtest.h" +#include "mediapipe/framework/port/parse_text_proto.h" +#include "mediapipe/framework/port/status_macros.h" #include "mediapipe/framework/port/status_matchers.h" namespace mediapipe { @@ -23,6 +25,7 @@ namespace { using ::testing::ElementsAre; using ::testing::ElementsAreArray; +using ::testing::HasSubstr; testing::Matcher EqRotatedRect(float width, float height, float center_x, float center_y, @@ -157,5 +160,95 @@ TEST(GetValueRangeTransformation, FloatToPixel) { EqValueTransformation(/*scale=*/255.0f, /*offset=*/0.0f)); } +constexpr char kValidFloatProto[] = R"( + output_tensor_float_range { min: 0.0 max: 1.0 } + output_tensor_width: 100 + output_tensor_height: 200 +)"; + +constexpr char kValidIntProto[] = R"( + output_tensor_float_range { min: 0 max: 255 } + output_tensor_width: 100 + output_tensor_height: 200 +)"; + +TEST(ValidateOptionOutputDims, ValidProtos) { + const auto float_options = + mediapipe::ParseTextProtoOrDie( + kValidFloatProto); + MP_EXPECT_OK(ValidateOptionOutputDims(float_options)); +} + +TEST(ValidateOptionOutputDims, EmptyProto) { + mediapipe::ImageToTensorCalculatorOptions options; + // No output tensor range set. + EXPECT_THAT(ValidateOptionOutputDims(options), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Output tensor range is required"))); + + // Invalid output float tensor range. + options.mutable_output_tensor_float_range()->set_min(1.0); + options.mutable_output_tensor_float_range()->set_max(0.0); + EXPECT_THAT( + ValidateOptionOutputDims(options), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Valid output float tensor range is required"))); + + // Output width/height is not set. + options.mutable_output_tensor_float_range()->set_min(0.0); + options.mutable_output_tensor_float_range()->set_max(1.0); + EXPECT_THAT(ValidateOptionOutputDims(options), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("Valid output tensor width is required"))); +} + +TEST(GetOutputTensorParams, SetValues) { + // Test int range with ImageToTensorCalculatorOptions. + const auto int_options = + mediapipe::ParseTextProtoOrDie( + kValidIntProto); + const auto params2 = GetOutputTensorParams(int_options); + EXPECT_EQ(params2.range_min, 0.0f); + EXPECT_EQ(params2.range_max, 255.0f); + EXPECT_EQ(params2.output_batch, 1); + EXPECT_EQ(params2.output_width, 100); + EXPECT_EQ(params2.output_height, 200); +} + +TEST(GetBorderMode, GetBorderMode) { + // Default to REPLICATE. + auto border_mode = + mediapipe::ImageToTensorCalculatorOptions_BorderMode_BORDER_UNSPECIFIED; + EXPECT_EQ(BorderMode::kReplicate, GetBorderMode(border_mode)); + + // Set to ZERO. + border_mode = + mediapipe::ImageToTensorCalculatorOptions_BorderMode_BORDER_ZERO; + EXPECT_EQ(BorderMode::kZero, GetBorderMode(border_mode)); +} + +TEST(GetOutputTensorType, GetOutputTensorType) { + OutputTensorParams params; + // Return float32 when GPU is enabled. + EXPECT_EQ(Tensor::ElementType::kFloat32, + GetOutputTensorType(/*uses_gpu=*/true, params)); + + // Return float32 when is_float_output is set to true. + params.is_float_output = true; + EXPECT_EQ(Tensor::ElementType::kFloat32, + GetOutputTensorType(/*uses_gpu=*/false, params)); + + // Return int8 when range_min is negative. + params.is_float_output = false; + params.range_min = -255.0f; + EXPECT_EQ(Tensor::ElementType::kInt8, + GetOutputTensorType(/*uses_gpu=*/false, params)); + + // Return 8int8 when range_min is non-negative. + params.range_min = 0.0f; + EXPECT_EQ(Tensor::ElementType::kUInt8, + GetOutputTensorType(/*uses_gpu=*/false, params)); +} + } // namespace } // namespace mediapipe diff --git a/mediapipe/calculators/tensor/inference_calculator_cpu.cc b/mediapipe/calculators/tensor/inference_calculator_cpu.cc index 2e90c7cc..79df9763 100644 --- a/mediapipe/calculators/tensor/inference_calculator_cpu.cc +++ b/mediapipe/calculators/tensor/inference_calculator_cpu.cc @@ -72,7 +72,7 @@ absl::Status InferenceCalculatorCpuImpl::Process(CalculatorContext* cc) { RET_CHECK(!input_tensors.empty()); ASSIGN_OR_RETURN(std::vector output_tensors, - inference_runner_->Run(input_tensors)); + inference_runner_->Run(cc, input_tensors)); kOutTensors(cc).Send(std::move(output_tensors)); return absl::OkStatus(); } diff --git a/mediapipe/calculators/tensor/inference_calculator_gl.cc b/mediapipe/calculators/tensor/inference_calculator_gl.cc index 1f3768ee..bd8eb3ee 100644 --- a/mediapipe/calculators/tensor/inference_calculator_gl.cc +++ b/mediapipe/calculators/tensor/inference_calculator_gl.cc @@ -26,6 +26,8 @@ #include "mediapipe/gpu/gl_calculator_helper.h" #include "tensorflow/lite/delegates/gpu/gl_delegate.h" +#define PERFETTO_TRACK_EVENT_NAMESPACE mediapipe + namespace mediapipe { namespace api2 { @@ -191,7 +193,7 @@ absl::Status InferenceCalculatorGlImpl::GpuInferenceRunner::Process( CalculatorContext* cc, const std::vector& input_tensors, std::vector& output_tensors) { return gpu_helper_.RunInGlContext( - [this, &input_tensors, &output_tensors]() -> absl::Status { + [this, cc, &input_tensors, &output_tensors]() -> absl::Status { // Explicitly copy input. for (int i = 0; i < input_tensors.size(); ++i) { glBindBuffer(GL_COPY_READ_BUFFER, @@ -203,7 +205,10 @@ absl::Status InferenceCalculatorGlImpl::GpuInferenceRunner::Process( } // Run inference. - RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk); + { + MEDIAPIPE_PROFILING(GPU_TASK_INVOKE, cc); + RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk); + } output_tensors.reserve(output_size_); for (int i = 0; i < output_size_; ++i) { diff --git a/mediapipe/calculators/tensor/inference_calculator_gl_advanced.cc b/mediapipe/calculators/tensor/inference_calculator_gl_advanced.cc index 7e11ee07..52359f7f 100644 --- a/mediapipe/calculators/tensor/inference_calculator_gl_advanced.cc +++ b/mediapipe/calculators/tensor/inference_calculator_gl_advanced.cc @@ -32,6 +32,8 @@ #include "mediapipe/util/android/file/base/helpers.h" #endif // MEDIAPIPE_ANDROID +#define PERFETTO_TRACK_EVENT_NAMESPACE mediapipe + namespace mediapipe { namespace api2 { @@ -83,7 +85,7 @@ class InferenceCalculatorGlAdvancedImpl const mediapipe::InferenceCalculatorOptions::Delegate& delegate); absl::StatusOr> Process( - const std::vector& input_tensors); + CalculatorContext* cc, const std::vector& input_tensors); absl::Status Close(); @@ -121,11 +123,11 @@ absl::Status InferenceCalculatorGlAdvancedImpl::GpuInferenceRunner::Init( absl::StatusOr> InferenceCalculatorGlAdvancedImpl::GpuInferenceRunner::Process( - const std::vector& input_tensors) { + CalculatorContext* cc, const std::vector& input_tensors) { std::vector output_tensors; MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext( - [this, &input_tensors, &output_tensors]() -> absl::Status { + [this, cc, &input_tensors, &output_tensors]() -> absl::Status { for (int i = 0; i < input_tensors.size(); ++i) { MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToInputTensor( input_tensors[i].GetOpenGlBufferReadView().name(), i)); @@ -138,7 +140,10 @@ InferenceCalculatorGlAdvancedImpl::GpuInferenceRunner::Process( output_tensors.back().GetOpenGlBufferWriteView().name(), i)); } // Run inference. - return tflite_gpu_runner_->Invoke(); + { + MEDIAPIPE_PROFILING(GPU_TASK_INVOKE, cc); + return tflite_gpu_runner_->Invoke(); + } })); return output_tensors; @@ -354,7 +359,7 @@ absl::Status InferenceCalculatorGlAdvancedImpl::Process(CalculatorContext* cc) { auto output_tensors = absl::make_unique>(); ASSIGN_OR_RETURN(*output_tensors, - gpu_inference_runner_->Process(input_tensors)); + gpu_inference_runner_->Process(cc, input_tensors)); kOutTensors(cc).Send(std::move(output_tensors)); return absl::OkStatus(); diff --git a/mediapipe/calculators/tensor/inference_calculator_xnnpack.cc b/mediapipe/calculators/tensor/inference_calculator_xnnpack.cc index 384c753f..a9417d50 100644 --- a/mediapipe/calculators/tensor/inference_calculator_xnnpack.cc +++ b/mediapipe/calculators/tensor/inference_calculator_xnnpack.cc @@ -70,7 +70,7 @@ absl::Status InferenceCalculatorXnnpackImpl::Process(CalculatorContext* cc) { RET_CHECK(!input_tensors.empty()); ASSIGN_OR_RETURN(std::vector output_tensors, - inference_runner_->Run(input_tensors)); + inference_runner_->Run(cc, input_tensors)); kOutTensors(cc).Send(std::move(output_tensors)); return absl::OkStatus(); } diff --git a/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.cc b/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.cc index 1d216daf..9ef4e822 100644 --- a/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.cc +++ b/mediapipe/calculators/tensor/inference_interpreter_delegate_runner.cc @@ -20,12 +20,15 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" #include "mediapipe/framework/formats/tensor.h" +#include "mediapipe/framework/mediapipe_profiling.h" #include "mediapipe/framework/port/ret_check.h" #include "tensorflow/lite/c/c_api_types.h" #include "tensorflow/lite/interpreter.h" #include "tensorflow/lite/interpreter_builder.h" #include "tensorflow/lite/string_util.h" +#define PERFETTO_TRACK_EVENT_NAMESPACE mediapipe + namespace mediapipe { namespace { @@ -79,7 +82,7 @@ class InferenceInterpreterDelegateRunner : public InferenceRunner { delegate_(std::move(delegate)) {} absl::StatusOr> Run( - const std::vector& input_tensors) override; + CalculatorContext* cc, const std::vector& input_tensors) override; private: api2::Packet model_; @@ -88,7 +91,7 @@ class InferenceInterpreterDelegateRunner : public InferenceRunner { }; absl::StatusOr> InferenceInterpreterDelegateRunner::Run( - const std::vector& input_tensors) { + CalculatorContext* cc, const std::vector& input_tensors) { // Read CPU input into tensors. RET_CHECK_EQ(interpreter_->inputs().size(), input_tensors.size()); for (int i = 0; i < input_tensors.size(); ++i) { @@ -131,8 +134,10 @@ absl::StatusOr> InferenceInterpreterDelegateRunner::Run( } // Run inference. - RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk); - + { + MEDIAPIPE_PROFILING(CPU_TASK_INVOKE, cc); + RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk); + } // Output result tensors (CPU). const auto& tensor_indexes = interpreter_->outputs(); std::vector output_tensors; diff --git a/mediapipe/calculators/tensor/inference_runner.h b/mediapipe/calculators/tensor/inference_runner.h index ec9d17b8..2283a6c1 100644 --- a/mediapipe/calculators/tensor/inference_runner.h +++ b/mediapipe/calculators/tensor/inference_runner.h @@ -2,6 +2,7 @@ #define MEDIAPIPE_CALCULATORS_TENSOR_INFERENCE_RUNNER_H_ #include "absl/status/statusor.h" +#include "mediapipe/framework/calculator_context.h" #include "mediapipe/framework/formats/tensor.h" namespace mediapipe { @@ -11,7 +12,7 @@ class InferenceRunner { public: virtual ~InferenceRunner() = default; virtual absl::StatusOr> Run( - const std::vector& inputs) = 0; + CalculatorContext* cc, const std::vector& inputs) = 0; }; } // namespace mediapipe diff --git a/mediapipe/calculators/tensor/tensors_to_audio_calculator.cc b/mediapipe/calculators/tensor/tensors_to_audio_calculator.cc new file mode 100644 index 00000000..8da29bb6 --- /dev/null +++ b/mediapipe/calculators/tensor/tensors_to_audio_calculator.cc @@ -0,0 +1,197 @@ +// Copyright 2022 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/status/status.h" +#include "audio/dsp/window_functions.h" +#include "mediapipe/calculators/tensor/tensors_to_audio_calculator.pb.h" +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/matrix.h" +#include "mediapipe/framework/formats/tensor.h" +#include "mediapipe/framework/port/ret_check.h" +#include "pffft.h" + +namespace mediapipe { +namespace api2 { +namespace { + +std::vector HannWindow(int window_size, bool sqrt_hann) { + std::vector hann_window(window_size); + audio_dsp::HannWindow().GetPeriodicSamples(window_size, &hann_window); + if (sqrt_hann) { + absl::c_transform(hann_window, hann_window.begin(), + [](double x) { return std::sqrt(x); }); + } + return hann_window; +} + +// Note that the InvHannWindow function may only work for 50% overlapping case. +std::vector InvHannWindow(int window_size, bool sqrt_hann) { + std::vector window = HannWindow(window_size, sqrt_hann); + std::vector inv_window(window.size()); + if (sqrt_hann) { + absl::c_copy(window, inv_window.begin()); + } else { + const int kHalfWindowSize = window.size() / 2; + absl::c_transform(window, inv_window.begin(), + [](double x) { return x * x; }); + for (int i = 0; i < kHalfWindowSize; ++i) { + double sum = inv_window[i] + inv_window[kHalfWindowSize + i]; + inv_window[i] = window[i] / sum; + inv_window[kHalfWindowSize + i] = window[kHalfWindowSize + i] / sum; + } + } + return inv_window; +} + +// PFFFT only supports transforms for inputs of length N of the form +// N = (2^a)*(3^b)*(5^c) where b >=0 and c >= 0 and a >= 5 for the real FFT. +bool IsValidFftSize(int size) { + if (size <= 0) { + return false; + } + constexpr int kFactors[] = {2, 3, 5}; + int factorization[] = {0, 0, 0}; + int n = static_cast(size); + for (int i = 0; i < 3; ++i) { + while (n % kFactors[i] == 0) { + n = n / kFactors[i]; + ++factorization[i]; + } + } + return factorization[0] >= 5 && n == 1; +} + +} // namespace + +// Converts 2D MediaPipe float Tensors to audio buffers. +// The calculator will perform ifft on the complex DFT and apply the window +// function (Inverse Hann) afterwards. The input 2D MediaPipe Tensor must +// have the DFT real parts in its first row and the DFT imagery parts in its +// second row. A valid "fft_size" must be set in the CalculatorOptions. +// +// Inputs: +// TENSORS - std::vector +// Vector containing a single Tensor that represents the audio's complex DFT +// results. +// DC_AND_NYQUIST - std::pair +// A pair of dc component and nyquist component. +// +// Outputs: +// AUDIO - mediapipe::Matrix +// The audio data represented as mediapipe::Matrix. +// +// Example: +// node { +// calculator: "TensorsToAudioCalculator" +// input_stream: "TENSORS:tensors" +// input_stream: "DC_AND_NYQUIST:dc_and_nyquist" +// output_stream: "AUDIO:audio" +// options { +// [mediapipe.AudioToTensorCalculatorOptions.ext] { +// fft_size: 256 +// } +// } +// } +class TensorsToAudioCalculator : public Node { + public: + static constexpr Input> kTensorsIn{"TENSORS"}; + static constexpr Input> kDcAndNyquistIn{ + "DC_AND_NYQUIST"}; + static constexpr Output kAudioOut{"AUDIO"}; + MEDIAPIPE_NODE_CONTRACT(kTensorsIn, kDcAndNyquistIn, kAudioOut); + + absl::Status Open(CalculatorContext* cc) override; + absl::Status Process(CalculatorContext* cc) override; + absl::Status Close(CalculatorContext* cc) override; + + private: + // The internal state of the FFT library. + PFFFT_Setup* fft_state_ = nullptr; + int fft_size_ = 0; + float inverse_fft_size_ = 0; + std::vector> input_dft_; + std::vector inv_fft_window_; + std::vector> fft_input_buffer_; + // pffft requires memory to work with to avoid using the stack. + std::vector> fft_workplace_; + std::vector> fft_output_; +}; + +absl::Status TensorsToAudioCalculator::Open(CalculatorContext* cc) { + const auto& options = + cc->Options(); + RET_CHECK(options.has_fft_size()) << "FFT size must be specified."; + RET_CHECK(IsValidFftSize(options.fft_size())) + << "FFT size must be of the form fft_size = (2^a)*(3^b)*(5^c) where b " + ">=0 and c >= 0 and a >= 5, the requested fft size is " + << options.fft_size(); + fft_size_ = options.fft_size(); + inverse_fft_size_ = 1.0f / fft_size_; + fft_state_ = pffft_new_setup(fft_size_, PFFFT_REAL); + input_dft_.resize(fft_size_); + inv_fft_window_ = InvHannWindow(fft_size_, /* sqrt_hann = */ false); + fft_input_buffer_.resize(fft_size_); + fft_workplace_.resize(fft_size_); + fft_output_.resize(fft_size_); + return absl::OkStatus(); +} + +absl::Status TensorsToAudioCalculator::Process(CalculatorContext* cc) { + if (kTensorsIn(cc).IsEmpty() || kDcAndNyquistIn(cc).IsEmpty()) { + return absl::OkStatus(); + } + const auto& input_tensors = *kTensorsIn(cc); + RET_CHECK_EQ(input_tensors.size(), 1); + RET_CHECK(input_tensors[0].element_type() == Tensor::ElementType::kFloat32); + auto view = input_tensors[0].GetCpuReadView(); + // DC's real part. + input_dft_[0] = kDcAndNyquistIn(cc)->first; + // Nyquist's real part is the penultimate element of the tensor buffer. + // pffft ignores the Nyquist's imagery part. No need to fetch the last value + // from the tensor buffer. + input_dft_[1] = *(view.buffer() + (fft_size_ - 2)); + std::memcpy(input_dft_.data() + 2, view.buffer(), + (fft_size_ - 2) * sizeof(float)); + pffft_transform_ordered(fft_state_, input_dft_.data(), fft_output_.data(), + fft_workplace_.data(), PFFFT_BACKWARD); + // Applies the inverse window function. + std::transform( + fft_output_.begin(), fft_output_.end(), inv_fft_window_.begin(), + fft_output_.begin(), + [this](float a, float b) { return a * b * inverse_fft_size_; }); + Matrix matrix = Eigen::Map(fft_output_.data(), 1, fft_output_.size()); + kAudioOut(cc).Send(std::move(matrix)); + return absl::OkStatus(); +} + +absl::Status TensorsToAudioCalculator::Close(CalculatorContext* cc) { + if (fft_state_) { + pffft_destroy_setup(fft_state_); + } + return absl::OkStatus(); +} + +MEDIAPIPE_REGISTER_NODE(TensorsToAudioCalculator); + +} // namespace api2 +} // namespace mediapipe diff --git a/mediapipe/calculators/tensor/tensors_to_audio_calculator.proto b/mediapipe/calculators/tensor/tensors_to_audio_calculator.proto new file mode 100644 index 00000000..90762712 --- /dev/null +++ b/mediapipe/calculators/tensor/tensors_to_audio_calculator.proto @@ -0,0 +1,29 @@ +// Copyright 2022 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +syntax = "proto2"; + +package mediapipe; + +import "mediapipe/framework/calculator.proto"; + +message TensorsToAudioCalculatorOptions { + extend mediapipe.CalculatorOptions { + optional TensorsToAudioCalculatorOptions ext = 484297136; + } + + // Size of the fft in number of bins. If set, the calculator will do ifft + // on the input tensor. + optional int64 fft_size = 1; +} diff --git a/mediapipe/calculators/tensor/tensors_to_audio_calculator_test.cc b/mediapipe/calculators/tensor/tensors_to_audio_calculator_test.cc new file mode 100644 index 00000000..b332381c --- /dev/null +++ b/mediapipe/calculators/tensor/tensors_to_audio_calculator_test.cc @@ -0,0 +1,149 @@ +// Copyright 2022 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/strings/substitute.h" +#include "mediapipe/calculators/tensor/audio_to_tensor_calculator.pb.h" +#include "mediapipe/calculators/tensor/tensors_to_audio_calculator.pb.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/matrix.h" +#include "mediapipe/framework/port/gtest.h" +#include "mediapipe/framework/port/parse_text_proto.h" +#include "mediapipe/framework/port/status_matchers.h" + +namespace mediapipe { +namespace { + +class TensorsToAudioCalculatorFftTest : public ::testing::Test { + protected: + // Creates an audio matrix containing a single sample of 1.0 at a specified + // offset. + Matrix CreateImpulseSignalData(int64 num_samples, int impulse_offset_idx) { + Matrix impulse = Matrix::Zero(1, num_samples); + impulse(0, impulse_offset_idx) = 1.0; + return impulse; + } + + void ConfigGraph(int num_samples, double sample_rate, int fft_size) { + graph_config_ = ParseTextProtoOrDie( + absl::Substitute(R"( + input_stream: "audio_in" + input_stream: "sample_rate" + output_stream: "audio_out" + node { + calculator: "AudioToTensorCalculator" + input_stream: "AUDIO:audio_in" + input_stream: "SAMPLE_RATE:sample_rate" + output_stream: "TENSORS:tensors" + output_stream: "DC_AND_NYQUIST:dc_and_nyquist" + options { + [mediapipe.AudioToTensorCalculatorOptions.ext] { + num_channels: 1 + num_samples: $0 + num_overlapping_samples: 0 + target_sample_rate: $1 + fft_size: $2 + } + } + } + node { + calculator: "TensorsToAudioCalculator" + input_stream: "TENSORS:tensors" + input_stream: "DC_AND_NYQUIST:dc_and_nyquist" + output_stream: "AUDIO:audio_out" + options { + [mediapipe.TensorsToAudioCalculatorOptions.ext] { + fft_size: $2 + } + } + } + )", + /*$0=*/num_samples, + /*$1=*/sample_rate, + /*$2=*/fft_size)); + tool::AddVectorSink("audio_out", &graph_config_, &audio_out_packets_); + } + + void RunGraph(const Matrix& input_data, double sample_rate) { + MP_ASSERT_OK(graph_.Initialize(graph_config_)); + MP_ASSERT_OK(graph_.StartRun({})); + MP_ASSERT_OK(graph_.AddPacketToInputStream( + "sample_rate", MakePacket(sample_rate).At(Timestamp(0)))); + MP_ASSERT_OK(graph_.AddPacketToInputStream( + "audio_in", MakePacket(input_data).At(Timestamp(0)))); + MP_ASSERT_OK(graph_.CloseAllInputStreams()); + MP_ASSERT_OK(graph_.WaitUntilDone()); + } + + std::vector audio_out_packets_; + CalculatorGraphConfig graph_config_; + CalculatorGraph graph_; +}; + +TEST_F(TensorsToAudioCalculatorFftTest, TestInvalidFftSize) { + ConfigGraph(320, 16000, 103); + MP_ASSERT_OK(graph_.Initialize(graph_config_)); + MP_ASSERT_OK(graph_.StartRun({})); + auto status = graph_.WaitUntilIdle(); + EXPECT_EQ(status.code(), absl::StatusCode::kInternal); + EXPECT_THAT(status.message(), + ::testing::HasSubstr("FFT size must be of the form")); +} + +TEST_F(TensorsToAudioCalculatorFftTest, TestImpulseSignalAtTheCenter) { + constexpr int sample_size = 320; + constexpr double sample_rate = 16000; + ConfigGraph(sample_size, sample_rate, 320); + + Matrix impulse_data = CreateImpulseSignalData(sample_size, sample_size / 2); + RunGraph(impulse_data, sample_rate); + ASSERT_EQ(1, audio_out_packets_.size()); + MP_ASSERT_OK(audio_out_packets_[0].ValidateAsType()); + // The impulse signal at the center is not affected by the window function. + EXPECT_EQ(audio_out_packets_[0].Get(), impulse_data); +} + +TEST_F(TensorsToAudioCalculatorFftTest, TestWindowedImpulseSignal) { + constexpr int sample_size = 320; + constexpr double sample_rate = 16000; + ConfigGraph(sample_size, sample_rate, 320); + Matrix impulse_data = CreateImpulseSignalData(sample_size, sample_size / 4); + RunGraph(impulse_data, sample_rate); + ASSERT_EQ(1, audio_out_packets_.size()); + MP_ASSERT_OK(audio_out_packets_[0].ValidateAsType()); + // As the impulse signal sits at the 1/4 of the hann window, the inverse + // window function reduces it by half. + EXPECT_EQ(audio_out_packets_[0].Get(), impulse_data / 2); +} + +TEST_F(TensorsToAudioCalculatorFftTest, TestImpulseSignalAtBeginning) { + constexpr int sample_size = 320; + constexpr double sample_rate = 16000; + ConfigGraph(sample_size, sample_rate, 320); + Matrix impulse_data = CreateImpulseSignalData(sample_size, 0); + RunGraph(impulse_data, sample_rate); + ASSERT_EQ(1, audio_out_packets_.size()); + MP_ASSERT_OK(audio_out_packets_[0].ValidateAsType()); + // As the impulse signal sits at the beginning of the hann window, the inverse + // window function completely removes it. + EXPECT_EQ(audio_out_packets_[0].Get(), Matrix::Zero(1, sample_size)); +} + +} // namespace +} // namespace mediapipe diff --git a/mediapipe/framework/api2/builder.h b/mediapipe/framework/api2/builder.h index 82905d2f..7dce211c 100644 --- a/mediapipe/framework/api2/builder.h +++ b/mediapipe/framework/api2/builder.h @@ -289,8 +289,15 @@ class NodeBase { template T& GetOptions() { + return GetOptions(T::ext); + } + + // Use this API when the proto extension does not follow the "ext" naming + // convention. + template + auto& GetOptions(const E& extension) { options_used_ = true; - return *options_.MutableExtension(T::ext); + return *options_.MutableExtension(extension); } protected: @@ -386,8 +393,15 @@ class PacketGenerator { template T& GetOptions() { + return GetOptions(T::ext); + } + + // Use this API when the proto extension does not follow the "ext" naming + // convention. + template + auto& GetOptions(const E& extension) { options_used_ = true; - return *options_.MutableExtension(T::ext); + return *options_.MutableExtension(extension); } template diff --git a/mediapipe/framework/calculator_base.h b/mediapipe/framework/calculator_base.h index f9f0d7a8..19f37f9d 100644 --- a/mediapipe/framework/calculator_base.h +++ b/mediapipe/framework/calculator_base.h @@ -185,7 +185,7 @@ class CalculatorBaseFactory { // Functions for checking that the calculator has the required GetContract. template constexpr bool CalculatorHasGetContract(decltype(&T::GetContract) /*unused*/) { - typedef absl::Status (*GetContractType)(CalculatorContract * cc); + typedef absl::Status (*GetContractType)(CalculatorContract* cc); return std::is_same::value; } template diff --git a/mediapipe/framework/calculator_profile.proto b/mediapipe/framework/calculator_profile.proto index 06ec678a..082d0a86 100644 --- a/mediapipe/framework/calculator_profile.proto +++ b/mediapipe/framework/calculator_profile.proto @@ -133,7 +133,13 @@ message GraphTrace { TPU_TASK = 13; GPU_CALIBRATION = 14; PACKET_QUEUED = 15; + GPU_TASK_INVOKE = 16; + TPU_TASK_INVOKE = 17; + CPU_TASK_INVOKE = 18; } + // //depot/mediapipe/framework/mediapipe_profiling.h:profiler_census_tags, + // //depot/mediapipe/framework/profiler/trace_buffer.h:event_type_list, + // ) // The timing for one packet set being processed at one caclulator node. message CalculatorTrace { diff --git a/mediapipe/framework/formats/BUILD b/mediapipe/framework/formats/BUILD index b967b27f..c3241d91 100644 --- a/mediapipe/framework/formats/BUILD +++ b/mediapipe/framework/formats/BUILD @@ -293,7 +293,6 @@ mediapipe_proto_library( name = "rect_proto", srcs = ["rect.proto"], visibility = ["//visibility:public"], - deps = ["//mediapipe/framework/formats:location_data_proto"], ) mediapipe_register_type( diff --git a/mediapipe/framework/profiler/trace_buffer.h b/mediapipe/framework/profiler/trace_buffer.h index 069f0961..9ad4314f 100644 --- a/mediapipe/framework/profiler/trace_buffer.h +++ b/mediapipe/framework/profiler/trace_buffer.h @@ -109,6 +109,13 @@ struct TraceEvent { static constexpr EventType TPU_TASK = GraphTrace::TPU_TASK; static constexpr EventType GPU_CALIBRATION = GraphTrace::GPU_CALIBRATION; static constexpr EventType PACKET_QUEUED = GraphTrace::PACKET_QUEUED; + static constexpr EventType GPU_TASK_INVOKE = GraphTrace::GPU_TASK_INVOKE; + static constexpr EventType TPU_TASK_INVOKE = GraphTrace::TPU_TASK_INVOKE; + static constexpr EventType CPU_TASK_INVOKE = GraphTrace::CPU_TASK_INVOKE; + + // //depot/mediapipe/framework/mediapipe_profiling.h:profiler_census_tags, + // //depot/mediapipe/framework/calculator_profile.proto:event_type, + // ) }; // Packet trace log buffer. diff --git a/mediapipe/framework/tool/switch_container.cc b/mediapipe/framework/tool/switch_container.cc index 9439acf9..daa12992 100644 --- a/mediapipe/framework/tool/switch_container.cc +++ b/mediapipe/framework/tool/switch_container.cc @@ -105,10 +105,10 @@ CalculatorGraphConfig::Node* BuildMuxNode( // Returns a PacketSequencerCalculator node. CalculatorGraphConfig::Node* BuildTimestampNode(CalculatorGraphConfig* config, - bool synchronize_io) { + bool async_selection) { CalculatorGraphConfig::Node* result = config->add_node(); *result->mutable_calculator() = "PacketSequencerCalculator"; - if (synchronize_io) { + if (!async_selection) { *result->mutable_input_stream_handler()->mutable_input_stream_handler() = "DefaultInputStreamHandler"; } @@ -239,6 +239,15 @@ bool HasTag(const proto_ns::RepeatedPtrField& streams, return tags.count({tag, 0}) > 0; } +// Returns true if a set of "TAG::index" includes a TagIndex. +bool ContainsTag(const proto_ns::RepeatedPtrField& tags, + TagIndex item) { + for (const std::string& t : tags) { + if (ParseTagIndex(t) == item) return true; + } + return false; +} + absl::StatusOr SwitchContainer::GetConfig( const Subgraph::SubgraphOptions& options) { CalculatorGraphConfig config; @@ -263,17 +272,17 @@ absl::StatusOr SwitchContainer::GetConfig( std::string enable_stream = "ENABLE:gate_enable"; // Add a PacketSequencerCalculator node for "SELECT" or "ENABLE" streams. - bool synchronize_io = - Subgraph::GetOptions(options) - .synchronize_io(); + const auto& switch_options = + Subgraph::GetOptions(options); + bool async_selection = switch_options.async_selection(); if (HasTag(container_node.input_stream(), "SELECT")) { - select_node = BuildTimestampNode(&config, synchronize_io); + select_node = BuildTimestampNode(&config, async_selection); select_node->add_input_stream("INPUT:gate_select"); select_node->add_output_stream("OUTPUT:gate_select_timed"); select_stream = "SELECT:gate_select_timed"; } if (HasTag(container_node.input_stream(), "ENABLE")) { - enable_node = BuildTimestampNode(&config, synchronize_io); + enable_node = BuildTimestampNode(&config, async_selection); enable_node->add_input_stream("INPUT:gate_enable"); enable_node->add_output_stream("OUTPUT:gate_enable_timed"); enable_stream = "ENABLE:gate_enable_timed"; @@ -296,7 +305,7 @@ absl::StatusOr SwitchContainer::GetConfig( mux->add_input_side_packet("SELECT:gate_select"); mux->add_input_side_packet("ENABLE:gate_enable"); - // Add input streams for graph and demux and the timestamper. + // Add input streams for graph and demux. config.add_input_stream("SELECT:gate_select"); config.add_input_stream("ENABLE:gate_enable"); config.add_input_side_packet("SELECT:gate_select"); @@ -306,6 +315,12 @@ absl::StatusOr SwitchContainer::GetConfig( std::string stream = CatStream(p.first, p.second); config.add_input_stream(stream); demux->add_input_stream(stream); + } + + // Add input streams for the timestamper. + auto& tick_streams = switch_options.tick_input_stream(); + for (const auto& p : input_tags) { + if (!tick_streams.empty() && !ContainsTag(tick_streams, p.first)) continue; TagIndex tick_tag{"TICK", tick_index++}; if (select_node) { select_node->add_input_stream(CatStream(tick_tag, p.second)); diff --git a/mediapipe/framework/tool/switch_container.proto b/mediapipe/framework/tool/switch_container.proto index a9c2d909..16050ef6 100644 --- a/mediapipe/framework/tool/switch_container.proto +++ b/mediapipe/framework/tool/switch_container.proto @@ -25,6 +25,14 @@ message SwitchContainerOptions { // Activates channel 1 for enable = true, channel 0 otherwise. optional bool enable = 4; - // Use DefaultInputStreamHandler for muxing & demuxing. + // Use DefaultInputStreamHandler for demuxing. optional bool synchronize_io = 5; + + // Use ImmediateInputStreamHandler for channel selection. + optional bool async_selection = 6; + + // Specifies an input stream, "TAG:index", that defines the processed + // timestamps. SwitchContainer awaits output at the last processed + // timestamp before advancing from one selected channel to the next. + repeated string tick_input_stream = 7; } diff --git a/mediapipe/framework/tool/switch_container_test.cc b/mediapipe/framework/tool/switch_container_test.cc index de4aa0b1..b20979b1 100644 --- a/mediapipe/framework/tool/switch_container_test.cc +++ b/mediapipe/framework/tool/switch_container_test.cc @@ -252,6 +252,9 @@ TEST(SwitchContainerTest, ApplyToSubnodes) { input_stream: "INPUT:enable" input_stream: "TICK:foo" output_stream: "OUTPUT:switchcontainer__gate_enable_timed" + input_stream_handler { + input_stream_handler: "DefaultInputStreamHandler" + } } node { name: "switchcontainer__SwitchDemuxCalculator" @@ -306,7 +309,8 @@ TEST(SwitchContainerTest, ApplyToSubnodes) { // Shows the SwitchContainer container runs with a pair of simple subnodes. TEST(SwitchContainerTest, RunsWithSubnodes) { EXPECT_TRUE(SubgraphRegistry::IsRegistered("SwitchContainer")); - CalculatorGraphConfig supergraph = SubnodeContainerExample(); + CalculatorGraphConfig supergraph = + SubnodeContainerExample("async_selection: true"); MP_EXPECT_OK(tool::ExpandSubgraphs(&supergraph)); RunTestContainer(supergraph); } diff --git a/mediapipe/framework/tool/switch_demux_calculator.cc b/mediapipe/framework/tool/switch_demux_calculator.cc index b9ba2a0f..c066d470 100644 --- a/mediapipe/framework/tool/switch_demux_calculator.cc +++ b/mediapipe/framework/tool/switch_demux_calculator.cc @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -54,21 +55,47 @@ namespace mediapipe { // contained subgraph or calculator nodes. // class SwitchDemuxCalculator : public CalculatorBase { - static constexpr char kSelectTag[] = "SELECT"; - static constexpr char kEnableTag[] = "ENABLE"; - public: static absl::Status GetContract(CalculatorContract* cc); absl::Status Open(CalculatorContext* cc) override; absl::Status Process(CalculatorContext* cc) override; + private: + absl::Status RecordPackets(CalculatorContext* cc); + int ChannelIndex(Timestamp timestamp); + absl::Status SendActivePackets(CalculatorContext* cc); + private: int channel_index_; std::set channel_tags_; + using PacketQueue = std::map>; + PacketQueue input_queue_; + std::map channel_history_; }; REGISTER_CALCULATOR(SwitchDemuxCalculator); +namespace { +static constexpr char kSelectTag[] = "SELECT"; +static constexpr char kEnableTag[] = "ENABLE"; + +// Returns the last received timestamp for an input stream. +inline Timestamp SettledTimestamp(const InputStreamShard& input) { + return input.Value().Timestamp(); +} + +// Returns the last received timestamp for channel selection. +inline Timestamp ChannelSettledTimestamp(CalculatorContext* cc) { + Timestamp result = Timestamp::Done(); + if (cc->Inputs().HasTag(kEnableTag)) { + result = SettledTimestamp(cc->Inputs().Tag(kEnableTag)); + } else if (cc->Inputs().HasTag(kSelectTag)) { + result = SettledTimestamp(cc->Inputs().Tag(kSelectTag)); + } + return result; +} +} // namespace + absl::Status SwitchDemuxCalculator::GetContract(CalculatorContract* cc) { // Allow any one of kSelectTag, kEnableTag. cc->Inputs().Tag(kSelectTag).Set().Optional(); @@ -125,6 +152,7 @@ absl::Status SwitchDemuxCalculator::GetContract(CalculatorContract* cc) { absl::Status SwitchDemuxCalculator::Open(CalculatorContext* cc) { channel_index_ = tool::GetChannelIndex(*cc, channel_index_); channel_tags_ = ChannelTags(cc->Outputs().TagMap()); + channel_history_[Timestamp::Unstarted()] = channel_index_; // Relay side packets to all channels. // Note: This is necessary because Calculator::Open only proceeds when every @@ -164,21 +192,77 @@ absl::Status SwitchDemuxCalculator::Open(CalculatorContext* cc) { } absl::Status SwitchDemuxCalculator::Process(CalculatorContext* cc) { - // Update the input channel index if specified. - channel_index_ = tool::GetChannelIndex(*cc, channel_index_); + MP_RETURN_IF_ERROR(RecordPackets(cc)); + MP_RETURN_IF_ERROR(SendActivePackets(cc)); + return absl::OkStatus(); +} - // Relay packets and timestamps only to channel_index_. +// Enqueue all arriving packets and bounds. +absl::Status SwitchDemuxCalculator::RecordPackets(CalculatorContext* cc) { + // Enqueue any new arriving packets. for (const std::string& tag : channel_tags_) { for (int index = 0; index < cc->Inputs().NumEntries(tag); ++index) { - auto& input = cc->Inputs().Get(tag, index); - std::string output_tag = tool::ChannelTag(tag, channel_index_); - auto output_id = cc->Outputs().GetId(output_tag, index); - if (output_id.IsValid()) { - auto& output = cc->Outputs().Get(output_tag, index); - tool::Relay(input, &output); + auto input_id = cc->Inputs().GetId(tag, index); + Packet packet = cc->Inputs().Get(input_id).Value(); + if (packet.Timestamp() == cc->InputTimestamp()) { + input_queue_[input_id].push(packet); } } } + + // Enque any new input channel and its activation timestamp. + Timestamp channel_settled = ChannelSettledTimestamp(cc); + int new_channel_index = tool::GetChannelIndex(*cc, channel_index_); + if (channel_settled == cc->InputTimestamp() && + new_channel_index != channel_index_) { + channel_index_ = new_channel_index; + channel_history_[channel_settled] = channel_index_; + } + return absl::OkStatus(); +} + +// Returns the channel index for a Timestamp. +int SwitchDemuxCalculator::ChannelIndex(Timestamp timestamp) { + auto it = std::prev(channel_history_.upper_bound(timestamp)); + return it->second; +} + +// Dispatches all queued input packets with known channels. +absl::Status SwitchDemuxCalculator::SendActivePackets(CalculatorContext* cc) { + // Dispatch any queued input packets with a defined channel_index. + Timestamp channel_settled = ChannelSettledTimestamp(cc); + for (const std::string& tag : channel_tags_) { + for (int index = 0; index < cc->Inputs().NumEntries(tag); ++index) { + auto input_id = cc->Inputs().GetId(tag, index); + auto& queue = input_queue_[input_id]; + while (!queue.empty() && queue.front().Timestamp() <= channel_settled) { + int channel_index = ChannelIndex(queue.front().Timestamp()); + std::string output_tag = tool::ChannelTag(tag, channel_index); + auto output_id = cc->Outputs().GetId(output_tag, index); + if (output_id.IsValid()) { + cc->Outputs().Get(output_id).AddPacket(queue.front()); + } + queue.pop(); + } + } + } + + // Discard all select packets not needed for any remaining input packets. + Timestamp input_settled = Timestamp::Done(); + for (const std::string& tag : channel_tags_) { + for (int index = 0; index < cc->Inputs().NumEntries(tag); ++index) { + auto input_id = cc->Inputs().GetId(tag, index); + Timestamp stream_settled = SettledTimestamp(cc->Inputs().Get(input_id)); + if (!input_queue_[input_id].empty()) { + Timestamp stream_bound = input_queue_[input_id].front().Timestamp(); + stream_settled = + std::min(stream_settled, stream_bound.PreviousAllowedInStream()); + } + } + } + Timestamp input_bound = input_settled.NextAllowedInStream(); + auto history_bound = std::prev(channel_history_.upper_bound(input_bound)); + channel_history_.erase(channel_history_.begin(), history_bound); return absl::OkStatus(); } diff --git a/mediapipe/framework/tool/switch_mux_calculator.cc b/mediapipe/framework/tool/switch_mux_calculator.cc index 1a313662..230544b6 100644 --- a/mediapipe/framework/tool/switch_mux_calculator.cc +++ b/mediapipe/framework/tool/switch_mux_calculator.cc @@ -164,7 +164,7 @@ absl::Status SwitchMuxCalculator::Open(CalculatorContext* cc) { options_ = cc->Options(); channel_index_ = tool::GetChannelIndex(*cc, channel_index_); channel_tags_ = ChannelTags(cc->Inputs().TagMap()); - channel_history_[Timestamp::Unset()] = channel_index_; + channel_history_[Timestamp::Unstarted()] = channel_index_; // Relay side packets only from channel_index_. for (const std::string& tag : ChannelTags(cc->InputSidePackets().TagMap())) { diff --git a/mediapipe/gpu/gl_context_egl.cc b/mediapipe/gpu/gl_context_egl.cc index 13710a68..78b196b0 100644 --- a/mediapipe/gpu/gl_context_egl.cc +++ b/mediapipe/gpu/gl_context_egl.cc @@ -38,13 +38,20 @@ static pthread_key_t egl_release_thread_key; static pthread_once_t egl_release_key_once = PTHREAD_ONCE_INIT; static void EglThreadExitCallback(void* key_value) { +#if defined(__ANDROID__) + eglMakeCurrent(EGL_NO_DISPLAY, EGL_NO_SURFACE, EGL_NO_SURFACE, + EGL_NO_CONTEXT); +#else // Some implementations have chosen to allow EGL_NO_DISPLAY as a valid display // parameter for eglMakeCurrent. This behavior is not portable to all EGL // implementations, and should be considered as an undocumented vendor // extension. // https://www.khronos.org/registry/EGL/sdk/docs/man/html/eglMakeCurrent.xhtml + // + // NOTE: crashes on some Android devices (occurs with libGLES_meow.so). eglMakeCurrent(eglGetDisplay(EGL_DEFAULT_DISPLAY), EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT); +#endif eglReleaseThread(); } diff --git a/mediapipe/java/com/google/mediapipe/framework/AndroidPacketCreator.java b/mediapipe/java/com/google/mediapipe/framework/AndroidPacketCreator.java index 4af9dae7..05700ba1 100644 --- a/mediapipe/java/com/google/mediapipe/framework/AndroidPacketCreator.java +++ b/mediapipe/java/com/google/mediapipe/framework/AndroidPacketCreator.java @@ -17,8 +17,8 @@ package com.google.mediapipe.framework; import android.graphics.Bitmap; import com.google.mediapipe.framework.image.BitmapExtractor; import com.google.mediapipe.framework.image.ByteBufferExtractor; -import com.google.mediapipe.framework.image.Image; -import com.google.mediapipe.framework.image.ImageProperties; +import com.google.mediapipe.framework.image.MPImage; +import com.google.mediapipe.framework.image.MPImageProperties; import java.nio.ByteBuffer; // TODO: use Preconditions in this file. @@ -60,24 +60,24 @@ public class AndroidPacketCreator extends PacketCreator { } /** - * Creates an Image packet from an {@link Image}. + * Creates a MediaPipe Image packet from a {@link MPImage}. * *

The ImageContainerType must be IMAGE_CONTAINER_BYTEBUFFER or IMAGE_CONTAINER_BITMAP. */ - public Packet createImage(Image image) { + public Packet createImage(MPImage image) { // TODO: Choose the best storage from multiple containers. - ImageProperties properties = image.getContainedImageProperties().get(0); - if (properties.getStorageType() == Image.STORAGE_TYPE_BYTEBUFFER) { + MPImageProperties properties = image.getContainedImageProperties().get(0); + if (properties.getStorageType() == MPImage.STORAGE_TYPE_BYTEBUFFER) { ByteBuffer buffer = ByteBufferExtractor.extract(image); int numChannels = 0; switch (properties.getImageFormat()) { - case Image.IMAGE_FORMAT_RGBA: + case MPImage.IMAGE_FORMAT_RGBA: numChannels = 4; break; - case Image.IMAGE_FORMAT_RGB: + case MPImage.IMAGE_FORMAT_RGB: numChannels = 3; break; - case Image.IMAGE_FORMAT_ALPHA: + case MPImage.IMAGE_FORMAT_ALPHA: numChannels = 1; break; default: // fall out @@ -90,7 +90,7 @@ public class AndroidPacketCreator extends PacketCreator { int height = image.getHeight(); return createImage(buffer, width, height, numChannels); } - if (properties.getImageFormat() == Image.STORAGE_TYPE_BITMAP) { + if (properties.getStorageType() == MPImage.STORAGE_TYPE_BITMAP) { Bitmap bitmap = BitmapExtractor.extract(image); if (bitmap.getConfig() != Bitmap.Config.ARGB_8888) { throw new UnsupportedOperationException("bitmap must use ARGB_8888 config."); @@ -100,7 +100,7 @@ public class AndroidPacketCreator extends PacketCreator { // Unsupported type. throw new UnsupportedOperationException( - "Unsupported Image container type: " + properties.getImageFormat()); + "Unsupported Image container type: " + properties.getStorageType()); } /** diff --git a/mediapipe/java/com/google/mediapipe/framework/image/BitmapExtractor.java b/mediapipe/java/com/google/mediapipe/framework/image/BitmapExtractor.java index 4c6cebd4..d6f50bf3 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/BitmapExtractor.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/BitmapExtractor.java @@ -18,29 +18,29 @@ package com.google.mediapipe.framework.image; import android.graphics.Bitmap; /** - * Utility for extracting {@link android.graphics.Bitmap} from {@link Image}. + * Utility for extracting {@link android.graphics.Bitmap} from {@link MPImage}. * - *

Currently it only supports {@link Image} with {@link Image#STORAGE_TYPE_BITMAP}, otherwise + *

Currently it only supports {@link MPImage} with {@link MPImage#STORAGE_TYPE_BITMAP}, otherwise * {@link IllegalArgumentException} will be thrown. */ public final class BitmapExtractor { /** - * Extracts a {@link android.graphics.Bitmap} from an {@link Image}. + * Extracts a {@link android.graphics.Bitmap} from a {@link MPImage}. * * @param image the image to extract {@link android.graphics.Bitmap} from. - * @return the {@link android.graphics.Bitmap} stored in {@link Image} + * @return the {@link android.graphics.Bitmap} stored in {@link MPImage} * @throws IllegalArgumentException when the extraction requires unsupported format or data type * conversions. */ - public static Bitmap extract(Image image) { - ImageContainer imageContainer = image.getContainer(Image.STORAGE_TYPE_BITMAP); + public static Bitmap extract(MPImage image) { + MPImageContainer imageContainer = image.getContainer(MPImage.STORAGE_TYPE_BITMAP); if (imageContainer != null) { return ((BitmapImageContainer) imageContainer).getBitmap(); } else { // TODO: Support ByteBuffer -> Bitmap conversion. throw new IllegalArgumentException( - "Extracting Bitmap from an Image created by objects other than Bitmap is not" + "Extracting Bitmap from a MPImage created by objects other than Bitmap is not" + " supported"); } } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/BitmapImageBuilder.java b/mediapipe/java/com/google/mediapipe/framework/image/BitmapImageBuilder.java index ea2ca6b1..988cdf54 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/BitmapImageBuilder.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/BitmapImageBuilder.java @@ -22,7 +22,7 @@ import android.provider.MediaStore; import java.io.IOException; /** - * Builds {@link Image} from {@link android.graphics.Bitmap}. + * Builds {@link MPImage} from {@link android.graphics.Bitmap}. * *

You can pass in either mutable or immutable {@link android.graphics.Bitmap}. However once * {@link android.graphics.Bitmap} is passed in, to keep data integrity you shouldn't modify content @@ -49,7 +49,7 @@ public class BitmapImageBuilder { } /** - * Creates the builder to build {@link Image} from a file. + * Creates the builder to build {@link MPImage} from a file. * * @param context the application context. * @param uri the path to the resource file. @@ -58,15 +58,15 @@ public class BitmapImageBuilder { this(MediaStore.Images.Media.getBitmap(context.getContentResolver(), uri)); } - /** Sets value for {@link Image#getTimestamp()}. */ + /** Sets value for {@link MPImage#getTimestamp()}. */ BitmapImageBuilder setTimestamp(long timestamp) { this.timestamp = timestamp; return this; } - /** Builds an {@link Image} instance. */ - public Image build() { - return new Image( + /** Builds a {@link MPImage} instance. */ + public MPImage build() { + return new MPImage( new BitmapImageContainer(bitmap), timestamp, bitmap.getWidth(), bitmap.getHeight()); } } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/BitmapImageContainer.java b/mediapipe/java/com/google/mediapipe/framework/image/BitmapImageContainer.java index 0457e1e9..6fbcac21 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/BitmapImageContainer.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/BitmapImageContainer.java @@ -16,19 +16,19 @@ limitations under the License. package com.google.mediapipe.framework.image; import android.graphics.Bitmap; -import com.google.mediapipe.framework.image.Image.ImageFormat; +import com.google.mediapipe.framework.image.MPImage.MPImageFormat; -class BitmapImageContainer implements ImageContainer { +class BitmapImageContainer implements MPImageContainer { private final Bitmap bitmap; - private final ImageProperties properties; + private final MPImageProperties properties; public BitmapImageContainer(Bitmap bitmap) { this.bitmap = bitmap; this.properties = - ImageProperties.builder() + MPImageProperties.builder() .setImageFormat(convertFormatCode(bitmap.getConfig())) - .setStorageType(Image.STORAGE_TYPE_BITMAP) + .setStorageType(MPImage.STORAGE_TYPE_BITMAP) .build(); } @@ -37,7 +37,7 @@ class BitmapImageContainer implements ImageContainer { } @Override - public ImageProperties getImageProperties() { + public MPImageProperties getImageProperties() { return properties; } @@ -46,15 +46,15 @@ class BitmapImageContainer implements ImageContainer { bitmap.recycle(); } - @ImageFormat + @MPImageFormat static int convertFormatCode(Bitmap.Config config) { switch (config) { case ALPHA_8: - return Image.IMAGE_FORMAT_ALPHA; + return MPImage.IMAGE_FORMAT_ALPHA; case ARGB_8888: - return Image.IMAGE_FORMAT_RGBA; + return MPImage.IMAGE_FORMAT_RGBA; default: - return Image.IMAGE_FORMAT_UNKNOWN; + return MPImage.IMAGE_FORMAT_UNKNOWN; } } } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferExtractor.java b/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferExtractor.java index a0e8c3df..748a1066 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferExtractor.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferExtractor.java @@ -21,45 +21,45 @@ import android.graphics.Bitmap.Config; import android.os.Build.VERSION; import android.os.Build.VERSION_CODES; import com.google.auto.value.AutoValue; -import com.google.mediapipe.framework.image.Image.ImageFormat; +import com.google.mediapipe.framework.image.MPImage.MPImageFormat; import java.nio.ByteBuffer; import java.nio.ByteOrder; import java.util.Locale; /** - * Utility for extracting {@link ByteBuffer} from {@link Image}. + * Utility for extracting {@link ByteBuffer} from {@link MPImage}. * - *

Currently it only supports {@link Image} with {@link Image#STORAGE_TYPE_BYTEBUFFER}, otherwise - * {@link IllegalArgumentException} will be thrown. + *

Currently it only supports {@link MPImage} with {@link MPImage#STORAGE_TYPE_BYTEBUFFER}, + * otherwise {@link IllegalArgumentException} will be thrown. */ public class ByteBufferExtractor { /** - * Extracts a {@link ByteBuffer} from an {@link Image}. + * Extracts a {@link ByteBuffer} from a {@link MPImage}. * *

The returned {@link ByteBuffer} is a read-only view, with the first available {@link - * ImageProperties} whose storage type is {@code Image.STORAGE_TYPE_BYTEBUFFER}. + * MPImageProperties} whose storage type is {@code MPImage.STORAGE_TYPE_BYTEBUFFER}. * - * @see Image#getContainedImageProperties() + * @see MPImage#getContainedImageProperties() * @return A read-only {@link ByteBuffer}. * @throws IllegalArgumentException when the image doesn't contain a {@link ByteBuffer} storage. */ @SuppressLint("SwitchIntDef") - public static ByteBuffer extract(Image image) { - ImageContainer container = image.getContainer(); + public static ByteBuffer extract(MPImage image) { + MPImageContainer container = image.getContainer(); switch (container.getImageProperties().getStorageType()) { - case Image.STORAGE_TYPE_BYTEBUFFER: + case MPImage.STORAGE_TYPE_BYTEBUFFER: ByteBufferImageContainer byteBufferImageContainer = (ByteBufferImageContainer) container; return byteBufferImageContainer.getByteBuffer().asReadOnlyBuffer(); default: throw new IllegalArgumentException( - "Extract ByteBuffer from an Image created by objects other than Bytebuffer is not" + "Extract ByteBuffer from a MPImage created by objects other than Bytebuffer is not" + " supported"); } } /** - * Extracts a readonly {@link ByteBuffer} in given {@code targetFormat} from an {@link Image}. + * Extracts a readonly {@link ByteBuffer} in given {@code targetFormat} from a {@link MPImage}. * *

Format conversion spec: * @@ -70,26 +70,26 @@ public class ByteBufferExtractor { * * @param image the image to extract buffer from. * @param targetFormat the image format of the result bytebuffer. - * @return the readonly {@link ByteBuffer} stored in {@link Image} + * @return the readonly {@link ByteBuffer} stored in {@link MPImage} * @throws IllegalArgumentException when the extraction requires unsupported format or data type * conversions. */ - static ByteBuffer extract(Image image, @ImageFormat int targetFormat) { - ImageContainer container; - ImageProperties byteBufferProperties = - ImageProperties.builder() - .setStorageType(Image.STORAGE_TYPE_BYTEBUFFER) + static ByteBuffer extract(MPImage image, @MPImageFormat int targetFormat) { + MPImageContainer container; + MPImageProperties byteBufferProperties = + MPImageProperties.builder() + .setStorageType(MPImage.STORAGE_TYPE_BYTEBUFFER) .setImageFormat(targetFormat) .build(); if ((container = image.getContainer(byteBufferProperties)) != null) { ByteBufferImageContainer byteBufferImageContainer = (ByteBufferImageContainer) container; return byteBufferImageContainer.getByteBuffer().asReadOnlyBuffer(); - } else if ((container = image.getContainer(Image.STORAGE_TYPE_BYTEBUFFER)) != null) { + } else if ((container = image.getContainer(MPImage.STORAGE_TYPE_BYTEBUFFER)) != null) { ByteBufferImageContainer byteBufferImageContainer = (ByteBufferImageContainer) container; - @ImageFormat int sourceFormat = byteBufferImageContainer.getImageFormat(); + @MPImageFormat int sourceFormat = byteBufferImageContainer.getImageFormat(); return convertByteBuffer(byteBufferImageContainer.getByteBuffer(), sourceFormat, targetFormat) .asReadOnlyBuffer(); - } else if ((container = image.getContainer(Image.STORAGE_TYPE_BITMAP)) != null) { + } else if ((container = image.getContainer(MPImage.STORAGE_TYPE_BITMAP)) != null) { BitmapImageContainer bitmapImageContainer = (BitmapImageContainer) container; ByteBuffer byteBuffer = extractByteBufferFromBitmap(bitmapImageContainer.getBitmap(), targetFormat) @@ -98,85 +98,89 @@ public class ByteBufferExtractor { return byteBuffer; } else { throw new IllegalArgumentException( - "Extracting ByteBuffer from an Image created by objects other than Bitmap or" + "Extracting ByteBuffer from a MPImage created by objects other than Bitmap or" + " Bytebuffer is not supported"); } } - /** A wrapper for a {@link ByteBuffer} and its {@link ImageFormat}. */ + /** A wrapper for a {@link ByteBuffer} and its {@link MPImageFormat}. */ @AutoValue abstract static class Result { - /** Gets the {@link ByteBuffer} in the result of {@link ByteBufferExtractor#extract(Image)}. */ + /** + * Gets the {@link ByteBuffer} in the result of {@link ByteBufferExtractor#extract(MPImage)}. + */ public abstract ByteBuffer buffer(); - /** Gets the {@link ImageFormat} in the result of {@link ByteBufferExtractor#extract(Image)}. */ - @ImageFormat + /** + * Gets the {@link MPImageFormat} in the result of {@link ByteBufferExtractor#extract(MPImage)}. + */ + @MPImageFormat public abstract int format(); - static Result create(ByteBuffer buffer, @ImageFormat int imageFormat) { + static Result create(ByteBuffer buffer, @MPImageFormat int imageFormat) { return new AutoValue_ByteBufferExtractor_Result(buffer, imageFormat); } } /** - * Extracts a {@link ByteBuffer} in any available {@code imageFormat} from an {@link Image}. + * Extracts a {@link ByteBuffer} in any available {@code imageFormat} from a {@link MPImage}. * *

It will make the best effort to return an already existed {@link ByteBuffer} to avoid copy. * - * @return the readonly {@link ByteBuffer} stored in {@link Image} + * @return the readonly {@link ByteBuffer} stored in {@link MPImage} * @throws IllegalArgumentException when {@code image} doesn't contain {@link ByteBuffer} with * given {@code imageFormat} */ - static Result extractInRecommendedFormat(Image image) { - ImageContainer container; - if ((container = image.getContainer(Image.STORAGE_TYPE_BITMAP)) != null) { + static Result extractInRecommendedFormat(MPImage image) { + MPImageContainer container; + if ((container = image.getContainer(MPImage.STORAGE_TYPE_BITMAP)) != null) { Bitmap bitmap = ((BitmapImageContainer) container).getBitmap(); - @ImageFormat int format = adviseImageFormat(bitmap); + @MPImageFormat int format = adviseImageFormat(bitmap); Result result = Result.create(extractByteBufferFromBitmap(bitmap, format).asReadOnlyBuffer(), format); boolean unused = image.addContainer(new ByteBufferImageContainer(result.buffer(), result.format())); return result; - } else if ((container = image.getContainer(Image.STORAGE_TYPE_BYTEBUFFER)) != null) { + } else if ((container = image.getContainer(MPImage.STORAGE_TYPE_BYTEBUFFER)) != null) { ByteBufferImageContainer byteBufferImageContainer = (ByteBufferImageContainer) container; return Result.create( byteBufferImageContainer.getByteBuffer().asReadOnlyBuffer(), byteBufferImageContainer.getImageFormat()); } else { throw new IllegalArgumentException( - "Extract ByteBuffer from an Image created by objects other than Bitmap or Bytebuffer" + "Extract ByteBuffer from a MPImage created by objects other than Bitmap or Bytebuffer" + " is not supported"); } } - @ImageFormat + @MPImageFormat private static int adviseImageFormat(Bitmap bitmap) { if (bitmap.getConfig() == Config.ARGB_8888) { - return Image.IMAGE_FORMAT_RGBA; + return MPImage.IMAGE_FORMAT_RGBA; } else { throw new IllegalArgumentException( String.format( - "Extracting ByteBuffer from an Image created by a Bitmap in config %s is not" + "Extracting ByteBuffer from a MPImage created by a Bitmap in config %s is not" + " supported", bitmap.getConfig())); } } private static ByteBuffer extractByteBufferFromBitmap( - Bitmap bitmap, @ImageFormat int imageFormat) { + Bitmap bitmap, @MPImageFormat int imageFormat) { if (VERSION.SDK_INT >= VERSION_CODES.JELLY_BEAN_MR1 && bitmap.isPremultiplied()) { throw new IllegalArgumentException( - "Extracting ByteBuffer from an Image created by a premultiplied Bitmap is not" + "Extracting ByteBuffer from a MPImage created by a premultiplied Bitmap is not" + " supported"); } if (bitmap.getConfig() == Config.ARGB_8888) { - if (imageFormat == Image.IMAGE_FORMAT_RGBA) { + if (imageFormat == MPImage.IMAGE_FORMAT_RGBA) { ByteBuffer buffer = ByteBuffer.allocateDirect(bitmap.getByteCount()); bitmap.copyPixelsToBuffer(buffer); buffer.rewind(); return buffer; - } else if (imageFormat == Image.IMAGE_FORMAT_RGB) { + } else if (imageFormat == MPImage.IMAGE_FORMAT_RGB) { // TODO: Try Use RGBA buffer to create RGB buffer which might be faster. int w = bitmap.getWidth(); int h = bitmap.getHeight(); @@ -196,14 +200,14 @@ public class ByteBufferExtractor { } throw new IllegalArgumentException( String.format( - "Extracting ByteBuffer from an Image created by Bitmap and convert from %s to format" + "Extracting ByteBuffer from a MPImage created by Bitmap and convert from %s to format" + " %d is not supported", bitmap.getConfig(), imageFormat)); } private static ByteBuffer convertByteBuffer( - ByteBuffer source, @ImageFormat int sourceFormat, @ImageFormat int targetFormat) { - if (sourceFormat == Image.IMAGE_FORMAT_RGB && targetFormat == Image.IMAGE_FORMAT_RGBA) { + ByteBuffer source, @MPImageFormat int sourceFormat, @MPImageFormat int targetFormat) { + if (sourceFormat == MPImage.IMAGE_FORMAT_RGB && targetFormat == MPImage.IMAGE_FORMAT_RGBA) { ByteBuffer target = ByteBuffer.allocateDirect(source.capacity() / 3 * 4); // Extend the buffer when the target is longer than the source. Use two cursors and sweep the // array reversely to convert in-place. @@ -221,7 +225,8 @@ public class ByteBufferExtractor { target.put(array, 0, target.capacity()); target.rewind(); return target; - } else if (sourceFormat == Image.IMAGE_FORMAT_RGBA && targetFormat == Image.IMAGE_FORMAT_RGB) { + } else if (sourceFormat == MPImage.IMAGE_FORMAT_RGBA + && targetFormat == MPImage.IMAGE_FORMAT_RGB) { ByteBuffer target = ByteBuffer.allocateDirect(source.capacity() / 4 * 3); // Shrink the buffer when the target is shorter than the source. Use two cursors and sweep the // array to convert in-place. diff --git a/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferImageBuilder.java b/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferImageBuilder.java index 07871da3..a650e4c3 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferImageBuilder.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferImageBuilder.java @@ -15,11 +15,11 @@ limitations under the License. package com.google.mediapipe.framework.image; -import com.google.mediapipe.framework.image.Image.ImageFormat; +import com.google.mediapipe.framework.image.MPImage.MPImageFormat; import java.nio.ByteBuffer; /** - * Builds a {@link Image} from a {@link ByteBuffer}. + * Builds a {@link MPImage} from a {@link ByteBuffer}. * *

You can pass in either mutable or immutable {@link ByteBuffer}. However once {@link * ByteBuffer} is passed in, to keep data integrity you shouldn't modify content in it. @@ -32,7 +32,7 @@ public class ByteBufferImageBuilder { private final ByteBuffer buffer; private final int width; private final int height; - @ImageFormat private final int imageFormat; + @MPImageFormat private final int imageFormat; // Optional fields. private long timestamp; @@ -49,7 +49,7 @@ public class ByteBufferImageBuilder { * @param imageFormat how the data encode the image. */ public ByteBufferImageBuilder( - ByteBuffer byteBuffer, int width, int height, @ImageFormat int imageFormat) { + ByteBuffer byteBuffer, int width, int height, @MPImageFormat int imageFormat) { this.buffer = byteBuffer; this.width = width; this.height = height; @@ -58,14 +58,14 @@ public class ByteBufferImageBuilder { this.timestamp = 0; } - /** Sets value for {@link Image#getTimestamp()}. */ + /** Sets value for {@link MPImage#getTimestamp()}. */ ByteBufferImageBuilder setTimestamp(long timestamp) { this.timestamp = timestamp; return this; } - /** Builds an {@link Image} instance. */ - public Image build() { - return new Image(new ByteBufferImageContainer(buffer, imageFormat), timestamp, width, height); + /** Builds a {@link MPImage} instance. */ + public MPImage build() { + return new MPImage(new ByteBufferImageContainer(buffer, imageFormat), timestamp, width, height); } } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferImageContainer.java b/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferImageContainer.java index 1c24c1df..82dbe32c 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferImageContainer.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/ByteBufferImageContainer.java @@ -15,21 +15,19 @@ limitations under the License. package com.google.mediapipe.framework.image; -import com.google.mediapipe.framework.image.Image.ImageFormat; +import com.google.mediapipe.framework.image.MPImage.MPImageFormat; import java.nio.ByteBuffer; -class ByteBufferImageContainer implements ImageContainer { +class ByteBufferImageContainer implements MPImageContainer { private final ByteBuffer buffer; - private final ImageProperties properties; + private final MPImageProperties properties; - public ByteBufferImageContainer( - ByteBuffer buffer, - @ImageFormat int imageFormat) { + public ByteBufferImageContainer(ByteBuffer buffer, @MPImageFormat int imageFormat) { this.buffer = buffer; this.properties = - ImageProperties.builder() - .setStorageType(Image.STORAGE_TYPE_BYTEBUFFER) + MPImageProperties.builder() + .setStorageType(MPImage.STORAGE_TYPE_BYTEBUFFER) .setImageFormat(imageFormat) .build(); } @@ -39,14 +37,12 @@ class ByteBufferImageContainer implements ImageContainer { } @Override - public ImageProperties getImageProperties() { + public MPImageProperties getImageProperties() { return properties; } - /** - * Returns the image format. - */ - @ImageFormat + /** Returns the image format. */ + @MPImageFormat public int getImageFormat() { return properties.getImageFormat(); } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/Image.java b/mediapipe/java/com/google/mediapipe/framework/image/MPImage.java similarity index 76% rename from mediapipe/java/com/google/mediapipe/framework/image/Image.java rename to mediapipe/java/com/google/mediapipe/framework/image/MPImage.java index 49e63bcc..e17cc4d3 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/Image.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/MPImage.java @@ -29,10 +29,10 @@ import java.util.Map.Entry; /** * The wrapper class for image objects. * - *

{@link Image} is designed to be an immutable image container, which could be shared + *

{@link MPImage} is designed to be an immutable image container, which could be shared * cross-platforms. * - *

To construct an {@link Image}, use the provided builders: + *

To construct a {@link MPImage}, use the provided builders: * *

    *
  • {@link ByteBufferImageBuilder} @@ -40,7 +40,7 @@ import java.util.Map.Entry; *
  • {@link MediaImageBuilder} *
* - *

{@link Image} uses reference counting to maintain internal storage. When it is created the + *

{@link MPImage} uses reference counting to maintain internal storage. When it is created the * reference count is 1. Developer can call {@link #close()} to reduce reference count to release * internal storage earlier, otherwise Java garbage collection will release the storage eventually. * @@ -53,7 +53,7 @@ import java.util.Map.Entry; *

  • {@link MediaImageExtractor} * */ -public class Image implements Closeable { +public class MPImage implements Closeable { /** Specifies the image format of an image. */ @IntDef({ @@ -69,7 +69,7 @@ public class Image implements Closeable { IMAGE_FORMAT_JPEG, }) @Retention(RetentionPolicy.SOURCE) - public @interface ImageFormat {} + public @interface MPImageFormat {} public static final int IMAGE_FORMAT_UNKNOWN = 0; public static final int IMAGE_FORMAT_RGBA = 1; @@ -98,14 +98,14 @@ public class Image implements Closeable { public static final int STORAGE_TYPE_IMAGE_PROXY = 4; /** - * Returns a list of supported image properties for this {@link Image}. + * Returns a list of supported image properties for this {@link MPImage}. * - *

    Currently {@link Image} only support single storage type so the size of return list will + *

    Currently {@link MPImage} only support single storage type so the size of return list will * always be 1. * - * @see ImageProperties + * @see MPImageProperties */ - public List getContainedImageProperties() { + public List getContainedImageProperties() { return Collections.singletonList(getContainer().getImageProperties()); } @@ -124,7 +124,7 @@ public class Image implements Closeable { return height; } - /** Acquires a reference on this {@link Image}. This will increase the reference count by 1. */ + /** Acquires a reference on this {@link MPImage}. This will increase the reference count by 1. */ private synchronized void acquire() { referenceCount += 1; } @@ -132,7 +132,7 @@ public class Image implements Closeable { /** * Removes a reference that was previously acquired or init. * - *

    When {@link Image} is created, it has 1 reference count. + *

    When {@link MPImage} is created, it has 1 reference count. * *

    When the reference count becomes 0, it will release the resource under the hood. */ @@ -141,24 +141,24 @@ public class Image implements Closeable { public synchronized void close() { referenceCount -= 1; if (referenceCount == 0) { - for (ImageContainer imageContainer : containerMap.values()) { + for (MPImageContainer imageContainer : containerMap.values()) { imageContainer.close(); } } } - /** Advanced API access for {@link Image}. */ + /** Advanced API access for {@link MPImage}. */ static final class Internal { /** - * Acquires a reference on this {@link Image}. This will increase the reference count by 1. + * Acquires a reference on this {@link MPImage}. This will increase the reference count by 1. * *

    This method is more useful for image consumer to acquire a reference so image resource * will not be closed accidentally. As image creator, normal developer doesn't need to call this * method. * - *

    The reference count is 1 when {@link Image} is created. Developer can call {@link - * #close()} to indicate it doesn't need this {@link Image} anymore. + *

    The reference count is 1 when {@link MPImage} is created. Developer can call {@link + * #close()} to indicate it doesn't need this {@link MPImage} anymore. * * @see #close() */ @@ -166,10 +166,10 @@ public class Image implements Closeable { image.acquire(); } - private final Image image; + private final MPImage image; - // Only Image creates the internal helper. - private Internal(Image image) { + // Only MPImage creates the internal helper. + private Internal(MPImage image) { this.image = image; } } @@ -179,15 +179,15 @@ public class Image implements Closeable { return new Internal(this); } - private final Map containerMap; + private final Map containerMap; private final long timestamp; private final int width; private final int height; private int referenceCount; - /** Constructs an {@link Image} with a built container. */ - Image(ImageContainer container, long timestamp, int width, int height) { + /** Constructs a {@link MPImage} with a built container. */ + MPImage(MPImageContainer container, long timestamp, int width, int height) { this.containerMap = new HashMap<>(); containerMap.put(container.getImageProperties(), container); this.timestamp = timestamp; @@ -201,10 +201,10 @@ public class Image implements Closeable { * * @return the current container. */ - ImageContainer getContainer() { + MPImageContainer getContainer() { // According to the design, in the future we will support multiple containers in one image. // Currently just return the original container. - // TODO: Cache multiple containers in Image. + // TODO: Cache multiple containers in MPImage. return containerMap.values().iterator().next(); } @@ -214,8 +214,8 @@ public class Image implements Closeable { *

    If there are multiple containers with required {@code storageType}, returns the first one. */ @Nullable - ImageContainer getContainer(@StorageType int storageType) { - for (Entry entry : containerMap.entrySet()) { + MPImageContainer getContainer(@StorageType int storageType) { + for (Entry entry : containerMap.entrySet()) { if (entry.getKey().getStorageType() == storageType) { return entry.getValue(); } @@ -225,13 +225,13 @@ public class Image implements Closeable { /** Gets container from required {@code imageProperties}. Returns {@code null} if non existed. */ @Nullable - ImageContainer getContainer(ImageProperties imageProperties) { + MPImageContainer getContainer(MPImageProperties imageProperties) { return containerMap.get(imageProperties); } /** Adds a new container if it doesn't exist. Returns {@code true} if it succeeds. */ - boolean addContainer(ImageContainer container) { - ImageProperties imageProperties = container.getImageProperties(); + boolean addContainer(MPImageContainer container) { + MPImageProperties imageProperties = container.getImageProperties(); if (containerMap.containsKey(imageProperties)) { return false; } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/ImageConsumer.java b/mediapipe/java/com/google/mediapipe/framework/image/MPImageConsumer.java similarity index 87% rename from mediapipe/java/com/google/mediapipe/framework/image/ImageConsumer.java rename to mediapipe/java/com/google/mediapipe/framework/image/MPImageConsumer.java index 18eed68c..f9f343e9 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/ImageConsumer.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/MPImageConsumer.java @@ -14,14 +14,14 @@ limitations under the License. ==============================================================================*/ package com.google.mediapipe.framework.image; -/** Lightweight abstraction for an object that can receive {@link Image} */ -public interface ImageConsumer { +/** Lightweight abstraction for an object that can receive {@link MPImage} */ +public interface MPImageConsumer { /** - * Called when an {@link Image} is available. + * Called when a {@link MPImage} is available. * *

    The argument is only guaranteed to be available until this method returns. if you need to * extend its life time, acquire it, then release it when done. */ - void onNewImage(Image image); + void onNewMPImage(MPImage image); } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/ImageContainer.java b/mediapipe/java/com/google/mediapipe/framework/image/MPImageContainer.java similarity index 93% rename from mediapipe/java/com/google/mediapipe/framework/image/ImageContainer.java rename to mediapipe/java/com/google/mediapipe/framework/image/MPImageContainer.java index 727ec089..674073b5 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/ImageContainer.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/MPImageContainer.java @@ -16,9 +16,9 @@ limitations under the License. package com.google.mediapipe.framework.image; /** Manages internal image data storage. The interface is package-private. */ -interface ImageContainer { +interface MPImageContainer { /** Returns the properties of the contained image. */ - ImageProperties getImageProperties(); + MPImageProperties getImageProperties(); /** Close the image container and releases the image resource inside. */ void close(); diff --git a/mediapipe/java/com/google/mediapipe/framework/image/ImageProducer.java b/mediapipe/java/com/google/mediapipe/framework/image/MPImageProducer.java similarity index 75% rename from mediapipe/java/com/google/mediapipe/framework/image/ImageProducer.java rename to mediapipe/java/com/google/mediapipe/framework/image/MPImageProducer.java index 4f3641d6..9783935d 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/ImageProducer.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/MPImageProducer.java @@ -14,9 +14,9 @@ limitations under the License. ==============================================================================*/ package com.google.mediapipe.framework.image; -/** Lightweight abstraction for an object that produce {@link Image} */ -public interface ImageProducer { +/** Lightweight abstraction for an object that produce {@link MPImage} */ +public interface MPImageProducer { - /** Sets the consumer that receives the {@link Image}. */ - void setImageConsumer(ImageConsumer imageConsumer); + /** Sets the consumer that receives the {@link MPImage}. */ + void setMPImageConsumer(MPImageConsumer imageConsumer); } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/ImageProperties.java b/mediapipe/java/com/google/mediapipe/framework/image/MPImageProperties.java similarity index 63% rename from mediapipe/java/com/google/mediapipe/framework/image/ImageProperties.java rename to mediapipe/java/com/google/mediapipe/framework/image/MPImageProperties.java index e33b33e7..6005ce77 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/ImageProperties.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/MPImageProperties.java @@ -17,25 +17,25 @@ package com.google.mediapipe.framework.image; import com.google.auto.value.AutoValue; import com.google.auto.value.extension.memoized.Memoized; -import com.google.mediapipe.framework.image.Image.ImageFormat; -import com.google.mediapipe.framework.image.Image.StorageType; +import com.google.mediapipe.framework.image.MPImage.MPImageFormat; +import com.google.mediapipe.framework.image.MPImage.StorageType; /** Groups a set of properties to describe how an image is stored. */ @AutoValue -public abstract class ImageProperties { +public abstract class MPImageProperties { /** * Gets the pixel format of the image. * - * @see Image.ImageFormat + * @see MPImage.MPImageFormat */ - @ImageFormat + @MPImageFormat public abstract int getImageFormat(); /** * Gets the storage type of the image. * - * @see Image.StorageType + * @see MPImage.StorageType */ @StorageType public abstract int getStorageType(); @@ -45,36 +45,36 @@ public abstract class ImageProperties { public abstract int hashCode(); /** - * Creates a builder of {@link ImageProperties}. + * Creates a builder of {@link MPImageProperties}. * - * @see ImageProperties.Builder + * @see MPImageProperties.Builder */ static Builder builder() { - return new AutoValue_ImageProperties.Builder(); + return new AutoValue_MPImageProperties.Builder(); } - /** Builds a {@link ImageProperties}. */ + /** Builds a {@link MPImageProperties}. */ @AutoValue.Builder abstract static class Builder { /** - * Sets the {@link Image.ImageFormat}. + * Sets the {@link MPImage.MPImageFormat}. * - * @see ImageProperties#getImageFormat + * @see MPImageProperties#getImageFormat */ - abstract Builder setImageFormat(@ImageFormat int value); + abstract Builder setImageFormat(@MPImageFormat int value); /** - * Sets the {@link Image.StorageType}. + * Sets the {@link MPImage.StorageType}. * - * @see ImageProperties#getStorageType + * @see MPImageProperties#getStorageType */ abstract Builder setStorageType(@StorageType int value); - /** Builds the {@link ImageProperties}. */ - abstract ImageProperties build(); + /** Builds the {@link MPImageProperties}. */ + abstract MPImageProperties build(); } // Hide the constructor. - ImageProperties() {} + MPImageProperties() {} } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/MediaImageBuilder.java b/mediapipe/java/com/google/mediapipe/framework/image/MediaImageBuilder.java index e351a87f..9e719715 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/MediaImageBuilder.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/MediaImageBuilder.java @@ -15,11 +15,12 @@ limitations under the License. package com.google.mediapipe.framework.image; +import android.media.Image; import android.os.Build.VERSION_CODES; import androidx.annotation.RequiresApi; /** - * Builds {@link Image} from {@link android.media.Image}. + * Builds {@link MPImage} from {@link android.media.Image}. * *

    Once {@link android.media.Image} is passed in, to keep data integrity you shouldn't modify * content in it. @@ -30,7 +31,7 @@ import androidx.annotation.RequiresApi; public class MediaImageBuilder { // Mandatory fields. - private final android.media.Image mediaImage; + private final Image mediaImage; // Optional fields. private long timestamp; @@ -40,20 +41,20 @@ public class MediaImageBuilder { * * @param mediaImage image data object. */ - public MediaImageBuilder(android.media.Image mediaImage) { + public MediaImageBuilder(Image mediaImage) { this.mediaImage = mediaImage; this.timestamp = 0; } - /** Sets value for {@link Image#getTimestamp()}. */ + /** Sets value for {@link MPImage#getTimestamp()}. */ MediaImageBuilder setTimestamp(long timestamp) { this.timestamp = timestamp; return this; } - /** Builds an {@link Image} instance. */ - public Image build() { - return new Image( + /** Builds a {@link MPImage} instance. */ + public MPImage build() { + return new MPImage( new MediaImageContainer(mediaImage), timestamp, mediaImage.getWidth(), diff --git a/mediapipe/java/com/google/mediapipe/framework/image/MediaImageContainer.java b/mediapipe/java/com/google/mediapipe/framework/image/MediaImageContainer.java index 144b64de..864c76df 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/MediaImageContainer.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/MediaImageContainer.java @@ -15,33 +15,34 @@ limitations under the License. package com.google.mediapipe.framework.image; +import android.media.Image; import android.os.Build; import android.os.Build.VERSION; import android.os.Build.VERSION_CODES; import androidx.annotation.RequiresApi; -import com.google.mediapipe.framework.image.Image.ImageFormat; +import com.google.mediapipe.framework.image.MPImage.MPImageFormat; @RequiresApi(VERSION_CODES.KITKAT) -class MediaImageContainer implements ImageContainer { +class MediaImageContainer implements MPImageContainer { - private final android.media.Image mediaImage; - private final ImageProperties properties; + private final Image mediaImage; + private final MPImageProperties properties; - public MediaImageContainer(android.media.Image mediaImage) { + public MediaImageContainer(Image mediaImage) { this.mediaImage = mediaImage; this.properties = - ImageProperties.builder() - .setStorageType(Image.STORAGE_TYPE_MEDIA_IMAGE) + MPImageProperties.builder() + .setStorageType(MPImage.STORAGE_TYPE_MEDIA_IMAGE) .setImageFormat(convertFormatCode(mediaImage.getFormat())) .build(); } - public android.media.Image getImage() { + public Image getImage() { return mediaImage; } @Override - public ImageProperties getImageProperties() { + public MPImageProperties getImageProperties() { return properties; } @@ -50,24 +51,24 @@ class MediaImageContainer implements ImageContainer { mediaImage.close(); } - @ImageFormat + @MPImageFormat static int convertFormatCode(int graphicsFormat) { // We only cover the format mentioned in // https://developer.android.com/reference/android/media/Image#getFormat() if (VERSION.SDK_INT >= Build.VERSION_CODES.M) { if (graphicsFormat == android.graphics.ImageFormat.FLEX_RGBA_8888) { - return Image.IMAGE_FORMAT_RGBA; + return MPImage.IMAGE_FORMAT_RGBA; } else if (graphicsFormat == android.graphics.ImageFormat.FLEX_RGB_888) { - return Image.IMAGE_FORMAT_RGB; + return MPImage.IMAGE_FORMAT_RGB; } } switch (graphicsFormat) { case android.graphics.ImageFormat.JPEG: - return Image.IMAGE_FORMAT_JPEG; + return MPImage.IMAGE_FORMAT_JPEG; case android.graphics.ImageFormat.YUV_420_888: - return Image.IMAGE_FORMAT_YUV_420_888; + return MPImage.IMAGE_FORMAT_YUV_420_888; default: - return Image.IMAGE_FORMAT_UNKNOWN; + return MPImage.IMAGE_FORMAT_UNKNOWN; } } } diff --git a/mediapipe/java/com/google/mediapipe/framework/image/MediaImageExtractor.java b/mediapipe/java/com/google/mediapipe/framework/image/MediaImageExtractor.java index 718cb471..76bb5a5e 100644 --- a/mediapipe/java/com/google/mediapipe/framework/image/MediaImageExtractor.java +++ b/mediapipe/java/com/google/mediapipe/framework/image/MediaImageExtractor.java @@ -15,13 +15,14 @@ limitations under the License. package com.google.mediapipe.framework.image; +import android.media.Image; import android.os.Build.VERSION_CODES; import androidx.annotation.RequiresApi; /** - * Utility for extracting {@link android.media.Image} from {@link Image}. + * Utility for extracting {@link android.media.Image} from {@link MPImage}. * - *

    Currently it only supports {@link Image} with {@link Image#STORAGE_TYPE_MEDIA_IMAGE}, + *

    Currently it only supports {@link MPImage} with {@link MPImage#STORAGE_TYPE_MEDIA_IMAGE}, * otherwise {@link IllegalArgumentException} will be thrown. */ @RequiresApi(VERSION_CODES.KITKAT) @@ -30,20 +31,20 @@ public class MediaImageExtractor { private MediaImageExtractor() {} /** - * Extracts a {@link android.media.Image} from an {@link Image}. Currently it only works for - * {@link Image} that built from {@link MediaImageBuilder}. + * Extracts a {@link android.media.Image} from a {@link MPImage}. Currently it only works for + * {@link MPImage} that built from {@link MediaImageBuilder}. * * @param image the image to extract {@link android.media.Image} from. - * @return {@link android.media.Image} that stored in {@link Image}. + * @return {@link android.media.Image} that stored in {@link MPImage}. * @throws IllegalArgumentException if the extraction failed. */ - public static android.media.Image extract(Image image) { - ImageContainer container; - if ((container = image.getContainer(Image.STORAGE_TYPE_MEDIA_IMAGE)) != null) { + public static Image extract(MPImage image) { + MPImageContainer container; + if ((container = image.getContainer(MPImage.STORAGE_TYPE_MEDIA_IMAGE)) != null) { return ((MediaImageContainer) container).getImage(); } throw new IllegalArgumentException( - "Extract Media Image from an Image created by objects other than Media Image" + "Extract Media Image from a MPImage created by objects other than Media Image" + " is not supported"); } } diff --git a/mediapipe/java/com/google/mediapipe/mediapipe_aar.bzl b/mediapipe/java/com/google/mediapipe/mediapipe_aar.bzl index 9b01e2f0..645e8b72 100644 --- a/mediapipe/java/com/google/mediapipe/mediapipe_aar.bzl +++ b/mediapipe/java/com/google/mediapipe/mediapipe_aar.bzl @@ -1,4 +1,4 @@ -# Copyright 2019-2020 The MediaPipe Authors. +# Copyright 2019-2022 The MediaPipe Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -209,9 +209,9 @@ def _mediapipe_jni(name, gen_libmediapipe, calculators = []): def mediapipe_build_aar_with_jni(name, android_library): """Builds MediaPipe AAR with jni. - Args: - name: The bazel target name. - android_library: the android library that contains jni. + Args: + name: The bazel target name. + android_library: the android library that contains jni. """ # Generates dummy AndroidManifest.xml for dummy apk usage @@ -328,19 +328,14 @@ def mediapipe_java_proto_srcs(name = ""): src_out = "com/google/mediapipe/proto/MediaPipeOptionsProto.java", )) - proto_src_list.append(mediapipe_java_proto_src_extractor( - target = "//mediapipe/framework/formats:landmark_java_proto_lite", - src_out = "com/google/mediapipe/formats/proto/LandmarkProto.java", - )) - proto_src_list.append(mediapipe_java_proto_src_extractor( target = "//mediapipe/framework/formats/annotation:rasterization_java_proto_lite", src_out = "com/google/mediapipe/formats/annotation/proto/RasterizationProto.java", )) proto_src_list.append(mediapipe_java_proto_src_extractor( - target = "//mediapipe/framework/formats:location_data_java_proto_lite", - src_out = "com/google/mediapipe/formats/proto/LocationDataProto.java", + target = "//mediapipe/framework/formats:classification_java_proto_lite", + src_out = "com/google/mediapipe/formats/proto/ClassificationProto.java", )) proto_src_list.append(mediapipe_java_proto_src_extractor( @@ -349,8 +344,18 @@ def mediapipe_java_proto_srcs(name = ""): )) proto_src_list.append(mediapipe_java_proto_src_extractor( - target = "//mediapipe/framework/formats:classification_java_proto_lite", - src_out = "com/google/mediapipe/formats/proto/ClassificationProto.java", + target = "//mediapipe/framework/formats:landmark_java_proto_lite", + src_out = "com/google/mediapipe/formats/proto/LandmarkProto.java", + )) + + proto_src_list.append(mediapipe_java_proto_src_extractor( + target = "//mediapipe/framework/formats:location_data_java_proto_lite", + src_out = "com/google/mediapipe/formats/proto/LocationDataProto.java", + )) + + proto_src_list.append(mediapipe_java_proto_src_extractor( + target = "//mediapipe/framework/formats:rect_java_proto_lite", + src_out = "com/google/mediapipe/formats/proto/RectProto.java", )) return proto_src_list diff --git a/mediapipe/model_maker/python/core/BUILD b/mediapipe/model_maker/python/core/BUILD index 9f205bb1..636a1a72 100644 --- a/mediapipe/model_maker/python/core/BUILD +++ b/mediapipe/model_maker/python/core/BUILD @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Placeholder for internal Python strict library compatibility macro. +# Placeholder for internal Python strict library and test compatibility macro. package( default_visibility = ["//mediapipe:__subpackages__"], diff --git a/mediapipe/model_maker/python/core/data/BUILD b/mediapipe/model_maker/python/core/data/BUILD index c4c659d5..70a62e8f 100644 --- a/mediapipe/model_maker/python/core/data/BUILD +++ b/mediapipe/model_maker/python/core/data/BUILD @@ -13,6 +13,7 @@ # limitations under the License. # Placeholder for internal Python strict library and test compatibility macro. +# Placeholder for internal Python strict test compatibility macro. licenses(["notice"]) @@ -23,15 +24,12 @@ package( py_library( name = "data_util", srcs = ["data_util.py"], - srcs_version = "PY3", ) py_test( name = "data_util_test", srcs = ["data_util_test.py"], data = ["//mediapipe/model_maker/python/core/data/testdata"], - python_version = "PY3", - srcs_version = "PY3", deps = [":data_util"], ) @@ -44,8 +42,6 @@ py_library( py_test( name = "dataset_test", srcs = ["dataset_test.py"], - python_version = "PY3", - srcs_version = "PY3", deps = [ ":dataset", "//mediapipe/model_maker/python/core/utils:test_util", @@ -55,14 +51,11 @@ py_test( py_library( name = "classification_dataset", srcs = ["classification_dataset.py"], - srcs_version = "PY3", deps = [":dataset"], ) py_test( name = "classification_dataset_test", srcs = ["classification_dataset_test.py"], - python_version = "PY3", - srcs_version = "PY3", deps = [":classification_dataset"], ) diff --git a/mediapipe/model_maker/python/core/data/classification_dataset.py b/mediapipe/model_maker/python/core/data/classification_dataset.py index 9075e46e..073e7963 100644 --- a/mediapipe/model_maker/python/core/data/classification_dataset.py +++ b/mediapipe/model_maker/python/core/data/classification_dataset.py @@ -13,7 +13,7 @@ # limitations under the License. """Common classification dataset library.""" -from typing import Any, Tuple +from typing import List, Tuple import tensorflow as tf @@ -21,15 +21,20 @@ from mediapipe.model_maker.python.core.data import dataset as ds class ClassificationDataset(ds.Dataset): - """DataLoader for classification models.""" + """Dataset Loader for classification models.""" - def __init__(self, dataset: tf.data.Dataset, size: int, index_to_label: Any): + def __init__(self, dataset: tf.data.Dataset, size: int, + label_names: List[str]): super().__init__(dataset, size) - self.index_to_label = index_to_label + self._label_names = label_names @property def num_classes(self: ds._DatasetT) -> int: - return len(self.index_to_label) + return len(self._label_names) + + @property + def label_names(self: ds._DatasetT) -> List[str]: + return self._label_names def split(self: ds._DatasetT, fraction: float) -> Tuple[ds._DatasetT, ds._DatasetT]: @@ -44,4 +49,4 @@ class ClassificationDataset(ds.Dataset): Returns: The splitted two sub datasets. """ - return self._split(fraction, self.index_to_label) + return self._split(fraction, self._label_names) diff --git a/mediapipe/model_maker/python/core/data/classification_dataset_test.py b/mediapipe/model_maker/python/core/data/classification_dataset_test.py index f8688ab1..82e74b04 100644 --- a/mediapipe/model_maker/python/core/data/classification_dataset_test.py +++ b/mediapipe/model_maker/python/core/data/classification_dataset_test.py @@ -12,45 +12,55 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any, List, Tuple, TypeVar + # Dependency imports import tensorflow as tf from mediapipe.model_maker.python.core.data import classification_dataset +_DatasetT = TypeVar( + '_DatasetT', bound='ClassificationDatasetTest.MagicClassificationDataset') -class ClassificationDataLoaderTest(tf.test.TestCase): + +class ClassificationDatasetTest(tf.test.TestCase): def test_split(self): - class MagicClassificationDataLoader( + class MagicClassificationDataset( classification_dataset.ClassificationDataset): + """A mock classification dataset class for testing purpose. - def __init__(self, dataset, size, index_to_label, value): - super(MagicClassificationDataLoader, - self).__init__(dataset, size, index_to_label) + Attributes: + value: A value variable stored by the mock dataset class for testing. + """ + + def __init__(self, dataset: tf.data.Dataset, size: int, + label_names: List[str], value: Any): + super().__init__(dataset=dataset, size=size, label_names=label_names) self.value = value - def split(self, fraction): - return self._split(fraction, self.index_to_label, self.value) + def split(self, fraction: float) -> Tuple[_DatasetT, _DatasetT]: + return self._split(fraction, self.label_names, self.value) # Some dummy inputs. magic_value = 42 num_classes = 2 - index_to_label = (False, True) + label_names = ['foo', 'bar'] # Create data loader from sample data. ds = tf.data.Dataset.from_tensor_slices([[0, 1], [1, 1], [0, 0], [1, 0]]) - data = MagicClassificationDataLoader(ds, len(ds), index_to_label, - magic_value) + data = MagicClassificationDataset( + dataset=ds, size=len(ds), label_names=label_names, value=magic_value) # Train/Test data split. fraction = .25 - train_data, test_data = data.split(fraction) + train_data, test_data = data.split(fraction=fraction) # `split` should return instances of child DataLoader. - self.assertIsInstance(train_data, MagicClassificationDataLoader) - self.assertIsInstance(test_data, MagicClassificationDataLoader) + self.assertIsInstance(train_data, MagicClassificationDataset) + self.assertIsInstance(test_data, MagicClassificationDataset) # Make sure number of entries are right. self.assertEqual(len(train_data.gen_tf_dataset()), len(train_data)) @@ -59,7 +69,7 @@ class ClassificationDataLoaderTest(tf.test.TestCase): # Make sure attributes propagated correctly. self.assertEqual(train_data.num_classes, num_classes) - self.assertEqual(test_data.index_to_label, index_to_label) + self.assertEqual(test_data.label_names, label_names) self.assertEqual(train_data.value, magic_value) self.assertEqual(test_data.value, magic_value) diff --git a/mediapipe/model_maker/python/core/tasks/BUILD b/mediapipe/model_maker/python/core/tasks/BUILD index b3588f0b..124de621 100644 --- a/mediapipe/model_maker/python/core/tasks/BUILD +++ b/mediapipe/model_maker/python/core/tasks/BUILD @@ -13,6 +13,7 @@ # limitations under the License. # Placeholder for internal Python strict library and test compatibility macro. +# Placeholder for internal Python strict test compatibility macro. package( default_visibility = ["//mediapipe:__subpackages__"], @@ -23,7 +24,6 @@ licenses(["notice"]) py_library( name = "custom_model", srcs = ["custom_model.py"], - srcs_version = "PY3", deps = [ "//mediapipe/model_maker/python/core/data:dataset", "//mediapipe/model_maker/python/core/utils:model_util", @@ -34,8 +34,6 @@ py_library( py_test( name = "custom_model_test", srcs = ["custom_model_test.py"], - python_version = "PY3", - srcs_version = "PY3", deps = [ ":custom_model", "//mediapipe/model_maker/python/core/utils:test_util", @@ -45,7 +43,6 @@ py_test( py_library( name = "classifier", srcs = ["classifier.py"], - srcs_version = "PY3", deps = [ ":custom_model", "//mediapipe/model_maker/python/core/data:dataset", @@ -55,8 +52,6 @@ py_library( py_test( name = "classifier_test", srcs = ["classifier_test.py"], - python_version = "PY3", - srcs_version = "PY3", deps = [ ":classifier", "//mediapipe/model_maker/python/core/utils:test_util", diff --git a/mediapipe/model_maker/python/core/tasks/classifier.py b/mediapipe/model_maker/python/core/tasks/classifier.py index 6b366f6d..f83d4059 100644 --- a/mediapipe/model_maker/python/core/tasks/classifier.py +++ b/mediapipe/model_maker/python/core/tasks/classifier.py @@ -29,22 +29,22 @@ from mediapipe.model_maker.python.core.tasks import custom_model class Classifier(custom_model.CustomModel): """An abstract base class that represents a TensorFlow classifier.""" - def __init__(self, model_spec: Any, index_to_label: List[str], shuffle: bool, + def __init__(self, model_spec: Any, label_names: List[str], shuffle: bool, full_train: bool): """Initilizes a classifier with its specifications. Args: model_spec: Specification for the model. - index_to_label: A list that map from index to label class name. + label_names: A list of label names for the classes. shuffle: Whether the dataset should be shuffled. full_train: If true, train the model end-to-end including the backbone and the classification layers on top. Otherwise, only train the top classification layers. """ super(Classifier, self).__init__(model_spec, shuffle) - self._index_to_label = index_to_label + self._label_names = label_names self._full_train = full_train - self._num_classes = len(index_to_label) + self._num_classes = len(label_names) def evaluate(self, data: dataset.Dataset, batch_size: int = 32) -> Any: """Evaluates the classifier with the provided evaluation dataset. @@ -74,4 +74,4 @@ class Classifier(custom_model.CustomModel): label_filepath = os.path.join(export_dir, label_filename) tf.compat.v1.logging.info('Saving labels in %s', label_filepath) with tf.io.gfile.GFile(label_filepath, 'w') as f: - f.write('\n'.join(self._index_to_label)) + f.write('\n'.join(self._label_names)) diff --git a/mediapipe/model_maker/python/core/tasks/classifier_test.py b/mediapipe/model_maker/python/core/tasks/classifier_test.py index fbf231d8..52a3b97d 100644 --- a/mediapipe/model_maker/python/core/tasks/classifier_test.py +++ b/mediapipe/model_maker/python/core/tasks/classifier_test.py @@ -36,10 +36,10 @@ class ClassifierTest(tf.test.TestCase): def setUp(self): super(ClassifierTest, self).setUp() - index_to_label = ['cat', 'dog'] + label_names = ['cat', 'dog'] self.model = MockClassifier( model_spec=None, - index_to_label=index_to_label, + label_names=label_names, shuffle=False, full_train=False) self.model.model = test_util.build_model(input_shape=[4], num_classes=2) diff --git a/mediapipe/model_maker/python/core/tasks/custom_model.py b/mediapipe/model_maker/python/core/tasks/custom_model.py index 2cea4e0a..66d1494d 100644 --- a/mediapipe/model_maker/python/core/tasks/custom_model.py +++ b/mediapipe/model_maker/python/core/tasks/custom_model.py @@ -21,8 +21,6 @@ import abc import os from typing import Any, Callable, Optional -# Dependency imports - import tensorflow as tf from mediapipe.model_maker.python.core.data import dataset @@ -77,9 +75,9 @@ class CustomModel(abc.ABC): tflite_filepath = os.path.join(export_dir, tflite_filename) # TODO: Populate metadata to the exported TFLite model. model_util.export_tflite( - self._model, - tflite_filepath, - quantization_config, + model=self._model, + tflite_filepath=tflite_filepath, + quantization_config=quantization_config, preprocess=preprocess) tf.compat.v1.logging.info( 'TensorFlow Lite model exported successfully: %s' % tflite_filepath) diff --git a/mediapipe/model_maker/python/core/tasks/custom_model_test.py b/mediapipe/model_maker/python/core/tasks/custom_model_test.py index e693e127..ad77d4ec 100644 --- a/mediapipe/model_maker/python/core/tasks/custom_model_test.py +++ b/mediapipe/model_maker/python/core/tasks/custom_model_test.py @@ -40,8 +40,8 @@ class CustomModelTest(tf.test.TestCase): def setUp(self): super(CustomModelTest, self).setUp() - self.model = MockCustomModel(model_spec=None, shuffle=False) - self.model._model = test_util.build_model(input_shape=[4], num_classes=2) + self._model = MockCustomModel(model_spec=None, shuffle=False) + self._model._model = test_util.build_model(input_shape=[4], num_classes=2) def _check_nonempty_file(self, filepath): self.assertTrue(os.path.isfile(filepath)) @@ -49,7 +49,7 @@ class CustomModelTest(tf.test.TestCase): def test_export_tflite(self): export_path = os.path.join(self.get_temp_dir(), 'export/') - self.model.export_tflite(export_dir=export_path) + self._model.export_tflite(export_dir=export_path) self._check_nonempty_file(os.path.join(export_path, 'model.tflite')) if __name__ == '__main__': diff --git a/mediapipe/model_maker/python/core/utils/BUILD b/mediapipe/model_maker/python/core/utils/BUILD index e4b18b39..a2ec5204 100644 --- a/mediapipe/model_maker/python/core/utils/BUILD +++ b/mediapipe/model_maker/python/core/utils/BUILD @@ -13,6 +13,7 @@ # limitations under the License. # Placeholder for internal Python strict library and test compatibility macro. +# Placeholder for internal Python strict test compatibility macro. licenses(["notice"]) @@ -24,31 +25,15 @@ py_library( name = "test_util", testonly = 1, srcs = ["test_util.py"], - srcs_version = "PY3", deps = [ ":model_util", "//mediapipe/model_maker/python/core/data:dataset", ], ) -py_library( - name = "image_preprocessing", - srcs = ["image_preprocessing.py"], - srcs_version = "PY3", -) - -py_test( - name = "image_preprocessing_test", - srcs = ["image_preprocessing_test.py"], - python_version = "PY3", - srcs_version = "PY3", - deps = [":image_preprocessing"], -) - py_library( name = "model_util", srcs = ["model_util.py"], - srcs_version = "PY3", deps = [ ":quantization", "//mediapipe/model_maker/python/core/data:dataset", @@ -58,8 +43,6 @@ py_library( py_test( name = "model_util_test", srcs = ["model_util_test.py"], - python_version = "PY3", - srcs_version = "PY3", deps = [ ":model_util", ":quantization", @@ -76,8 +59,6 @@ py_library( py_test( name = "loss_functions_test", srcs = ["loss_functions_test.py"], - python_version = "PY3", - srcs_version = "PY3", deps = [":loss_functions"], ) @@ -91,8 +72,6 @@ py_library( py_test( name = "quantization_test", srcs = ["quantization_test.py"], - python_version = "PY3", - srcs_version = "PY3", deps = [ ":quantization", ":test_util", diff --git a/mediapipe/model_maker/python/core/utils/loss_functions.py b/mediapipe/model_maker/python/core/utils/loss_functions.py index 17c738a1..5b0aa32b 100644 --- a/mediapipe/model_maker/python/core/utils/loss_functions.py +++ b/mediapipe/model_maker/python/core/utils/loss_functions.py @@ -56,7 +56,7 @@ class FocalLoss(tf.keras.losses.Loss): class_weight: A weight to apply to the loss, one for each class. The weight is applied for each input where the ground truth label matches. """ - super(tf.keras.losses.Loss, self).__init__() + super().__init__() # Used for clipping min/max values of probability values in y_pred to avoid # NaNs and Infs in computation. self._epsilon = 1e-7 diff --git a/mediapipe/model_maker/python/core/utils/model_util.py b/mediapipe/model_maker/python/core/utils/model_util.py index 0899a9b1..e1228eb6 100644 --- a/mediapipe/model_maker/python/core/utils/model_util.py +++ b/mediapipe/model_maker/python/core/utils/model_util.py @@ -104,8 +104,8 @@ def export_tflite( quantization_config: Configuration for post-training quantization. supported_ops: A list of supported ops in the converted TFLite file. preprocess: A callable to preprocess the representative dataset for - quantization. The callable takes three arguments in order: feature, - label, and is_training. + quantization. The callable takes three arguments in order: feature, label, + and is_training. """ if tflite_filepath is None: raise ValueError( diff --git a/mediapipe/model_maker/python/core/utils/model_util_test.py b/mediapipe/model_maker/python/core/utils/model_util_test.py index ce31c187..35b52eb7 100644 --- a/mediapipe/model_maker/python/core/utils/model_util_test.py +++ b/mediapipe/model_maker/python/core/utils/model_util_test.py @@ -100,7 +100,8 @@ class ModelUtilTest(tf.test.TestCase, parameterized.TestCase): model = test_util.build_model(input_shape=[input_dim], num_classes=2) tflite_file = os.path.join(self.get_temp_dir(), 'model.tflite') model_util.export_tflite(model, tflite_file) - self._test_tflite(model, tflite_file, input_dim) + test_util.test_tflite( + keras_model=model, tflite_file=tflite_file, size=[1, input_dim]) @parameterized.named_parameters( dict( @@ -121,27 +122,20 @@ class ModelUtilTest(tf.test.TestCase, parameterized.TestCase): input_dim = 16 num_classes = 2 max_input_value = 5 - model = test_util.build_model([input_dim], num_classes) + model = test_util.build_model( + input_shape=[input_dim], num_classes=num_classes) tflite_file = os.path.join(self.get_temp_dir(), 'model_quantized.tflite') - model_util.export_tflite(model, tflite_file, config) - self._test_tflite( - model, tflite_file, input_dim, max_input_value, atol=1e-00) - self.assertNear(os.path.getsize(tflite_file), model_size, 300) - - def _test_tflite(self, - keras_model: tf.keras.Model, - tflite_model_file: str, - input_dim: int, - max_input_value: int = 1000, - atol: float = 1e-04): - random_input = test_util.create_random_sample( - size=[1, input_dim], high=max_input_value) - random_input = tf.convert_to_tensor(random_input) - + model_util.export_tflite( + model=model, tflite_filepath=tflite_file, quantization_config=config) self.assertTrue( - test_util.is_same_output( - tflite_model_file, keras_model, random_input, atol=atol)) + test_util.test_tflite( + keras_model=model, + tflite_file=tflite_file, + size=[1, input_dim], + high=max_input_value, + atol=1e-00)) + self.assertNear(os.path.getsize(tflite_file), model_size, 300) if __name__ == '__main__': diff --git a/mediapipe/model_maker/python/core/utils/test_util.py b/mediapipe/model_maker/python/core/utils/test_util.py index cac2a0e1..b402d379 100644 --- a/mediapipe/model_maker/python/core/utils/test_util.py +++ b/mediapipe/model_maker/python/core/utils/test_util.py @@ -92,3 +92,32 @@ def is_same_output(tflite_file: str, keras_output = keras_model.predict_on_batch(input_tensors) return np.allclose(lite_output, keras_output, atol=atol) + + +def test_tflite(keras_model: tf.keras.Model, + tflite_file: str, + size: Union[int, List[int]], + high: float = 1, + atol: float = 1e-04) -> bool: + """Verifies if the output of TFLite model and TF Keras model are identical. + + Args: + keras_model: Input TensorFlow Keras model. + tflite_file: Input TFLite model file. + size: Size of the input tesnor. + high: Higher boundary of the values in input tensors. + atol: Absolute tolerance of the difference between the outputs of Keras + model and TFLite model. + + Returns: + True if the output of TFLite model and TF Keras model are identical. + Otherwise, False. + """ + random_input = create_random_sample(size=size, high=high) + random_input = tf.convert_to_tensor(random_input) + + return is_same_output( + tflite_file=tflite_file, + keras_model=keras_model, + input_tensors=random_input, + atol=atol) diff --git a/mediapipe/model_maker/python/internal/README.md b/mediapipe/model_maker/python/internal/README.md new file mode 100644 index 00000000..100d6a52 --- /dev/null +++ b/mediapipe/model_maker/python/internal/README.md @@ -0,0 +1,4 @@ +# MediaPipe Model Maker Internal Library + +This directory contains model maker library for internal users and experimental +purposes. diff --git a/mediapipe/model_maker/python/internal/__init__.py b/mediapipe/model_maker/python/internal/__init__.py new file mode 100644 index 00000000..05f41d8a --- /dev/null +++ b/mediapipe/model_maker/python/internal/__init__.py @@ -0,0 +1 @@ +"""Model maker internal library.""" diff --git a/mediapipe/model_maker/python/vision/core/BUILD b/mediapipe/model_maker/python/vision/core/BUILD new file mode 100644 index 00000000..0b15a027 --- /dev/null +++ b/mediapipe/model_maker/python/vision/core/BUILD @@ -0,0 +1,33 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Placeholder for internal Python strict library and test compatibility macro. +# Placeholder for internal Python strict test compatibility macro. + +licenses(["notice"]) + +package( + default_visibility = ["//mediapipe:__subpackages__"], +) + +py_library( + name = "image_preprocessing", + srcs = ["image_preprocessing.py"], +) + +py_test( + name = "image_preprocessing_test", + srcs = ["image_preprocessing_test.py"], + deps = [":image_preprocessing"], +) diff --git a/mediapipe/model_maker/python/vision/core/__init__.py b/mediapipe/model_maker/python/vision/core/__init__.py new file mode 100644 index 00000000..7ca2f921 --- /dev/null +++ b/mediapipe/model_maker/python/vision/core/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/mediapipe/model_maker/python/core/utils/image_preprocessing.py b/mediapipe/model_maker/python/vision/core/image_preprocessing.py similarity index 98% rename from mediapipe/model_maker/python/core/utils/image_preprocessing.py rename to mediapipe/model_maker/python/vision/core/image_preprocessing.py index 62b34fb2..104ccd9c 100644 --- a/mediapipe/model_maker/python/core/utils/image_preprocessing.py +++ b/mediapipe/model_maker/python/vision/core/image_preprocessing.py @@ -13,11 +13,7 @@ # limitations under the License. # ============================================================================== """ImageNet preprocessing.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function -# Dependency imports import tensorflow as tf IMAGE_SIZE = 224 diff --git a/mediapipe/model_maker/python/core/utils/image_preprocessing_test.py b/mediapipe/model_maker/python/vision/core/image_preprocessing_test.py similarity index 94% rename from mediapipe/model_maker/python/core/utils/image_preprocessing_test.py rename to mediapipe/model_maker/python/vision/core/image_preprocessing_test.py index bc4b4456..0594b437 100644 --- a/mediapipe/model_maker/python/core/utils/image_preprocessing_test.py +++ b/mediapipe/model_maker/python/vision/core/image_preprocessing_test.py @@ -12,15 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -# Dependency imports import numpy as np import tensorflow as tf -from mediapipe.model_maker.python.core.utils import image_preprocessing +from mediapipe.model_maker.python.vision.core import image_preprocessing def _get_preprocessed_image(preprocessor, is_training=False): diff --git a/mediapipe/model_maker/python/vision/image_classifier/BUILD b/mediapipe/model_maker/python/vision/image_classifier/BUILD index a9386d56..a2268059 100644 --- a/mediapipe/model_maker/python/vision/image_classifier/BUILD +++ b/mediapipe/model_maker/python/vision/image_classifier/BUILD @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Placeholder for internal Python library rule. # Placeholder for internal Python strict library and test compatibility macro. +# Placeholder for internal Python library rule. licenses(["notice"]) @@ -78,9 +78,9 @@ py_library( ":train_image_classifier_lib", "//mediapipe/model_maker/python/core/data:classification_dataset", "//mediapipe/model_maker/python/core/tasks:classifier", - "//mediapipe/model_maker/python/core/utils:image_preprocessing", "//mediapipe/model_maker/python/core/utils:model_util", "//mediapipe/model_maker/python/core/utils:quantization", + "//mediapipe/model_maker/python/vision/core:image_preprocessing", ], ) diff --git a/mediapipe/model_maker/python/vision/image_classifier/dataset.py b/mediapipe/model_maker/python/vision/image_classifier/dataset.py index 4d808500..071fe483 100644 --- a/mediapipe/model_maker/python/vision/image_classifier/dataset.py +++ b/mediapipe/model_maker/python/vision/image_classifier/dataset.py @@ -16,7 +16,7 @@ import os import random -from typing import List, Optional, Tuple +from typing import List, Optional import tensorflow as tf import tensorflow_datasets as tfds @@ -84,10 +84,10 @@ class Dataset(classification_dataset.ClassificationDataset): name for name in os.listdir(data_root) if os.path.isdir(os.path.join(data_root, name))) all_label_size = len(label_names) - label_to_index = dict( + index_by_label = dict( (name, index) for index, name in enumerate(label_names)) all_image_labels = [ - label_to_index[os.path.basename(os.path.dirname(path))] + index_by_label[os.path.basename(os.path.dirname(path))] for path in all_image_paths ] @@ -106,33 +106,4 @@ class Dataset(classification_dataset.ClassificationDataset): 'Load image with size: %d, num_label: %d, labels: %s.', all_image_size, all_label_size, ', '.join(label_names)) return Dataset( - dataset=image_label_ds, size=all_image_size, index_to_label=label_names) - - @classmethod - def load_tf_dataset( - cls, name: str - ) -> Tuple[Optional[classification_dataset.ClassificationDataset], - Optional[classification_dataset.ClassificationDataset], - Optional[classification_dataset.ClassificationDataset]]: - """Loads data from tensorflow_datasets. - - Args: - name: the registered name of the tfds.core.DatasetBuilder. Refer to the - documentation of tfds.load for more details. - - Returns: - A tuple of Datasets for the train/validation/test. - - Raises: - ValueError: if the input tf dataset does not have train/validation/test - labels. - """ - data, info = tfds.load(name, with_info=True) - if 'label' not in info.features: - raise ValueError('info.features need to contain \'label\' key.') - label_names = info.features['label'].names - - train_data = _create_data('train', data, info, label_names) - validation_data = _create_data('validation', data, info, label_names) - test_data = _create_data('test', data, info, label_names) - return train_data, validation_data, test_data + dataset=image_label_ds, size=all_image_size, label_names=label_names) diff --git a/mediapipe/model_maker/python/vision/image_classifier/dataset_test.py b/mediapipe/model_maker/python/vision/image_classifier/dataset_test.py index 3a5d198b..0eed547e 100644 --- a/mediapipe/model_maker/python/vision/image_classifier/dataset_test.py +++ b/mediapipe/model_maker/python/vision/image_classifier/dataset_test.py @@ -49,27 +49,27 @@ class DatasetTest(tf.test.TestCase): def test_split(self): ds = tf.data.Dataset.from_tensor_slices([[0, 1], [1, 1], [0, 0], [1, 0]]) - data = dataset.Dataset(ds, 4, ['pos', 'neg']) - train_data, test_data = data.split(0.5) + data = dataset.Dataset(dataset=ds, size=4, label_names=['pos', 'neg']) + train_data, test_data = data.split(fraction=0.5) self.assertLen(train_data, 2) for i, elem in enumerate(train_data._dataset): self.assertTrue((elem.numpy() == np.array([i, 1])).all()) self.assertEqual(train_data.num_classes, 2) - self.assertEqual(train_data.index_to_label, ['pos', 'neg']) + self.assertEqual(train_data.label_names, ['pos', 'neg']) self.assertLen(test_data, 2) for i, elem in enumerate(test_data._dataset): self.assertTrue((elem.numpy() == np.array([i, 0])).all()) self.assertEqual(test_data.num_classes, 2) - self.assertEqual(test_data.index_to_label, ['pos', 'neg']) + self.assertEqual(test_data.label_names, ['pos', 'neg']) def test_from_folder(self): - data = dataset.Dataset.from_folder(self.image_path) + data = dataset.Dataset.from_folder(dirname=self.image_path) self.assertLen(data, 2) self.assertEqual(data.num_classes, 2) - self.assertEqual(data.index_to_label, ['daisy', 'tulips']) + self.assertEqual(data.label_names, ['daisy', 'tulips']) for image, label in data.gen_tf_dataset(): self.assertTrue(label.numpy() == 1 or label.numpy() == 0) if label.numpy() == 0: @@ -88,19 +88,19 @@ class DatasetTest(tf.test.TestCase): self.assertIsInstance(train_data.gen_tf_dataset(), tf.data.Dataset) self.assertLen(train_data, 1034) self.assertEqual(train_data.num_classes, 3) - self.assertEqual(train_data.index_to_label, + self.assertEqual(train_data.label_names, ['angular_leaf_spot', 'bean_rust', 'healthy']) self.assertIsInstance(validation_data.gen_tf_dataset(), tf.data.Dataset) self.assertLen(validation_data, 133) self.assertEqual(validation_data.num_classes, 3) - self.assertEqual(validation_data.index_to_label, + self.assertEqual(validation_data.label_names, ['angular_leaf_spot', 'bean_rust', 'healthy']) self.assertIsInstance(test_data.gen_tf_dataset(), tf.data.Dataset) self.assertLen(test_data, 128) self.assertEqual(test_data.num_classes, 3) - self.assertEqual(test_data.index_to_label, + self.assertEqual(test_data.label_names, ['angular_leaf_spot', 'bean_rust', 'healthy']) diff --git a/mediapipe/model_maker/python/vision/image_classifier/image_classifier.py b/mediapipe/model_maker/python/vision/image_classifier/image_classifier.py index 7a99f9ae..a3282ebf 100644 --- a/mediapipe/model_maker/python/vision/image_classifier/image_classifier.py +++ b/mediapipe/model_maker/python/vision/image_classifier/image_classifier.py @@ -13,16 +13,16 @@ # limitations under the License. """APIs to train image classifier model.""" -from typing import Any, List, Optional +from typing import List, Optional import tensorflow as tf import tensorflow_hub as hub from mediapipe.model_maker.python.core.data import classification_dataset as classification_ds from mediapipe.model_maker.python.core.tasks import classifier -from mediapipe.model_maker.python.core.utils import image_preprocessing from mediapipe.model_maker.python.core.utils import model_util from mediapipe.model_maker.python.core.utils import quantization +from mediapipe.model_maker.python.vision.core import image_preprocessing from mediapipe.model_maker.python.vision.image_classifier import hyperparameters as hp from mediapipe.model_maker.python.vision.image_classifier import model_spec as ms from mediapipe.model_maker.python.vision.image_classifier import train_image_classifier_lib @@ -31,18 +31,18 @@ from mediapipe.model_maker.python.vision.image_classifier import train_image_cla class ImageClassifier(classifier.Classifier): """ImageClassifier for building image classification model.""" - def __init__(self, model_spec: ms.ModelSpec, index_to_label: List[Any], + def __init__(self, model_spec: ms.ModelSpec, label_names: List[str], hparams: hp.HParams): """Initializes ImageClassifier class. Args: model_spec: Specification for the model. - index_to_label: A list that maps from index to label class name. + label_names: A list of label names for the classes. hparams: The hyperparameters for training image classifier. """ - super(ImageClassifier, self).__init__( + super().__init__( model_spec=model_spec, - index_to_label=index_to_label, + label_names=label_names, shuffle=hparams.shuffle, full_train=hparams.do_fine_tuning) self._hparams = hparams @@ -80,9 +80,7 @@ class ImageClassifier(classifier.Classifier): spec = ms.SupportedModels.get(model_spec) image_classifier = cls( - model_spec=spec, - index_to_label=train_data.index_to_label, - hparams=hparams) + model_spec=spec, label_names=train_data.label_names, hparams=hparams) image_classifier._create_model() diff --git a/mediapipe/model_maker/python/vision/image_classifier/train_image_classifier_lib.py b/mediapipe/model_maker/python/vision/image_classifier/train_image_classifier_lib.py index 704d71a5..265c36a6 100644 --- a/mediapipe/model_maker/python/vision/image_classifier/train_image_classifier_lib.py +++ b/mediapipe/model_maker/python/vision/image_classifier/train_image_classifier_lib.py @@ -98,6 +98,5 @@ def train_model(model: tf.keras.Model, hparams: hp.HParams, return model.fit( x=train_ds, epochs=hparams.train_epochs, - steps_per_epoch=hparams.steps_per_epoch, validation_data=validation_ds, callbacks=callbacks) diff --git a/mediapipe/modules/face_geometry/libs/effect_renderer.cc b/mediapipe/modules/face_geometry/libs/effect_renderer.cc index 27a54e01..73f47308 100644 --- a/mediapipe/modules/face_geometry/libs/effect_renderer.cc +++ b/mediapipe/modules/face_geometry/libs/effect_renderer.cc @@ -161,7 +161,7 @@ class Texture { ~Texture() { if (is_owned_) { - glDeleteProgram(handle_); + glDeleteTextures(1, &handle_); } } diff --git a/mediapipe/python/BUILD b/mediapipe/python/BUILD index 3df0e279..dc205855 100644 --- a/mediapipe/python/BUILD +++ b/mediapipe/python/BUILD @@ -87,6 +87,9 @@ cc_library( cc_library( name = "builtin_task_graphs", deps = [ + "//mediapipe/tasks/cc/vision/gesture_recognizer:gesture_recognizer_graph", + "//mediapipe/tasks/cc/vision/image_classifier:image_classifier_graph", + "//mediapipe/tasks/cc/vision/image_segmenter:image_segmenter_graph", "//mediapipe/tasks/cc/vision/object_detector:object_detector_graph", "//mediapipe/tasks/cc/vision/image_embedder:image_embedder_graph", ], diff --git a/mediapipe/python/packet_getter.py b/mediapipe/python/packet_getter.py index 4d93e713..cf6e7574 100644 --- a/mediapipe/python/packet_getter.py +++ b/mediapipe/python/packet_getter.py @@ -14,7 +14,7 @@ """The public facing packet getter APIs.""" -from typing import List, Type +from typing import List from google.protobuf import message from google.protobuf import symbol_database @@ -39,7 +39,7 @@ get_image_frame = _packet_getter.get_image_frame get_matrix = _packet_getter.get_matrix -def get_proto(packet: mp_packet.Packet) -> Type[message.Message]: +def get_proto(packet: mp_packet.Packet) -> message.Message: """Get the content of a MediaPipe proto Packet as a proto message. Args: diff --git a/mediapipe/tasks/cc/components/BUILD b/mediapipe/tasks/cc/components/BUILD index e4905546..344fafb4 100644 --- a/mediapipe/tasks/cc/components/BUILD +++ b/mediapipe/tasks/cc/components/BUILD @@ -46,8 +46,10 @@ cc_library( "//mediapipe/framework/formats:image", "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/framework/formats:tensor", + "//mediapipe/gpu:gpu_origin_cc_proto", "//mediapipe/tasks/cc:common", "//mediapipe/tasks/cc/core:model_resources", + "//mediapipe/tasks/cc/core/proto:acceleration_cc_proto", "//mediapipe/tasks/cc/vision/utils:image_tensor_specs", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", diff --git a/mediapipe/tasks/cc/components/calculators/BUILD b/mediapipe/tasks/cc/components/calculators/BUILD index 7d01e4df..a688f291 100644 --- a/mediapipe/tasks/cc/components/calculators/BUILD +++ b/mediapipe/tasks/cc/components/calculators/BUILD @@ -44,6 +44,30 @@ cc_library( alwayslink = 1, ) +cc_test( + name = "classification_aggregation_calculator_test", + srcs = ["classification_aggregation_calculator_test.cc"], + deps = [ + ":classification_aggregation_calculator", + ":classification_aggregation_calculator_cc_proto", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework:output_stream_poller", + "//mediapipe/framework:packet", + "//mediapipe/framework:timestamp", + "//mediapipe/framework/api2:builder", + "//mediapipe/framework/api2:port", + "//mediapipe/framework/formats:classification_cc_proto", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:parse_text_proto", + "//mediapipe/framework/port:status", + "//mediapipe/tasks/cc/components/containers/proto:classifications_cc_proto", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:str_format", + "@org_tensorflow//tensorflow/lite/core/shims:cc_shims_test_util", + ], +) + mediapipe_proto_library( name = "score_calibration_calculator_proto", srcs = ["score_calibration_calculator.proto"], diff --git a/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.cc b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.cc index e1f69e60..1a83fdad 100644 --- a/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.cc +++ b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.cc @@ -31,37 +31,62 @@ namespace mediapipe { namespace api2 { -using ::mediapipe::tasks::ClassificationAggregationCalculatorOptions; using ::mediapipe::tasks::components::containers::proto::ClassificationResult; using ::mediapipe::tasks::components::containers::proto::Classifications; -// Aggregates ClassificationLists into a single ClassificationResult that has -// 3 dimensions: (classification head, classification timestamp, classification -// category). +// Aggregates ClassificationLists into either a ClassificationResult object +// representing the classification results aggregated by classifier head, or +// into an std::vector representing the classification +// results aggregated first by timestamp then by classifier head. // // Inputs: -// CLASSIFICATIONS - ClassificationList +// CLASSIFICATIONS - ClassificationList @Multiple // ClassificationList per classification head. // TIMESTAMPS - std::vector @Optional -// The collection of the timestamps that a single ClassificationResult -// should aggragate. This stream is optional, and the timestamp information -// will only be populated to the ClassificationResult proto when this stream -// is connected. +// The collection of the timestamps that this calculator should aggregate. +// This stream is optional: if provided then the TIMESTAMPED_CLASSIFICATIONS +// output is used for results. Otherwise as no timestamp aggregation is +// required the CLASSIFICATIONS output is used for results. // // Outputs: -// CLASSIFICATION_RESULT - ClassificationResult +// CLASSIFICATIONS - ClassificationResult @Optional +// The classification results aggregated by head. Must be connected if the +// TIMESTAMPS input is not connected, as it signals that timestamp +// aggregation is not required. +// TIMESTAMPED_CLASSIFICATIONS - std::vector @Optional +// The classification result aggregated by timestamp, then by head. Must be +// connected if the TIMESTAMPS input is connected, as it signals that +// timestamp aggregation is required. +// // TODO: remove output once migration is over. +// CLASSIFICATION_RESULT - (DEPRECATED) ClassificationResult @Optional // The aggregated classification result. // -// Example: +// Example without timestamp aggregation: +// node { +// calculator: "ClassificationAggregationCalculator" +// input_stream: "CLASSIFICATIONS:0:stream_a" +// input_stream: "CLASSIFICATIONS:1:stream_b" +// input_stream: "CLASSIFICATIONS:2:stream_c" +// output_stream: "CLASSIFICATIONS:classifications" +// options { +// [mediapipe.ClassificationAggregationCalculatorOptions.ext] { +// head_names: "head_name_a" +// head_names: "head_name_b" +// head_names: "head_name_c" +// } +// } +// } +// +// Example with timestamp aggregation: // node { // calculator: "ClassificationAggregationCalculator" // input_stream: "CLASSIFICATIONS:0:stream_a" // input_stream: "CLASSIFICATIONS:1:stream_b" // input_stream: "CLASSIFICATIONS:2:stream_c" // input_stream: "TIMESTAMPS:timestamps" -// output_stream: "CLASSIFICATION_RESULT:classification_result" +// output_stream: "TIMESTAMPED_CLASSIFICATIONS:timestamped_classifications" // options { -// [mediapipe.tasks.ClassificationAggregationCalculatorOptions.ext] { +// [mediapipe.ClassificationAggregationCalculatorOptions.ext] { // head_names: "head_name_a" // head_names: "head_name_b" // head_names: "head_name_c" @@ -74,8 +99,15 @@ class ClassificationAggregationCalculator : public Node { "CLASSIFICATIONS"}; static constexpr Input>::Optional kTimestampsIn{ "TIMESTAMPS"}; - static constexpr Output kOut{"CLASSIFICATION_RESULT"}; - MEDIAPIPE_NODE_CONTRACT(kClassificationListIn, kTimestampsIn, kOut); + static constexpr Output::Optional kClassificationsOut{ + "CLASSIFICATIONS"}; + static constexpr Output>::Optional + kTimestampedClassificationsOut{"TIMESTAMPED_CLASSIFICATIONS"}; + static constexpr Output::Optional + kClassificationResultOut{"CLASSIFICATION_RESULT"}; + MEDIAPIPE_NODE_CONTRACT(kClassificationListIn, kTimestampsIn, + kClassificationsOut, kTimestampedClassificationsOut, + kClassificationResultOut); static absl::Status UpdateContract(CalculatorContract* cc); absl::Status Open(CalculatorContext* cc); @@ -88,6 +120,11 @@ class ClassificationAggregationCalculator : public Node { cached_classifications_; ClassificationResult ConvertToClassificationResult(CalculatorContext* cc); + std::vector ConvertToTimestampedClassificationResults( + CalculatorContext* cc); + // TODO: deprecate this function once migration is over. + ClassificationResult LegacyConvertToClassificationResult( + CalculatorContext* cc); }; absl::Status ClassificationAggregationCalculator::UpdateContract( @@ -100,6 +137,10 @@ absl::Status ClassificationAggregationCalculator::UpdateContract( << "The size of classifications input streams should match the " "size of head names specified in the calculator options"; } + // TODO: enforce connecting TIMESTAMPED_CLASSIFICATIONS if + // TIMESTAMPS is connected, and connecting CLASSIFICATIONS if TIMESTAMPS is + // not connected. All dependent tasks must be updated to use these outputs + // first. return absl::OkStatus(); } @@ -124,10 +165,19 @@ absl::Status ClassificationAggregationCalculator::Process( [](const auto& elem) -> ClassificationList { return elem.Get(); }); cached_classifications_[cc->InputTimestamp().Value()] = std::move(classification_lists); - if (time_aggregation_enabled_ && kTimestampsIn(cc).IsEmpty()) { - return absl::OkStatus(); + ClassificationResult classification_result; + if (time_aggregation_enabled_) { + if (kTimestampsIn(cc).IsEmpty()) { + return absl::OkStatus(); + } + classification_result = LegacyConvertToClassificationResult(cc); + kTimestampedClassificationsOut(cc).Send( + ConvertToTimestampedClassificationResults(cc)); + } else { + classification_result = LegacyConvertToClassificationResult(cc); + kClassificationsOut(cc).Send(ConvertToClassificationResult(cc)); } - kOut(cc).Send(ConvertToClassificationResult(cc)); + kClassificationResultOut(cc).Send(classification_result); RET_CHECK(cached_classifications_.empty()); return absl::OkStatus(); } @@ -136,6 +186,50 @@ ClassificationResult ClassificationAggregationCalculator::ConvertToClassificationResult( CalculatorContext* cc) { ClassificationResult result; + auto& classification_lists = + cached_classifications_[cc->InputTimestamp().Value()]; + for (int i = 0; i < classification_lists.size(); ++i) { + auto classifications = result.add_classifications(); + classifications->set_head_index(i); + if (!head_names_.empty()) { + classifications->set_head_name(head_names_[i]); + } + *classifications->mutable_classification_list() = + std::move(classification_lists[i]); + } + cached_classifications_.erase(cc->InputTimestamp().Value()); + return result; +} + +std::vector +ClassificationAggregationCalculator::ConvertToTimestampedClassificationResults( + CalculatorContext* cc) { + auto timestamps = kTimestampsIn(cc).Get(); + std::vector results; + results.reserve(timestamps.size()); + for (const auto& timestamp : timestamps) { + ClassificationResult result; + result.set_timestamp_ms((timestamp.Value() - timestamps[0].Value()) / 1000); + auto& classification_lists = cached_classifications_[timestamp.Value()]; + for (int i = 0; i < classification_lists.size(); ++i) { + auto classifications = result.add_classifications(); + classifications->set_head_index(i); + if (!head_names_.empty()) { + classifications->set_head_name(head_names_[i]); + } + *classifications->mutable_classification_list() = + std::move(classification_lists[i]); + } + cached_classifications_.erase(timestamp.Value()); + results.push_back(std::move(result)); + } + return results; +} + +ClassificationResult +ClassificationAggregationCalculator::LegacyConvertToClassificationResult( + CalculatorContext* cc) { + ClassificationResult result; Timestamp first_timestamp(0); std::vector timestamps; if (time_aggregation_enabled_) { @@ -177,7 +271,6 @@ ClassificationAggregationCalculator::ConvertToClassificationResult( entry->set_timestamp_ms((timestamp.Value() - first_timestamp.Value()) / 1000); } - cached_classifications_.erase(timestamp.Value()); } return result; } diff --git a/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.proto b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.proto index c2a74a48..e2ed1788 100644 --- a/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.proto +++ b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.proto @@ -15,7 +15,7 @@ limitations under the License. syntax = "proto2"; -package mediapipe.tasks; +package mediapipe; import "mediapipe/framework/calculator.proto"; diff --git a/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator_test.cc b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator_test.cc new file mode 100644 index 00000000..1bc8cafd --- /dev/null +++ b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator_test.cc @@ -0,0 +1,213 @@ +/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_format.h" +#include "mediapipe/framework/api2/builder.h" +#include "mediapipe/framework/api2/port.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/classification.pb.h" +#include "mediapipe/framework/output_stream_poller.h" +#include "mediapipe/framework/packet.h" +#include "mediapipe/framework/port/gmock.h" +#include "mediapipe/framework/port/gtest.h" +#include "mediapipe/framework/port/parse_text_proto.h" +#include "mediapipe/framework/port/status_macros.h" +#include "mediapipe/framework/port/status_matchers.h" +#include "mediapipe/framework/timestamp.h" +#include "mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.pb.h" +#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h" +#include "tensorflow/lite/core/shims/cc/shims_test_util.h" + +namespace mediapipe { +namespace { + +using ::mediapipe::ParseTextProtoOrDie; +using ::mediapipe::api2::Input; +using ::mediapipe::api2::Output; +using ::mediapipe::api2::builder::Graph; +using ::mediapipe::api2::builder::Source; +using ::mediapipe::tasks::components::containers::proto::ClassificationResult; +using ::testing::Pointwise; + +constexpr char kClassificationInput0Tag[] = "CLASSIFICATIONS_0"; +constexpr char kClassificationInput0Name[] = "classifications_0"; +constexpr char kClassificationInput1Tag[] = "CLASSIFICATIONS_1"; +constexpr char kClassificationInput1Name[] = "classifications_1"; +constexpr char kTimestampsTag[] = "TIMESTAMPS"; +constexpr char kTimestampsName[] = "timestamps"; +constexpr char kClassificationsTag[] = "CLASSIFICATIONS"; +constexpr char kClassificationsName[] = "classifications"; +constexpr char kTimestampedClassificationsTag[] = "TIMESTAMPED_CLASSIFICATIONS"; +constexpr char kTimestampedClassificationsName[] = + "timestamped_classifications"; + +ClassificationList MakeClassificationList(int class_index) { + return ParseTextProtoOrDie(absl::StrFormat( + R"pb( + classification { index: %d } + )pb", + class_index)); +} + +class ClassificationAggregationCalculatorTest + : public tflite_shims::testing::Test { + protected: + absl::StatusOr BuildGraph( + bool connect_timestamps = false) { + Graph graph; + auto& calculator = graph.AddNode("ClassificationAggregationCalculator"); + calculator + .GetOptions() = + ParseTextProtoOrDie< + mediapipe::ClassificationAggregationCalculatorOptions>( + R"pb(head_names: "foo" head_names: "bar")pb"); + graph[Input(kClassificationInput0Tag)].SetName( + kClassificationInput0Name) >> + calculator.In(absl::StrFormat("%s:%d", kClassificationsTag, 0)); + graph[Input(kClassificationInput1Tag)].SetName( + kClassificationInput1Name) >> + calculator.In(absl::StrFormat("%s:%d", kClassificationsTag, 1)); + if (connect_timestamps) { + graph[Input>(kTimestampsTag)].SetName( + kTimestampsName) >> + calculator.In(kTimestampsTag); + calculator.Out(kTimestampedClassificationsTag) + .SetName(kTimestampedClassificationsName) >> + graph[Output>( + kTimestampedClassificationsTag)]; + } else { + calculator.Out(kClassificationsTag).SetName(kClassificationsName) >> + graph[Output(kClassificationsTag)]; + } + + MP_RETURN_IF_ERROR(calculator_graph_.Initialize(graph.GetConfig())); + if (connect_timestamps) { + ASSIGN_OR_RETURN(auto poller, calculator_graph_.AddOutputStreamPoller( + kTimestampedClassificationsName)); + MP_RETURN_IF_ERROR(calculator_graph_.StartRun(/*extra_side_packets=*/{})); + return poller; + } + ASSIGN_OR_RETURN(auto poller, calculator_graph_.AddOutputStreamPoller( + kClassificationsName)); + MP_RETURN_IF_ERROR(calculator_graph_.StartRun(/*extra_side_packets=*/{})); + return poller; + } + + absl::Status Send( + std::vector classifications, int timestamp = 0, + std::optional> aggregation_timestamps = std::nullopt) { + MP_RETURN_IF_ERROR(calculator_graph_.AddPacketToInputStream( + kClassificationInput0Name, + MakePacket(classifications[0]) + .At(Timestamp(timestamp)))); + MP_RETURN_IF_ERROR(calculator_graph_.AddPacketToInputStream( + kClassificationInput1Name, + MakePacket(classifications[1]) + .At(Timestamp(timestamp)))); + if (aggregation_timestamps.has_value()) { + auto packet = std::make_unique>(); + for (const auto& timestamp : *aggregation_timestamps) { + packet->emplace_back(Timestamp(timestamp)); + } + MP_RETURN_IF_ERROR(calculator_graph_.AddPacketToInputStream( + kTimestampsName, Adopt(packet.release()).At(Timestamp(timestamp)))); + } + return absl::OkStatus(); + } + + template + absl::StatusOr GetResult(OutputStreamPoller& poller) { + MP_RETURN_IF_ERROR(calculator_graph_.WaitUntilIdle()); + MP_RETURN_IF_ERROR(calculator_graph_.CloseAllInputStreams()); + + Packet packet; + if (!poller.Next(&packet)) { + return absl::InternalError("Unable to get output packet"); + } + auto result = packet.Get(); + MP_RETURN_IF_ERROR(calculator_graph_.WaitUntilDone()); + return result; + } + + private: + CalculatorGraph calculator_graph_; +}; + +TEST_F(ClassificationAggregationCalculatorTest, SucceedsWithoutTimestamps) { + MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph()); + MP_ASSERT_OK(Send({MakeClassificationList(0), MakeClassificationList(1)})); + MP_ASSERT_OK_AND_ASSIGN(auto result, GetResult(poller)); + + EXPECT_THAT(result, + EqualsProto(ParseTextProtoOrDie( + R"pb(classifications { + head_index: 0 + head_name: "foo" + classification_list { classification { index: 0 } } + } + classifications { + head_index: 1 + head_name: "bar" + classification_list { classification { index: 1 } } + })pb"))); +} + +TEST_F(ClassificationAggregationCalculatorTest, SucceedsWithTimestamps) { + MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(/*connect_timestamps=*/true)); + MP_ASSERT_OK(Send({MakeClassificationList(0), MakeClassificationList(1)})); + MP_ASSERT_OK(Send( + {MakeClassificationList(2), MakeClassificationList(3)}, + /*timestamp=*/1000, + /*aggregation_timestamps=*/std::optional>({0, 1000}))); + MP_ASSERT_OK_AND_ASSIGN(auto result, + GetResult>(poller)); + + EXPECT_THAT(result, + Pointwise(EqualsProto(), + {ParseTextProtoOrDie(R"pb( + timestamp_ms: 0, + classifications { + head_index: 0 + head_name: "foo" + classification_list { classification { index: 0 } } + } + classifications { + head_index: 1 + head_name: "bar" + classification_list { classification { index: 1 } } + } + )pb"), + ParseTextProtoOrDie(R"pb( + timestamp_ms: 1, + classifications { + head_index: 0 + head_name: "foo" + classification_list { classification { index: 2 } } + } + classifications { + head_index: 1 + head_name: "bar" + classification_list { classification { index: 3 } } + } + )pb")})); +} + +} // namespace +} // namespace mediapipe diff --git a/mediapipe/tasks/cc/components/containers/BUILD b/mediapipe/tasks/cc/components/containers/BUILD index af51d0c3..7a52f11e 100644 --- a/mediapipe/tasks/cc/components/containers/BUILD +++ b/mediapipe/tasks/cc/components/containers/BUILD @@ -29,3 +29,23 @@ cc_library( "//mediapipe/framework/formats:landmark_cc_proto", ], ) + +cc_library( + name = "category", + srcs = ["category.cc"], + hdrs = ["category.h"], + deps = [ + "//mediapipe/framework/formats:classification_cc_proto", + ], +) + +cc_library( + name = "classification_result", + srcs = ["classification_result.cc"], + hdrs = ["classification_result.h"], + deps = [ + ":category", + "//mediapipe/framework/formats:classification_cc_proto", + "//mediapipe/tasks/cc/components/containers/proto:classifications_cc_proto", + ], +) diff --git a/mediapipe/tasks/cc/components/containers/category.cc b/mediapipe/tasks/cc/components/containers/category.cc new file mode 100644 index 00000000..e07333a7 --- /dev/null +++ b/mediapipe/tasks/cc/components/containers/category.cc @@ -0,0 +1,38 @@ +/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "mediapipe/tasks/cc/components/containers/category.h" + +#include +#include + +#include "mediapipe/framework/formats/classification.pb.h" + +namespace mediapipe::tasks::components::containers { + +Category ConvertToCategory(const mediapipe::Classification& proto) { + Category category; + category.index = proto.index(); + category.score = proto.score(); + if (proto.has_label()) { + category.category_name = proto.label(); + } + if (proto.has_display_name()) { + category.display_name = proto.display_name(); + } + return category; +} + +} // namespace mediapipe::tasks::components::containers diff --git a/mediapipe/tasks/cc/components/containers/category.h b/mediapipe/tasks/cc/components/containers/category.h new file mode 100644 index 00000000..57b18e7e --- /dev/null +++ b/mediapipe/tasks/cc/components/containers/category.h @@ -0,0 +1,52 @@ +/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef MEDIAPIPE_TASKS_CC_COMPONENTS_CONTAINERS_CATEGORY_H_ +#define MEDIAPIPE_TASKS_CC_COMPONENTS_CONTAINERS_CATEGORY_H_ + +#include +#include + +#include "mediapipe/framework/formats/classification.pb.h" + +namespace mediapipe::tasks::components::containers { + +// Defines a single classification result. +// +// The label maps packed into the TFLite Model Metadata [1] are used to populate +// the 'category_name' and 'display_name' fields. +// +// [1]: https://www.tensorflow.org/lite/convert/metadata +struct Category { + // The index of the category in the classification model output. + int index; + // The score for this category, e.g. (but not necessarily) a probability in + // [0,1]. + float score; + // The optional ID for the category, read from the label map packed in the + // TFLite Model Metadata if present. Not necessarily human-readable. + std::optional category_name = std::nullopt; + // The optional human-readable name for the category, read from the label map + // packed in the TFLite Model Metadata if present. + std::optional display_name = std::nullopt; +}; + +// Utility function to convert from mediapipe::Classification proto to Category +// struct. +Category ConvertToCategory(const mediapipe::Classification& proto); + +} // namespace mediapipe::tasks::components::containers + +#endif // MEDIAPIPE_TASKS_CC_COMPONENTS_CONTAINERS_CATEGORY_H_ diff --git a/mediapipe/tasks/cc/components/containers/classification_result.cc b/mediapipe/tasks/cc/components/containers/classification_result.cc new file mode 100644 index 00000000..98583ff1 --- /dev/null +++ b/mediapipe/tasks/cc/components/containers/classification_result.cc @@ -0,0 +1,57 @@ +/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "mediapipe/tasks/cc/components/containers/classification_result.h" + +#include +#include +#include + +#include "mediapipe/framework/formats/classification.pb.h" +#include "mediapipe/tasks/cc/components/containers/category.h" +#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h" + +namespace mediapipe::tasks::components::containers { + +Classifications ConvertToClassifications(const proto::Classifications& proto) { + Classifications classifications; + classifications.categories.reserve( + proto.classification_list().classification_size()); + for (const auto& classification : + proto.classification_list().classification()) { + classifications.categories.push_back(ConvertToCategory(classification)); + } + classifications.head_index = proto.head_index(); + if (proto.has_head_name()) { + classifications.head_name = proto.head_name(); + } + return classifications; +} + +ClassificationResult ConvertToClassificationResult( + const proto::ClassificationResult& proto) { + ClassificationResult classification_result; + classification_result.classifications.reserve(proto.classifications_size()); + for (const auto& classifications : proto.classifications()) { + classification_result.classifications.push_back( + ConvertToClassifications(classifications)); + } + if (proto.has_timestamp_ms()) { + classification_result.timestamp_ms = proto.timestamp_ms(); + } + return classification_result; +} + +} // namespace mediapipe::tasks::components::containers diff --git a/mediapipe/tasks/cc/components/containers/classification_result.h b/mediapipe/tasks/cc/components/containers/classification_result.h new file mode 100644 index 00000000..88273fd0 --- /dev/null +++ b/mediapipe/tasks/cc/components/containers/classification_result.h @@ -0,0 +1,68 @@ +/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef MEDIAPIPE_TASKS_CC_COMPONENTS_CONTAINERS_CLASSIFICATION_RESULT_H_ +#define MEDIAPIPE_TASKS_CC_COMPONENTS_CONTAINERS_CLASSIFICATION_RESULT_H_ + +#include +#include +#include + +#include "mediapipe/tasks/cc/components/containers/category.h" +#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h" + +namespace mediapipe::tasks::components::containers { + +// Defines classification results for a given classifier head. +struct Classifications { + // The array of predicted categories, usually sorted by descending scores, + // e.g. from high to low probability. + std::vector categories; + // The index of the classifier head (i.e. output tensor) these categories + // refer to. This is useful for multi-head models. + int head_index; + // The optional name of the classifier head, as provided in the TFLite Model + // Metadata [1] if present. This is useful for multi-head models. + // + // [1]: https://www.tensorflow.org/lite/convert/metadata + std::optional head_name = std::nullopt; +}; + +// Defines classification results of a model. +struct ClassificationResult { + // The classification results for each head of the model. + std::vector classifications; + // The optional timestamp (in milliseconds) of the start of the chunk of data + // corresponding to these results. + // + // This is only used for classification on time series (e.g. audio + // classification). In these use cases, the amount of data to process might + // exceed the maximum size that the model can process: to solve this, the + // input data is split into multiple chunks starting at different timestamps. + std::optional timestamp_ms = std::nullopt; +}; + +// Utility function to convert from Classifications proto to +// Classifications struct. +Classifications ConvertToClassifications(const proto::Classifications& proto); + +// Utility function to convert from ClassificationResult proto to +// ClassificationResult struct. +ClassificationResult ConvertToClassificationResult( + const proto::ClassificationResult& proto); + +} // namespace mediapipe::tasks::components::containers + +#endif // MEDIAPIPE_TASKS_CC_COMPONENTS_CONTAINERS_CLASSIFICATION_RESULT_H_ diff --git a/mediapipe/tasks/cc/components/containers/proto/BUILD b/mediapipe/tasks/cc/components/containers/proto/BUILD index 633b5b36..7b455c0c 100644 --- a/mediapipe/tasks/cc/components/containers/proto/BUILD +++ b/mediapipe/tasks/cc/components/containers/proto/BUILD @@ -28,6 +28,7 @@ mediapipe_proto_library( srcs = ["classifications.proto"], deps = [ ":category_proto", + "//mediapipe/framework/formats:classification_proto", ], ) diff --git a/mediapipe/tasks/cc/components/containers/proto/category.proto b/mediapipe/tasks/cc/components/containers/proto/category.proto index a154e5f4..412e7142 100644 --- a/mediapipe/tasks/cc/components/containers/proto/category.proto +++ b/mediapipe/tasks/cc/components/containers/proto/category.proto @@ -17,9 +17,10 @@ syntax = "proto2"; package mediapipe.tasks.components.containers.proto; -option java_package = "com.google.mediapipe.tasks.components.container.proto"; +option java_package = "com.google.mediapipe.tasks.components.containers.proto"; option java_outer_classname = "CategoryProto"; +// TODO: deprecate this message once migration is over. // A single classification result. message Category { // The index of the category in the corresponding label map, usually packed in diff --git a/mediapipe/tasks/cc/components/containers/proto/classifications.proto b/mediapipe/tasks/cc/components/containers/proto/classifications.proto index 0f5086b9..f098ed0e 100644 --- a/mediapipe/tasks/cc/components/containers/proto/classifications.proto +++ b/mediapipe/tasks/cc/components/containers/proto/classifications.proto @@ -17,11 +17,13 @@ syntax = "proto2"; package mediapipe.tasks.components.containers.proto; +import "mediapipe/framework/formats/classification.proto"; import "mediapipe/tasks/cc/components/containers/proto/category.proto"; -option java_package = "com.google.mediapipe.tasks.components.container.proto"; +option java_package = "com.google.mediapipe.tasks.components.containers.proto"; option java_outer_classname = "ClassificationsProto"; +// TODO: deprecate this message once migration is over. // List of predicted categories with an optional timestamp. message ClassificationEntry { // The array of predicted categories, usually sorted by descending scores, @@ -33,9 +35,12 @@ message ClassificationEntry { optional int64 timestamp_ms = 2; } -// Classifications for a given classifier head. +// Classifications for a given classifier head, i.e. for a given output tensor. message Classifications { + // TODO: deprecate this field once migration is over. repeated ClassificationEntry entries = 1; + // The classification results for this head. + optional mediapipe.ClassificationList classification_list = 4; // The index of the classifier head these categories refer to. This is useful // for multi-head models. optional int32 head_index = 2; @@ -45,7 +50,17 @@ message Classifications { optional string head_name = 3; } -// Contains one set of results per classifier head. +// Classifications for a given classifier model. message ClassificationResult { + // The classification results for each model head, i.e. one for each output + // tensor. repeated Classifications classifications = 1; + // The optional timestamp (in milliseconds) of the start of the chunk of data + // corresponding to these results. + // + // This is only used for classification on time series (e.g. audio + // classification). In these use cases, the amount of data to process might + // exceed the maximum size that the model can process: to solve this, the + // input data is split into multiple chunks starting at different timestamps. + optional int64 timestamp_ms = 2; } diff --git a/mediapipe/tasks/cc/components/containers/proto/embeddings.proto b/mediapipe/tasks/cc/components/containers/proto/embeddings.proto index d57b08b5..39811e6c 100644 --- a/mediapipe/tasks/cc/components/containers/proto/embeddings.proto +++ b/mediapipe/tasks/cc/components/containers/proto/embeddings.proto @@ -17,6 +17,9 @@ syntax = "proto2"; package mediapipe.tasks.components.containers.proto; +option java_package = "com.google.mediapipe.tasks.components.containers.proto"; +option java_outer_classname = "EmbeddingsProto"; + // Defines a dense floating-point embedding. message FloatEmbedding { repeated float values = 1 [packed = true]; diff --git a/mediapipe/tasks/cc/components/image_preprocessing.cc b/mediapipe/tasks/cc/components/image_preprocessing.cc index 046a97e4..7940080e 100644 --- a/mediapipe/tasks/cc/components/image_preprocessing.cc +++ b/mediapipe/tasks/cc/components/image_preprocessing.cc @@ -30,9 +30,11 @@ limitations under the License. #include "mediapipe/framework/formats/image.h" #include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/formats/tensor.h" +#include "mediapipe/gpu/gpu_origin.pb.h" #include "mediapipe/tasks/cc/common.h" #include "mediapipe/tasks/cc/components/image_preprocessing_options.pb.h" #include "mediapipe/tasks/cc/core/model_resources.h" +#include "mediapipe/tasks/cc/core/proto/acceleration.pb.h" #include "mediapipe/tasks/cc/vision/utils/image_tensor_specs.h" #include "tensorflow/lite/schema/schema_generated.h" @@ -128,12 +130,21 @@ absl::Status ConfigureImageToTensorCalculator( options->mutable_output_tensor_float_range()->set_max((255.0f - mean) / std); } + // TODO: need to support different GPU origin on differnt + // platforms or applications. + options->set_gpu_origin(mediapipe::GpuOrigin::TOP_LEFT); return absl::OkStatus(); } } // namespace +bool DetermineImagePreprocessingGpuBackend( + const core::proto::Acceleration& acceleration) { + return acceleration.has_gpu(); +} + absl::Status ConfigureImagePreprocessing(const ModelResources& model_resources, + bool use_gpu, ImagePreprocessingOptions* options) { ASSIGN_OR_RETURN(auto image_tensor_specs, BuildImageTensorSpecs(model_resources)); @@ -141,7 +152,9 @@ absl::Status ConfigureImagePreprocessing(const ModelResources& model_resources, image_tensor_specs, options->mutable_image_to_tensor_options())); // The GPU backend isn't able to process int data. If the input tensor is // quantized, forces the image preprocessing graph to use CPU backend. - if (image_tensor_specs.tensor_type == tflite::TensorType_UINT8) { + if (use_gpu && image_tensor_specs.tensor_type != tflite::TensorType_UINT8) { + options->set_backend(ImagePreprocessingOptions::GPU_BACKEND); + } else { options->set_backend(ImagePreprocessingOptions::CPU_BACKEND); } return absl::OkStatus(); diff --git a/mediapipe/tasks/cc/components/image_preprocessing.h b/mediapipe/tasks/cc/components/image_preprocessing.h index a5b767f3..6963b655 100644 --- a/mediapipe/tasks/cc/components/image_preprocessing.h +++ b/mediapipe/tasks/cc/components/image_preprocessing.h @@ -19,20 +19,26 @@ limitations under the License. #include "absl/status/status.h" #include "mediapipe/tasks/cc/components/image_preprocessing_options.pb.h" #include "mediapipe/tasks/cc/core/model_resources.h" +#include "mediapipe/tasks/cc/core/proto/acceleration.pb.h" namespace mediapipe { namespace tasks { namespace components { -// Configures an ImagePreprocessing subgraph using the provided model resources. +// Configures an ImagePreprocessing subgraph using the provided model resources +// When use_gpu is true, use GPU as backend to convert image to tensor. // - Accepts CPU input images and outputs CPU tensors. // // Example usage: // // auto& preprocessing = // graph.AddNode("mediapipe.tasks.components.ImagePreprocessingSubgraph"); +// core::proto::Acceleration acceleration; +// acceleration.mutable_xnnpack(); +// bool use_gpu = DetermineImagePreprocessingGpuBackend(acceleration); // MP_RETURN_IF_ERROR(ConfigureImagePreprocessing( // model_resources, +// use_gpu, // &preprocessing.GetOptions())); // // The resulting ImagePreprocessing subgraph has the following I/O: @@ -56,9 +62,14 @@ namespace components { // The image that has the pixel data stored on the target storage (CPU vs // GPU). absl::Status ConfigureImagePreprocessing( - const core::ModelResources& model_resources, + const core::ModelResources& model_resources, bool use_gpu, ImagePreprocessingOptions* options); +// Determine if the image preprocessing subgraph should use GPU as the backend +// according to the given acceleration setting. +bool DetermineImagePreprocessingGpuBackend( + const core::proto::Acceleration& acceleration); + } // namespace components } // namespace tasks } // namespace mediapipe diff --git a/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph.cc b/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph.cc index b4fbf966..0fb62afa 100644 --- a/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph.cc +++ b/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph.cc @@ -78,6 +78,14 @@ constexpr char kClassificationsTag[] = "CLASSIFICATIONS"; constexpr char kScoresTag[] = "SCORES"; constexpr char kTensorsTag[] = "TENSORS"; constexpr char kTimestampsTag[] = "TIMESTAMPS"; +constexpr char kTimestampedClassificationsTag[] = "TIMESTAMPED_CLASSIFICATIONS"; + +// Struct holding the different output streams produced by the graph. +struct ClassificationPostprocessingOutputStreams { + Source classification_result; + Source classifications; + Source> timestamped_classifications; +}; // Performs sanity checks on provided ClassifierOptions. absl::Status SanityCheckClassifierOptions( @@ -286,7 +294,7 @@ absl::Status ConfigureScoreCalibrationIfAny( void ConfigureClassificationAggregationCalculator( const ModelMetadataExtractor& metadata_extractor, - ClassificationAggregationCalculatorOptions* options) { + mediapipe::ClassificationAggregationCalculatorOptions* options) { auto* output_tensors_metadata = metadata_extractor.GetOutputTensorMetadata(); if (output_tensors_metadata == nullptr) { return; @@ -378,12 +386,23 @@ absl::Status ConfigureClassificationPostprocessingGraph( // TENSORS - std::vector // The output tensors of an InferenceCalculator. // TIMESTAMPS - std::vector @Optional -// The collection of timestamps that a single ClassificationResult should -// aggregate. This is mostly useful for classifiers working on time series, -// e.g. audio or video classification. +// The collection of the timestamps that this calculator should aggregate. +// This stream is optional: if provided then the TIMESTAMPED_CLASSIFICATIONS +// output is used for results. Otherwise as no timestamp aggregation is +// required the CLASSIFICATIONS output is used for results. +// // Outputs: -// CLASSIFICATION_RESULT - ClassificationResult -// The output aggregated classification results. +// CLASSIFICATIONS - ClassificationResult @Optional +// The classification results aggregated by head. Must be connected if the +// TIMESTAMPS input is not connected, as it signals that timestamp +// aggregation is not required. +// TIMESTAMPED_CLASSIFICATIONS - std::vector @Optional +// The classification result aggregated by timestamp, then by head. Must be +// connected if the TIMESTAMPS input is connected, as it signals that +// timestamp aggregation is required. +// // TODO: remove output once migration is over. +// CLASSIFICATION_RESULT - (DEPRECATED) ClassificationResult @Optional +// The aggregated classification result. // // The recommended way of using this graph is through the GraphBuilder API // using the 'ConfigureClassificationPostprocessingGraph()' function. See header @@ -394,28 +413,39 @@ class ClassificationPostprocessingGraph : public mediapipe::Subgraph { mediapipe::SubgraphContext* sc) override { Graph graph; ASSIGN_OR_RETURN( - auto classification_result_out, + auto output_streams, BuildClassificationPostprocessing( sc->Options(), graph[Input>(kTensorsTag)], graph[Input>(kTimestampsTag)], graph)); - classification_result_out >> + output_streams.classification_result >> graph[Output(kClassificationResultTag)]; + output_streams.classifications >> + graph[Output(kClassificationsTag)]; + output_streams.timestamped_classifications >> + graph[Output>( + kTimestampedClassificationsTag)]; return graph.GetConfig(); } private: // Adds an on-device classification postprocessing graph into the provided // builder::Graph instance. The classification postprocessing graph takes - // tensors (std::vector) as input and returns one output - // stream containing the output classification results (ClassificationResult). + // tensors (std::vector) and optional timestamps + // (std::vector) as input and returns two output streams: + // - classification results aggregated by classifier head as a + // ClassificationResult proto, used when no timestamps are passed in + // the graph, + // - classification results aggregated by timestamp then by classifier head + // as a std::vector, used when timestamps are passed + // in the graph. // // options: the on-device ClassificationPostprocessingGraphOptions. // tensors_in: (std::vector>) tensors to postprocess. // timestamps_in: (std::vector) optional collection of - // timestamps that a single ClassificationResult should aggregate. + // timestamps that should be used to aggregate classification results. // graph: the mediapipe builder::Graph instance to be updated. - absl::StatusOr> + absl::StatusOr BuildClassificationPostprocessing( const proto::ClassificationPostprocessingGraphOptions& options, Source> tensors_in, @@ -494,7 +524,8 @@ class ClassificationPostprocessingGraph : public mediapipe::Subgraph { // Aggregates Classifications into a single ClassificationResult. auto& result_aggregation = graph.AddNode("ClassificationAggregationCalculator"); - result_aggregation.GetOptions() + result_aggregation + .GetOptions() .CopyFrom(options.classification_aggregation_options()); for (int i = 0; i < num_heads; ++i) { tensors_to_classification_nodes[i]->Out(kClassificationsTag) >> @@ -504,8 +535,15 @@ class ClassificationPostprocessingGraph : public mediapipe::Subgraph { timestamps_in >> result_aggregation.In(kTimestampsTag); // Connects output. - return result_aggregation[Output( - kClassificationResultTag)]; + ClassificationPostprocessingOutputStreams output_streams{ + /*classification_result=*/result_aggregation + [Output(kClassificationResultTag)], + /*classifications=*/ + result_aggregation[Output(kClassificationsTag)], + /*timestamped_classifications=*/ + result_aggregation[Output>( + kTimestampedClassificationsTag)]}; + return output_streams; } }; diff --git a/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph.h b/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph.h index be166982..48575ceb 100644 --- a/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph.h +++ b/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph.h @@ -45,12 +45,22 @@ namespace processors { // TENSORS - std::vector // The output tensors of an InferenceCalculator. // TIMESTAMPS - std::vector @Optional -// The collection of timestamps that a single ClassificationResult should -// aggregate. This is mostly useful for classifiers working on time series, -// e.g. audio or video classification. +// The collection of the timestamps that this calculator should aggregate. +// This stream is optional: if provided then the TIMESTAMPED_CLASSIFICATIONS +// output is used for results. Otherwise as no timestamp aggregation is +// required the CLASSIFICATIONS output is used for results. // Outputs: -// CLASSIFICATION_RESULT - ClassificationResult -// The output aggregated classification results. +// CLASSIFICATIONS - ClassificationResult @Optional +// The classification results aggregated by head. Must be connected if the +// TIMESTAMPS input is not connected, as it signals that timestamp +// aggregation is not required. +// TIMESTAMPED_CLASSIFICATIONS - std::vector @Optional +// The classification result aggregated by timestamp, then by head. Must be +// connected if the TIMESTAMPS input is connected, as it signals that +// timestamp aggregation is required. +// // TODO: remove output once migration is over. +// CLASSIFICATION_RESULT - (DEPRECATED) ClassificationResult @Optional +// The aggregated classification result. absl::Status ConfigureClassificationPostprocessingGraph( const tasks::core::ModelResources& model_resources, const proto::ClassifierOptions& classifier_options, diff --git a/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph_test.cc b/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph_test.cc index bb03e253..d4728e72 100644 --- a/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph_test.cc +++ b/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph_test.cc @@ -38,6 +38,7 @@ limitations under the License. #include "mediapipe/framework/output_stream_poller.h" #include "mediapipe/framework/port/gmock.h" #include "mediapipe/framework/port/gtest.h" +#include "mediapipe/framework/port/parse_text_proto.h" #include "mediapipe/framework/port/status_matchers.h" #include "mediapipe/framework/timestamp.h" #include "mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.pb.h" @@ -64,6 +65,7 @@ using ::mediapipe::file::JoinPath; using ::mediapipe::tasks::components::containers::proto::ClassificationResult; using ::mediapipe::tasks::core::ModelResources; using ::testing::HasSubstr; +using ::testing::Pointwise; using ::testing::proto::Approximately; constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/"; @@ -86,6 +88,11 @@ constexpr char kTimestampsTag[] = "TIMESTAMPS"; constexpr char kTimestampsName[] = "timestamps"; constexpr char kClassificationResultTag[] = "CLASSIFICATION_RESULT"; constexpr char kClassificationResultName[] = "classification_result"; +constexpr char kClassificationsTag[] = "CLASSIFICATIONS"; +constexpr char kClassificationsName[] = "classifications"; +constexpr char kTimestampedClassificationsTag[] = "TIMESTAMPED_CLASSIFICATIONS"; +constexpr char kTimestampedClassificationsName[] = + "timestamped_classifications"; // Helper function to get ModelResources. absl::StatusOr> CreateModelResourcesForModel( @@ -413,6 +420,316 @@ TEST_F(ConfigureTest, SucceedsWithMultipleHeads) { } class PostprocessingTest : public tflite_shims::testing::Test { + protected: + absl::StatusOr BuildGraph( + absl::string_view model_name, const proto::ClassifierOptions& options, + bool connect_timestamps = false) { + ASSIGN_OR_RETURN(auto model_resources, + CreateModelResourcesForModel(model_name)); + + Graph graph; + auto& postprocessing = graph.AddNode( + "mediapipe.tasks.components.processors." + "ClassificationPostprocessingGraph"); + MP_RETURN_IF_ERROR(ConfigureClassificationPostprocessingGraph( + *model_resources, options, + &postprocessing + .GetOptions())); + graph[Input>(kTensorsTag)].SetName(kTensorsName) >> + postprocessing.In(kTensorsTag); + if (connect_timestamps) { + graph[Input>(kTimestampsTag)].SetName( + kTimestampsName) >> + postprocessing.In(kTimestampsTag); + postprocessing.Out(kTimestampedClassificationsTag) + .SetName(kTimestampedClassificationsName) >> + graph[Output>( + kTimestampedClassificationsTag)]; + } else { + postprocessing.Out(kClassificationsTag).SetName(kClassificationsName) >> + graph[Output(kClassificationsTag)]; + } + + MP_RETURN_IF_ERROR(calculator_graph_.Initialize(graph.GetConfig())); + if (connect_timestamps) { + ASSIGN_OR_RETURN(auto poller, calculator_graph_.AddOutputStreamPoller( + kTimestampedClassificationsName)); + MP_RETURN_IF_ERROR(calculator_graph_.StartRun(/*extra_side_packets=*/{})); + return poller; + } + ASSIGN_OR_RETURN(auto poller, calculator_graph_.AddOutputStreamPoller( + kClassificationsName)); + MP_RETURN_IF_ERROR(calculator_graph_.StartRun(/*extra_side_packets=*/{})); + return poller; + } + + template + void AddTensor( + const std::vector& tensor, const Tensor::ElementType& element_type, + const Tensor::QuantizationParameters& quantization_parameters = {}) { + tensors_->emplace_back(element_type, + Tensor::Shape{1, static_cast(tensor.size())}, + quantization_parameters); + auto view = tensors_->back().GetCpuWriteView(); + T* buffer = view.buffer(); + std::copy(tensor.begin(), tensor.end(), buffer); + } + + absl::Status Run( + std::optional> aggregation_timestamps = std::nullopt, + int timestamp = 0) { + MP_RETURN_IF_ERROR(calculator_graph_.AddPacketToInputStream( + kTensorsName, Adopt(tensors_.release()).At(Timestamp(timestamp)))); + // Reset tensors for future calls. + tensors_ = absl::make_unique>(); + if (aggregation_timestamps.has_value()) { + auto packet = absl::make_unique>(); + for (const auto& timestamp : *aggregation_timestamps) { + packet->emplace_back(Timestamp(timestamp)); + } + MP_RETURN_IF_ERROR(calculator_graph_.AddPacketToInputStream( + kTimestampsName, Adopt(packet.release()).At(Timestamp(timestamp)))); + } + return absl::OkStatus(); + } + + template + absl::StatusOr GetResult(OutputStreamPoller& poller) { + MP_RETURN_IF_ERROR(calculator_graph_.WaitUntilIdle()); + MP_RETURN_IF_ERROR(calculator_graph_.CloseAllInputStreams()); + + Packet packet; + if (!poller.Next(&packet)) { + return absl::InternalError("Unable to get output packet"); + } + auto result = packet.Get(); + MP_RETURN_IF_ERROR(calculator_graph_.WaitUntilDone()); + return result; + } + + private: + CalculatorGraph calculator_graph_; + std::unique_ptr> tensors_ = + absl::make_unique>(); +}; + +TEST_F(PostprocessingTest, SucceedsWithoutMetadata) { + // Build graph. + proto::ClassifierOptions options; + options.set_max_results(3); + options.set_score_threshold(0.5); + MP_ASSERT_OK_AND_ASSIGN( + auto poller, + BuildGraph(kQuantizedImageClassifierWithoutMetadata, options)); + // Build input tensors. + std::vector tensor(kMobileNetNumClasses, 0); + tensor[1] = 18; + tensor[2] = 16; + + // Send tensors and get results. + AddTensor(tensor, Tensor::ElementType::kUInt8, + /*quantization_parameters=*/{0.1, 10}); + MP_ASSERT_OK(Run()); + MP_ASSERT_OK_AND_ASSIGN(auto results, + GetResult(poller)); + + // Validate results. + EXPECT_THAT(results, + EqualsProto(ParseTextProtoOrDie(R"pb( + classifications { + head_index: 0 + classification_list { + classification { index: 1 score: 0.8 } + classification { index: 2 score: 0.6 } + } + } + )pb"))); +} + +TEST_F(PostprocessingTest, SucceedsWithMetadata) { + // Build graph. + proto::ClassifierOptions options; + options.set_max_results(3); + MP_ASSERT_OK_AND_ASSIGN( + auto poller, BuildGraph(kQuantizedImageClassifierWithMetadata, options)); + // Build input tensors. + std::vector tensor(kMobileNetNumClasses, 0); + tensor[1] = 12; + tensor[2] = 14; + tensor[3] = 16; + tensor[4] = 18; + + // Send tensors and get results. + AddTensor(tensor, Tensor::ElementType::kUInt8, + /*quantization_parameters=*/{0.1, 10}); + MP_ASSERT_OK(Run()); + MP_ASSERT_OK_AND_ASSIGN(auto results, + GetResult(poller)); + + // Validate results. + EXPECT_THAT( + results, EqualsProto(ParseTextProtoOrDie(R"pb( + classifications { + head_index: 0 + head_name: "probability" + classification_list { + classification { index: 4 score: 0.8 label: "tiger shark" } + classification { index: 3 score: 0.6 label: "great white shark" } + classification { index: 2 score: 0.4 label: "goldfish" } + } + } + )pb"))); +} + +TEST_F(PostprocessingTest, SucceedsWithScoreCalibration) { + // Build graph. + proto::ClassifierOptions options; + options.set_max_results(3); + MP_ASSERT_OK_AND_ASSIGN( + auto poller, + BuildGraph(kQuantizedImageClassifierWithDummyScoreCalibration, options)); + // Build input tensors. + std::vector tensor(kMobileNetNumClasses, 0); + tensor[1] = 12; + tensor[2] = 14; + tensor[3] = 16; + tensor[4] = 18; + + // Send tensors and get results. + AddTensor(tensor, Tensor::ElementType::kUInt8, + /*quantization_parameters=*/{0.1, 10}); + MP_ASSERT_OK(Run()); + MP_ASSERT_OK_AND_ASSIGN(auto results, + GetResult(poller)); + + // Validate results. + EXPECT_THAT( + results, EqualsProto(ParseTextProtoOrDie(R"pb( + classifications { + head_index: 0 + head_name: "probability" + classification_list { + classification { index: 4 score: 0.6899744811 label: "tiger shark" } + classification { + index: 3 + score: 0.6456563062 + label: "great white shark" + } + classification { index: 2 score: 0.5986876601 label: "goldfish" } + } + } + )pb"))); +} + +TEST_F(PostprocessingTest, SucceedsWithMultipleHeads) { + // Build graph. + proto::ClassifierOptions options; + options.set_max_results(2); + MP_ASSERT_OK_AND_ASSIGN( + auto poller, + BuildGraph(kFloatTwoHeadsAudioClassifierWithMetadata, options)); + // Build input tensors. + std::vector tensor_0(kTwoHeadsNumClasses[0], 0); + tensor_0[1] = 0.2; + tensor_0[2] = 0.4; + tensor_0[3] = 0.6; + std::vector tensor_1(kTwoHeadsNumClasses[1], 0); + tensor_1[1] = 0.2; + tensor_1[2] = 0.4; + tensor_1[3] = 0.6; + + // Send tensors and get results. + AddTensor(tensor_0, Tensor::ElementType::kFloat32); + AddTensor(tensor_1, Tensor::ElementType::kFloat32); + MP_ASSERT_OK(Run()); + MP_ASSERT_OK_AND_ASSIGN(auto results, + GetResult(poller)); + + // Validate results. + EXPECT_THAT( + results, EqualsProto(ParseTextProtoOrDie(R"pb( + classifications { + head_index: 0 + head_name: "yamnet_classification" + classification_list { + classification { index: 3 score: 0.6 label: "Narration, monologue" } + classification { index: 2 score: 0.4 label: "Conversation" } + } + } + classifications { + head_index: 1 + head_name: "bird_classification" + classification_list { + classification { index: 3 score: 0.6 label: "Azara\'s Spinetail" } + classification { index: 2 score: 0.4 label: "House Sparrow" } + } + } + )pb"))); +} + +TEST_F(PostprocessingTest, SucceedsWithTimestamps) { + // Build graph. + proto::ClassifierOptions options; + options.set_max_results(2); + MP_ASSERT_OK_AND_ASSIGN( + auto poller, BuildGraph(kQuantizedImageClassifierWithMetadata, options, + /*connect_timestamps=*/true)); + // Build input tensors. + std::vector tensor_0(kMobileNetNumClasses, 0); + tensor_0[1] = 12; + tensor_0[2] = 14; + tensor_0[3] = 16; + std::vector tensor_1(kMobileNetNumClasses, 0); + tensor_1[5] = 12; + tensor_1[6] = 14; + tensor_1[7] = 16; + + // Send tensors and get results. + AddTensor(tensor_0, Tensor::ElementType::kUInt8, + /*quantization_parameters=*/{0.1, 10}); + MP_ASSERT_OK(Run()); + AddTensor(tensor_1, Tensor::ElementType::kUInt8, + /*quantization_parameters=*/{0.1, 10}); + MP_ASSERT_OK(Run( + /*aggregation_timestamps=*/std::optional>({0, 1000}), + /*timestamp=*/1000)); + + MP_ASSERT_OK_AND_ASSIGN(auto results, + GetResult>(poller)); + + // Validate results. + EXPECT_THAT( + results, + Pointwise( + EqualsProto(), + {ParseTextProtoOrDie(R"pb( + timestamp_ms: 0 + classifications { + head_index: 0 + head_name: "probability" + classification_list { + classification { + index: 3 + score: 0.6 + label: "great white shark" + } + classification { index: 2 score: 0.4 label: "goldfish" } + } + })pb"), + ParseTextProtoOrDie(R"pb( + timestamp_ms: 1 + classifications { + head_index: 0 + head_name: "probability" + classification_list { + classification { index: 7 score: 0.6 label: "stingray" } + classification { index: 6 score: 0.4 label: "electric ray" } + } + })pb")})); +} + +// TODO: remove these tests once migration is over. +class LegacyPostprocessingTest : public tflite_shims::testing::Test { protected: absl::StatusOr BuildGraph( absl::string_view model_name, const proto::ClassifierOptions& options, @@ -496,7 +813,7 @@ class PostprocessingTest : public tflite_shims::testing::Test { absl::make_unique>(); }; -TEST_F(PostprocessingTest, SucceedsWithoutMetadata) { +TEST_F(LegacyPostprocessingTest, SucceedsWithoutMetadata) { // Build graph. proto::ClassifierOptions options; options.set_max_results(3); @@ -525,7 +842,7 @@ TEST_F(PostprocessingTest, SucceedsWithoutMetadata) { })pb")); } -TEST_F(PostprocessingTest, SucceedsWithMetadata) { +TEST_F(LegacyPostprocessingTest, SucceedsWithMetadata) { // Build graph. proto::ClassifierOptions options; options.set_max_results(3); @@ -568,7 +885,7 @@ TEST_F(PostprocessingTest, SucceedsWithMetadata) { })pb")); } -TEST_F(PostprocessingTest, SucceedsWithScoreCalibration) { +TEST_F(LegacyPostprocessingTest, SucceedsWithScoreCalibration) { // Build graph. proto::ClassifierOptions options; options.set_max_results(3); @@ -614,7 +931,7 @@ TEST_F(PostprocessingTest, SucceedsWithScoreCalibration) { })pb")); } -TEST_F(PostprocessingTest, SucceedsWithMultipleHeads) { +TEST_F(LegacyPostprocessingTest, SucceedsWithMultipleHeads) { // Build graph. proto::ClassifierOptions options; options.set_max_results(2); @@ -674,7 +991,7 @@ TEST_F(PostprocessingTest, SucceedsWithMultipleHeads) { })pb")); } -TEST_F(PostprocessingTest, SucceedsWithTimestamps) { +TEST_F(LegacyPostprocessingTest, SucceedsWithTimestamps) { // Build graph. proto::ClassifierOptions options; options.set_max_results(2); diff --git a/mediapipe/tasks/cc/components/processors/proto/classification_postprocessing_graph_options.proto b/mediapipe/tasks/cc/components/processors/proto/classification_postprocessing_graph_options.proto index 1de788ea..84ba9522 100644 --- a/mediapipe/tasks/cc/components/processors/proto/classification_postprocessing_graph_options.proto +++ b/mediapipe/tasks/cc/components/processors/proto/classification_postprocessing_graph_options.proto @@ -38,7 +38,7 @@ message ClassificationPostprocessingGraphOptions { // Options for the ClassificationAggregationCalculator encapsulated by the // ClassificationPostprocessing subgraph. - optional ClassificationAggregationCalculatorOptions + optional mediapipe.ClassificationAggregationCalculatorOptions classification_aggregation_options = 2; // Whether output tensors are quantized (kTfLiteUint8) or not (kFloat32). diff --git a/mediapipe/tasks/cc/core/model_task_graph.cc b/mediapipe/tasks/cc/core/model_task_graph.cc index 47334b67..66434483 100644 --- a/mediapipe/tasks/cc/core/model_task_graph.cc +++ b/mediapipe/tasks/cc/core/model_task_graph.cc @@ -156,21 +156,24 @@ absl::StatusOr ModelTaskGraph::GetConfig( } absl::StatusOr ModelTaskGraph::CreateModelResources( - SubgraphContext* sc, std::unique_ptr external_file) { + SubgraphContext* sc, std::unique_ptr external_file, + const std::string tag_suffix) { auto model_resources_cache_service = sc->Service(kModelResourcesCacheService); if (!model_resources_cache_service.IsAvailable()) { - ASSIGN_OR_RETURN(local_model_resources_, + ASSIGN_OR_RETURN(auto local_model_resource, ModelResources::Create("", std::move(external_file))); LOG(WARNING) << "A local ModelResources object is created. Please consider using " "ModelResourcesCacheService to cache the created ModelResources " "object in the CalculatorGraph."; - return local_model_resources_.get(); + local_model_resources_.push_back(std::move(local_model_resource)); + return local_model_resources_.back().get(); } ASSIGN_OR_RETURN( auto op_resolver_packet, model_resources_cache_service.GetObject().GetGraphOpResolverPacket()); - const std::string tag = CreateModelResourcesTag(sc->OriginalNode()); + const std::string tag = + absl::StrCat(CreateModelResourcesTag(sc->OriginalNode()), tag_suffix); ASSIGN_OR_RETURN(auto model_resources, ModelResources::Create(tag, std::move(external_file), op_resolver_packet)); @@ -182,7 +185,8 @@ absl::StatusOr ModelTaskGraph::CreateModelResources( absl::StatusOr ModelTaskGraph::CreateModelAssetBundleResources( - SubgraphContext* sc, std::unique_ptr external_file) { + SubgraphContext* sc, std::unique_ptr external_file, + const std::string tag_suffix) { auto model_resources_cache_service = sc->Service(kModelResourcesCacheService); bool has_file_pointer_meta = external_file->has_file_pointer_meta(); // if external file is set by file pointer, no need to add the model asset @@ -190,7 +194,7 @@ ModelTaskGraph::CreateModelAssetBundleResources( // not owned by this model asset bundle resources. if (!model_resources_cache_service.IsAvailable() || has_file_pointer_meta) { ASSIGN_OR_RETURN( - local_model_asset_bundle_resources_, + auto local_model_asset_bundle_resource, ModelAssetBundleResources::Create("", std::move(external_file))); if (!has_file_pointer_meta) { LOG(WARNING) @@ -198,10 +202,12 @@ ModelTaskGraph::CreateModelAssetBundleResources( "ModelResourcesCacheService to cache the created ModelResources " "object in the CalculatorGraph."; } - return local_model_asset_bundle_resources_.get(); + local_model_asset_bundle_resources_.push_back( + std::move(local_model_asset_bundle_resource)); + return local_model_asset_bundle_resources_.back().get(); } - const std::string tag = - CreateModelAssetBundleResourcesTag(sc->OriginalNode()); + const std::string tag = absl::StrCat( + CreateModelAssetBundleResourcesTag(sc->OriginalNode()), tag_suffix); ASSIGN_OR_RETURN( auto model_bundle_resources, ModelAssetBundleResources::Create(tag, std::move(external_file))); diff --git a/mediapipe/tasks/cc/core/model_task_graph.h b/mediapipe/tasks/cc/core/model_task_graph.h index 5ee70e8f..50dcc903 100644 --- a/mediapipe/tasks/cc/core/model_task_graph.h +++ b/mediapipe/tasks/cc/core/model_task_graph.h @@ -19,6 +19,7 @@ limitations under the License. #include #include #include +#include #include "absl/status/statusor.h" #include "absl/strings/str_format.h" @@ -75,9 +76,14 @@ class ModelTaskGraph : public Subgraph { // construction stage. Note that the external file contents will be moved // into the model resources object on creation. The returned model resources // pointer will provide graph authors with the access to the metadata - // extractor and the tflite model. + // extractor and the tflite model. When the model resources graph service is + // available, a tag is generated internally asscoiated with the created model + // resource. If more than one model resources are created in a graph, the + // model resources graph service add the tag_suffix to support multiple + // resources. absl::StatusOr CreateModelResources( - SubgraphContext* sc, std::unique_ptr external_file); + SubgraphContext* sc, std::unique_ptr external_file, + const std::string tag_suffix = ""); // If the model resources graph service is available, creates a model asset // bundle resources object from the subgraph context, and caches the created @@ -103,10 +109,15 @@ class ModelTaskGraph : public Subgraph { // that can only be used in the graph construction stage. Note that the // external file contents will be moved into the model asset bundle resources // object on creation. The returned model asset bundle resources pointer will - // provide graph authors with the access to extracted model files. + // provide graph authors with the access to extracted model files. When the + // model resources graph service is available, a tag is generated internally + // asscoiated with the created model asset bundle resource. If more than one + // model asset bundle resources are created in a graph, the model resources + // graph service add the tag_suffix to support multiple resources. absl::StatusOr CreateModelAssetBundleResources( - SubgraphContext* sc, std::unique_ptr external_file); + SubgraphContext* sc, std::unique_ptr external_file, + const std::string tag_suffix = ""); // Inserts a mediapipe task inference subgraph into the provided // GraphBuilder. The returned node provides the following interfaces to the @@ -124,9 +135,9 @@ class ModelTaskGraph : public Subgraph { api2::builder::Graph& graph) const; private: - std::unique_ptr local_model_resources_; + std::vector> local_model_resources_; - std::unique_ptr + std::vector> local_model_asset_bundle_resources_; }; diff --git a/mediapipe/tasks/cc/text/text_classifier/BUILD b/mediapipe/tasks/cc/text/text_classifier/BUILD index a8553863..b2e1bed2 100644 --- a/mediapipe/tasks/cc/text/text_classifier/BUILD +++ b/mediapipe/tasks/cc/text/text_classifier/BUILD @@ -49,6 +49,8 @@ cc_library( ":text_classifier_graph", "//mediapipe/framework:packet", "//mediapipe/framework/api2:builder", + "//mediapipe/tasks/cc/components/containers:category", + "//mediapipe/tasks/cc/components/containers:classification_result", "//mediapipe/tasks/cc/components/containers/proto:classifications_cc_proto", "//mediapipe/tasks/cc/components/processors:classifier_options", "//mediapipe/tasks/cc/components/processors/proto:classifier_options_cc_proto", @@ -63,6 +65,30 @@ cc_library( ], ) +cc_test( + name = "text_classifier_test", + srcs = ["text_classifier_test.cc"], + data = [ + "//mediapipe/tasks/testdata/text:bert_text_classifier_models", + "//mediapipe/tasks/testdata/text:text_classifier_models", + ], + deps = [ + ":text_classifier", + ":text_classifier_test_utils", + "//mediapipe/framework/deps:file_path", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/tasks/cc:common", + "//mediapipe/tasks/cc/components/containers:category", + "//mediapipe/tasks/cc/components/containers:classification_result", + "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:cord", + "@org_tensorflow//tensorflow/lite/core/shims:cc_shims_test_util", + ], +) + cc_library( name = "text_classifier_test_utils", srcs = ["text_classifier_test_utils.cc"], diff --git a/mediapipe/tasks/cc/text/text_classifier/text_classifier.cc b/mediapipe/tasks/cc/text/text_classifier/text_classifier.cc index 699f15bc..d174fac4 100644 --- a/mediapipe/tasks/cc/text/text_classifier/text_classifier.cc +++ b/mediapipe/tasks/cc/text/text_classifier/text_classifier.cc @@ -24,6 +24,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "mediapipe/framework/api2/builder.h" #include "mediapipe/framework/packet.h" +#include "mediapipe/tasks/cc/components/containers/classification_result.h" #include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h" #include "mediapipe/tasks/cc/components/processors/proto/classifier_options.pb.h" #include "mediapipe/tasks/cc/core/task_api_factory.h" @@ -37,12 +38,13 @@ namespace text_classifier { namespace { +using ::mediapipe::tasks::components::containers::ConvertToClassificationResult; using ::mediapipe::tasks::components::containers::proto::ClassificationResult; constexpr char kTextStreamName[] = "text_in"; constexpr char kTextTag[] = "TEXT"; -constexpr char kClassificationResultStreamName[] = "classification_result_out"; -constexpr char kClassificationResultTag[] = "CLASSIFICATION_RESULT"; +constexpr char kClassificationsStreamName[] = "classifications_out"; +constexpr char kClassificationsTag[] = "CLASSIFICATIONS"; constexpr char kSubgraphTypeName[] = "mediapipe.tasks.text.text_classifier.TextClassifierGraph"; @@ -54,9 +56,8 @@ CalculatorGraphConfig CreateGraphConfig( auto& subgraph = graph.AddNode(kSubgraphTypeName); subgraph.GetOptions().Swap(options.get()); graph.In(kTextTag).SetName(kTextStreamName) >> subgraph.In(kTextTag); - subgraph.Out(kClassificationResultTag) - .SetName(kClassificationResultStreamName) >> - graph.Out(kClassificationResultTag); + subgraph.Out(kClassificationsTag).SetName(kClassificationsStreamName) >> + graph.Out(kClassificationsTag); return graph.GetConfig(); } @@ -88,14 +89,14 @@ absl::StatusOr> TextClassifier::Create( std::move(options->base_options.op_resolver)); } -absl::StatusOr TextClassifier::Classify( +absl::StatusOr TextClassifier::Classify( absl::string_view text) { ASSIGN_OR_RETURN( auto output_packets, runner_->Process( {{kTextStreamName, MakePacket(std::string(text))}})); - return output_packets[kClassificationResultStreamName] - .Get(); + return ConvertToClassificationResult( + output_packets[kClassificationsStreamName].Get()); } } // namespace text_classifier diff --git a/mediapipe/tasks/cc/text/text_classifier/text_classifier.h b/mediapipe/tasks/cc/text/text_classifier/text_classifier.h index b027a978..03569c5a 100644 --- a/mediapipe/tasks/cc/text/text_classifier/text_classifier.h +++ b/mediapipe/tasks/cc/text/text_classifier/text_classifier.h @@ -21,7 +21,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" -#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h" +#include "mediapipe/tasks/cc/components/containers/classification_result.h" #include "mediapipe/tasks/cc/components/processors/classifier_options.h" #include "mediapipe/tasks/cc/core/base_options.h" #include "mediapipe/tasks/cc/core/base_task_api.h" @@ -31,6 +31,10 @@ namespace tasks { namespace text { namespace text_classifier { +// Alias the shared ClassificationResult struct as result type. +using TextClassifierResult = + ::mediapipe::tasks::components::containers::ClassificationResult; + // The options for configuring a MediaPipe text classifier task. struct TextClassifierOptions { // Base options for configuring MediaPipe Tasks, such as specifying the model @@ -81,8 +85,7 @@ class TextClassifier : core::BaseTaskApi { std::unique_ptr options); // Performs classification on the input `text`. - absl::StatusOr Classify( - absl::string_view text); + absl::StatusOr Classify(absl::string_view text); // Shuts down the TextClassifier when all the work is done. absl::Status Close() { return runner_->Close(); } diff --git a/mediapipe/tasks/cc/text/text_classifier/text_classifier_graph.cc b/mediapipe/tasks/cc/text/text_classifier/text_classifier_graph.cc index 9706db4d..36ff68a0 100644 --- a/mediapipe/tasks/cc/text/text_classifier/text_classifier_graph.cc +++ b/mediapipe/tasks/cc/text/text_classifier/text_classifier_graph.cc @@ -47,10 +47,18 @@ using ::mediapipe::tasks::components::containers::proto::ClassificationResult; using ::mediapipe::tasks::core::ModelResources; constexpr char kClassificationResultTag[] = "CLASSIFICATION_RESULT"; +constexpr char kClassificationsTag[] = "CLASSIFICATIONS"; constexpr char kTextTag[] = "TEXT"; constexpr char kMetadataExtractorTag[] = "METADATA_EXTRACTOR"; constexpr char kTensorsTag[] = "TENSORS"; +// TODO: remove once Java API migration is over. +// Struct holding the different output streams produced by the text classifier. +struct TextClassifierOutputStreams { + Source classification_result; + Source classifications; +}; + } // namespace // A "TextClassifierGraph" performs Natural Language classification (including @@ -62,7 +70,10 @@ constexpr char kTensorsTag[] = "TENSORS"; // Input text to perform classification on. // // Outputs: -// CLASSIFICATION_RESULT - ClassificationResult +// CLASSIFICATIONS - ClassificationResult @Optional +// The classification results aggregated by classifier head. +// TODO: remove once Java API migration is over. +// CLASSIFICATION_RESULT - (DEPRECATED) ClassificationResult @Optional // The aggregated classification result object that has 3 dimensions: // (classification head, classification timestamp, classification category). // @@ -70,7 +81,7 @@ constexpr char kTensorsTag[] = "TENSORS"; // node { // calculator: "mediapipe.tasks.text.text_classifier.TextClassifierGraph" // input_stream: "TEXT:text_in" -// output_stream: "CLASSIFICATION_RESULT:classification_result_out" +// output_stream: "CLASSIFICATIONS:classifications_out" // options { // [mediapipe.tasks.text.text_classifier.proto.TextClassifierGraphOptions.ext] // { @@ -91,12 +102,14 @@ class TextClassifierGraph : public core::ModelTaskGraph { CreateModelResources(sc)); Graph graph; ASSIGN_OR_RETURN( - Source classification_result_out, + auto output_streams, BuildTextClassifierTask( sc->Options(), *model_resources, graph[Input(kTextTag)], graph)); - classification_result_out >> + output_streams.classification_result >> graph[Output(kClassificationResultTag)]; + output_streams.classifications >> + graph[Output(kClassificationsTag)]; return graph.GetConfig(); } @@ -111,7 +124,7 @@ class TextClassifierGraph : public core::ModelTaskGraph { // TextClassifier model file with model metadata. // text_in: (std::string) stream to run text classification on. // graph: the mediapipe builder::Graph instance to be updated. - absl::StatusOr> BuildTextClassifierTask( + absl::StatusOr BuildTextClassifierTask( const proto::TextClassifierGraphOptions& task_options, const ModelResources& model_resources, Source text_in, Graph& graph) { @@ -148,8 +161,11 @@ class TextClassifierGraph : public core::ModelTaskGraph { // Outputs the aggregated classification result as the subgraph output // stream. - return postprocessing[Output( - kClassificationResultTag)]; + return TextClassifierOutputStreams{ + /*classification_result=*/postprocessing[Output( + kClassificationResultTag)], + /*classifications=*/postprocessing[Output( + kClassificationsTag)]}; } }; diff --git a/mediapipe/tasks/cc/text/text_classifier/text_classifier_test.cc b/mediapipe/tasks/cc/text/text_classifier/text_classifier_test.cc index 5b33f660..8f73914f 100644 --- a/mediapipe/tasks/cc/text/text_classifier/text_classifier_test.cc +++ b/mediapipe/tasks/cc/text/text_classifier/text_classifier_test.cc @@ -33,7 +33,8 @@ limitations under the License. #include "mediapipe/framework/port/gtest.h" #include "mediapipe/framework/port/status_matchers.h" #include "mediapipe/tasks/cc/common.h" -#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h" +#include "mediapipe/tasks/cc/components/containers/category.h" +#include "mediapipe/tasks/cc/components/containers/classification_result.h" #include "mediapipe/tasks/cc/text/text_classifier/text_classifier_test_utils.h" #include "tensorflow/lite/core/shims/cc/shims_test_util.h" @@ -43,17 +44,13 @@ namespace text { namespace text_classifier { namespace { -using ::mediapipe::EqualsProto; using ::mediapipe::file::JoinPath; using ::mediapipe::tasks::kMediaPipeTasksPayload; -using ::mediapipe::tasks::components::containers::proto::ClassificationResult; +using ::mediapipe::tasks::components::containers::Category; +using ::mediapipe::tasks::components::containers::Classifications; using ::testing::HasSubstr; using ::testing::Optional; -using ::testing::proto::Approximately; -using ::testing::proto::IgnoringRepeatedFieldOrdering; -using ::testing::proto::Partially; -constexpr float kEpsilon = 0.001; constexpr int kMaxSeqLen = 128; constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/text/"; constexpr char kTestBertModelPath[] = "bert_text_classifier.tflite"; @@ -67,6 +64,30 @@ std::string GetFullPath(absl::string_view file_name) { return JoinPath("./", kTestDataDirectory, file_name); } +// Checks that the two provided `TextClassifierResult` are equal, with a +// tolerancy on floating-point score to account for numerical instabilities. +// TODO: create shared matcher for ClassificationResult. +void ExpectApproximatelyEqual(const TextClassifierResult& actual, + const TextClassifierResult& expected) { + const float kPrecision = 1e-6; + ASSERT_EQ(actual.classifications.size(), expected.classifications.size()); + for (int i = 0; i < actual.classifications.size(); ++i) { + const Classifications& a = actual.classifications[i]; + const Classifications& b = expected.classifications[i]; + EXPECT_EQ(a.head_index, b.head_index); + EXPECT_EQ(a.head_name, b.head_name); + EXPECT_EQ(a.categories.size(), b.categories.size()); + for (int j = 0; j < a.categories.size(); ++j) { + const Category& x = a.categories[j]; + const Category& y = b.categories[j]; + EXPECT_EQ(x.index, y.index); + EXPECT_NEAR(x.score, y.score, kPrecision); + EXPECT_EQ(x.category_name, y.category_name); + EXPECT_EQ(x.display_name, y.display_name); + } + } +} + class TextClassifierTest : public tflite_shims::testing::Test {}; TEST_F(TextClassifierTest, CreateSucceedsWithBertModel) { @@ -116,34 +137,29 @@ TEST_F(TextClassifierTest, TextClassifierWithBert) { MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr classifier, TextClassifier::Create(std::move(options))); MP_ASSERT_OK_AND_ASSIGN( - ClassificationResult negative_result, + TextClassifierResult negative_result, classifier->Classify("unflinchingly bleak and desperate")); - ASSERT_THAT(negative_result, - Partially(IgnoringRepeatedFieldOrdering(Approximately( - EqualsProto(R"pb( - classifications { - entries { - categories { category_name: "negative" score: 0.956 } - categories { category_name: "positive" score: 0.044 } - } - } - )pb"), - kEpsilon)))); + TextClassifierResult negative_expected; + negative_expected.classifications.emplace_back(Classifications{ + /*categories=*/{ + {/*index=*/0, /*score=*/0.956317, /*category_name=*/"negative"}, + {/*index=*/1, /*score=*/0.043683, /*category_name=*/"positive"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(negative_result, negative_expected); MP_ASSERT_OK_AND_ASSIGN( - ClassificationResult positive_result, + TextClassifierResult positive_result, classifier->Classify("it's a charming and often affecting journey")); - ASSERT_THAT(positive_result, - Partially(IgnoringRepeatedFieldOrdering(Approximately( - EqualsProto(R"pb( - classifications { - entries { - categories { category_name: "negative" score: 0.0 } - categories { category_name: "positive" score: 1.0 } - } - } - )pb"), - kEpsilon)))); + TextClassifierResult positive_expected; + positive_expected.classifications.emplace_back(Classifications{ + /*categories=*/{ + {/*index=*/1, /*score=*/0.999945, /*category_name=*/"positive"}, + {/*index=*/0, /*score=*/0.000056, /*category_name=*/"negative"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(positive_result, positive_expected); + MP_ASSERT_OK(classifier->Close()); } @@ -152,35 +168,30 @@ TEST_F(TextClassifierTest, TextClassifierWithIntInputs) { options->base_options.model_asset_path = GetFullPath(kTestRegexModelPath); MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr classifier, TextClassifier::Create(std::move(options))); - MP_ASSERT_OK_AND_ASSIGN(ClassificationResult negative_result, + MP_ASSERT_OK_AND_ASSIGN(TextClassifierResult negative_result, classifier->Classify("What a waste of my time.")); - ASSERT_THAT(negative_result, - Partially(IgnoringRepeatedFieldOrdering(Approximately( - EqualsProto(R"pb( - classifications { - entries { - categories { category_name: "Negative" score: 0.813 } - categories { category_name: "Positive" score: 0.187 } - } - } - )pb"), - kEpsilon)))); + TextClassifierResult negative_expected; + negative_expected.classifications.emplace_back(Classifications{ + /*categories=*/{ + {/*index=*/0, /*score=*/0.813130, /*category_name=*/"Negative"}, + {/*index=*/1, /*score=*/0.186870, /*category_name=*/"Positive"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(negative_result, negative_expected); MP_ASSERT_OK_AND_ASSIGN( - ClassificationResult positive_result, - classifier->Classify("This is the best movie I’ve seen in recent years. " + TextClassifierResult positive_result, + classifier->Classify("This is the best movie I’ve seen in recent years." "Strongly recommend it!")); - ASSERT_THAT(positive_result, - Partially(IgnoringRepeatedFieldOrdering(Approximately( - EqualsProto(R"pb( - classifications { - entries { - categories { category_name: "Negative" score: 0.487 } - categories { category_name: "Positive" score: 0.513 } - } - } - )pb"), - kEpsilon)))); + TextClassifierResult positive_expected; + positive_expected.classifications.emplace_back(Classifications{ + /*categories=*/{ + {/*index=*/1, /*score=*/0.513427, /*category_name=*/"Positive"}, + {/*index=*/0, /*score=*/0.486573, /*category_name=*/"Negative"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(positive_result, positive_expected); + MP_ASSERT_OK(classifier->Close()); } @@ -190,44 +201,19 @@ TEST_F(TextClassifierTest, TextClassifierWithStringToBool) { options->base_options.op_resolver = CreateCustomResolver(); MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr classifier, TextClassifier::Create(std::move(options))); - MP_ASSERT_OK_AND_ASSIGN(ClassificationResult result, + MP_ASSERT_OK_AND_ASSIGN(TextClassifierResult result, classifier->Classify("hello")); - ASSERT_THAT(result, Partially(IgnoringRepeatedFieldOrdering(EqualsProto(R"pb( - classifications { - entries { - categories { index: 1 score: 1 } - categories { index: 0 score: 1 } - categories { index: 2 score: 0 } - } - } - )pb")))); -} -TEST_F(TextClassifierTest, BertLongPositive) { - std::stringstream ss_for_positive_review; - ss_for_positive_review - << "it's a charming and often affecting journey and this is a long"; - for (int i = 0; i < kMaxSeqLen; ++i) { - ss_for_positive_review << " long"; - } - ss_for_positive_review << " movie review"; - auto options = std::make_unique(); - options->base_options.model_asset_path = GetFullPath(kTestBertModelPath); - MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr classifier, - TextClassifier::Create(std::move(options))); - MP_ASSERT_OK_AND_ASSIGN(ClassificationResult result, - classifier->Classify(ss_for_positive_review.str())); - ASSERT_THAT(result, - Partially(IgnoringRepeatedFieldOrdering(Approximately( - EqualsProto(R"pb( - classifications { - entries { - categories { category_name: "negative" score: 0.014 } - categories { category_name: "positive" score: 0.986 } - } - } - )pb"), - kEpsilon)))); + // Binary outputs causes flaky ordering, so we compare manually. + ASSERT_EQ(result.classifications.size(), 1); + ASSERT_EQ(result.classifications[0].head_index, 0); + ASSERT_EQ(result.classifications[0].categories.size(), 3); + ASSERT_EQ(result.classifications[0].categories[0].score, 1); + ASSERT_LT(result.classifications[0].categories[0].index, 2); // i.e O or 1. + ASSERT_EQ(result.classifications[0].categories[1].score, 1); + ASSERT_LT(result.classifications[0].categories[1].index, 2); // i.e 0 or 1. + ASSERT_EQ(result.classifications[0].categories[2].score, 0); + ASSERT_EQ(result.classifications[0].categories[2].index, 2); MP_ASSERT_OK(classifier->Close()); } diff --git a/mediapipe/tasks/cc/text/tokenizers/BUILD b/mediapipe/tasks/cc/text/tokenizers/BUILD index 048c7021..5ce08b2d 100644 --- a/mediapipe/tasks/cc/text/tokenizers/BUILD +++ b/mediapipe/tasks/cc/text/tokenizers/BUILD @@ -73,7 +73,18 @@ cc_library( ], ) -# TODO: This test fails in OSS +cc_test( + name = "sentencepiece_tokenizer_test", + srcs = ["sentencepiece_tokenizer_test.cc"], + data = [ + "//mediapipe/tasks/testdata/text:albert_model", + ], + deps = [ + ":sentencepiece_tokenizer", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/tasks/cc/core:utils", + ], +) cc_library( name = "tokenizer_utils", @@ -97,7 +108,32 @@ cc_library( ], ) -# TODO: This test fails in OSS +cc_test( + name = "tokenizer_utils_test", + srcs = ["tokenizer_utils_test.cc"], + data = [ + "//mediapipe/tasks/testdata/text:albert_model", + "//mediapipe/tasks/testdata/text:mobile_bert_model", + "//mediapipe/tasks/testdata/text:text_classifier_models", + ], + linkopts = ["-ldl"], + deps = [ + ":bert_tokenizer", + ":regex_tokenizer", + ":sentencepiece_tokenizer", + ":tokenizer_utils", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:status", + "//mediapipe/tasks/cc:common", + "//mediapipe/tasks/cc/core:utils", + "//mediapipe/tasks/cc/metadata:metadata_extractor", + "//mediapipe/tasks/metadata:metadata_schema_cc", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:cord", + ], +) cc_library( name = "regex_tokenizer", diff --git a/mediapipe/tasks/cc/vision/core/BUILD b/mediapipe/tasks/cc/vision/core/BUILD index 12d78990..e8e197a1 100644 --- a/mediapipe/tasks/cc/vision/core/BUILD +++ b/mediapipe/tasks/cc/vision/core/BUILD @@ -21,12 +21,23 @@ cc_library( hdrs = ["running_mode.h"], ) +cc_library( + name = "image_processing_options", + hdrs = ["image_processing_options.h"], + deps = [ + "//mediapipe/tasks/cc/components/containers:rect", + ], +) + cc_library( name = "base_vision_task_api", hdrs = ["base_vision_task_api.h"], deps = [ + ":image_processing_options", ":running_mode", "//mediapipe/calculators/core:flow_limiter_calculator", + "//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", "@com_google_absl//absl/status", 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 4586cbbd..c3c0a026 100644 --- a/mediapipe/tasks/cc/vision/core/base_vision_task_api.h +++ b/mediapipe/tasks/cc/vision/core/base_vision_task_api.h @@ -16,15 +16,20 @@ limitations under the License. #ifndef MEDIAPIPE_TASKS_CC_VISION_CORE_BASE_VISION_TASK_API_H_ #define MEDIAPIPE_TASKS_CC_VISION_CORE_BASE_VISION_TASK_API_H_ +#include #include +#include #include #include #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.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" namespace mediapipe { @@ -87,6 +92,60 @@ class BaseVisionTaskApi : public tasks::core::BaseTaskApi { return runner_->Send(std::move(inputs)); } + // Convert from ImageProcessingOptions to NormalizedRect, performing sanity + // checks on-the-fly. If the input ImageProcessingOptions is not present, + // returns a default NormalizedRect covering the whole image with rotation set + // to 0. If 'roi_allowed' is false, an error will be returned if the input + // ImageProcessingOptions has its 'region_or_interest' field set. + static absl::StatusOr ConvertToNormalizedRect( + std::optional options, bool roi_allowed = true) { + mediapipe::NormalizedRect normalized_rect; + normalized_rect.set_rotation(0); + normalized_rect.set_x_center(0.5); + normalized_rect.set_y_center(0.5); + normalized_rect.set_width(1.0); + normalized_rect.set_height(1.0); + if (!options.has_value()) { + return normalized_rect; + } + + if (options->rotation_degrees % 90 != 0) { + return CreateStatusWithPayload( + absl::StatusCode::kInvalidArgument, + "Expected rotation to be a multiple of 90°.", + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError); + } + // Convert to radians counter-clockwise. + normalized_rect.set_rotation(-options->rotation_degrees * M_PI / 180.0); + + if (options->region_of_interest.has_value()) { + if (!roi_allowed) { + return CreateStatusWithPayload( + absl::StatusCode::kInvalidArgument, + "This task doesn't support region-of-interest.", + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError); + } + auto& roi = *options->region_of_interest; + if (roi.left >= roi.right || roi.top >= roi.bottom) { + return CreateStatusWithPayload( + absl::StatusCode::kInvalidArgument, + "Expected Rect with left < right and top < bottom.", + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError); + } + if (roi.left < 0 || roi.top < 0 || roi.right > 1 || roi.bottom > 1) { + return CreateStatusWithPayload( + absl::StatusCode::kInvalidArgument, + "Expected Rect values to be in [0,1].", + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError); + } + normalized_rect.set_x_center((roi.left + roi.right) / 2.0); + normalized_rect.set_y_center((roi.top + roi.bottom) / 2.0); + normalized_rect.set_width(roi.right - roi.left); + normalized_rect.set_height(roi.bottom - roi.top); + } + return normalized_rect; + } + private: RunningMode running_mode_; }; diff --git a/mediapipe/tasks/cc/vision/core/image_processing_options.h b/mediapipe/tasks/cc/vision/core/image_processing_options.h new file mode 100644 index 00000000..7e764c1f --- /dev/null +++ b/mediapipe/tasks/cc/vision/core/image_processing_options.h @@ -0,0 +1,52 @@ +/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef MEDIAPIPE_TASKS_CC_VISION_CORE_IMAGE_PROCESSING_OPTIONS_H_ +#define MEDIAPIPE_TASKS_CC_VISION_CORE_IMAGE_PROCESSING_OPTIONS_H_ + +#include + +#include "mediapipe/tasks/cc/components/containers/rect.h" + +namespace mediapipe { +namespace tasks { +namespace vision { +namespace core { + +// Options for image processing. +// +// If both region-or-interest and rotation are specified, the crop around the +// region-of-interest is extracted first, the the specified rotation is applied +// to the crop. +struct ImageProcessingOptions { + // The optional region-of-interest to crop from the image. If not specified, + // the full image is used. + // + // Coordinates must be in [0,1] with 'left' < 'right' and 'top' < bottom. + std::optional region_of_interest = std::nullopt; + + // The rotation to apply to the image (or cropped region-of-interest), in + // degrees clockwise. + // + // The rotation must be a multiple (positive or negative) of 90°. + int rotation_degrees = 0; +}; + +} // namespace core +} // namespace vision +} // namespace tasks +} // namespace mediapipe + +#endif // MEDIAPIPE_TASKS_CC_VISION_CORE_IMAGE_PROCESSING_OPTIONS_H_ diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/BUILD b/mediapipe/tasks/cc/vision/gesture_recognizer/BUILD index 985c25cf..f32d4cc5 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/BUILD +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/BUILD @@ -56,22 +56,31 @@ cc_library( "//mediapipe/framework/formats:classification_cc_proto", "//mediapipe/framework/formats:landmark_cc_proto", "//mediapipe/framework/formats:matrix", + "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/framework/formats:tensor", "//mediapipe/tasks/cc:common", "//mediapipe/tasks/cc/components:image_preprocessing", "//mediapipe/tasks/cc/components/processors:classification_postprocessing_graph", "//mediapipe/tasks/cc/components/processors/proto:classification_postprocessing_graph_options_cc_proto", + "//mediapipe/tasks/cc/components/processors/proto:classifier_options_cc_proto", + "//mediapipe/tasks/cc/core:model_asset_bundle_resources", "//mediapipe/tasks/cc/core:model_resources", + "//mediapipe/tasks/cc/core:model_resources_cache", "//mediapipe/tasks/cc/core:model_task_graph", "//mediapipe/tasks/cc/core:utils", + "//mediapipe/tasks/cc/core/proto:base_options_cc_proto", + "//mediapipe/tasks/cc/core/proto:external_file_cc_proto", "//mediapipe/tasks/cc/core/proto:inference_subgraph_cc_proto", + "//mediapipe/tasks/cc/metadata/utils:zip_utils", + "//mediapipe/tasks/cc/vision/gesture_recognizer/calculators:combined_prediction_calculator", + "//mediapipe/tasks/cc/vision/gesture_recognizer/calculators:combined_prediction_calculator_cc_proto", "//mediapipe/tasks/cc/vision/gesture_recognizer/calculators:handedness_to_matrix_calculator", "//mediapipe/tasks/cc/vision/gesture_recognizer/calculators:landmarks_to_matrix_calculator", "//mediapipe/tasks/cc/vision/gesture_recognizer/calculators:landmarks_to_matrix_calculator_cc_proto", + "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_classifier_graph_options_cc_proto", + "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_embedder_graph_options_cc_proto", "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:hand_gesture_recognizer_graph_options_cc_proto", "//mediapipe/tasks/cc/vision/hand_landmarker:hand_landmarks_detector_graph", - "//mediapipe/tasks/cc/vision/hand_landmarker/proto:hand_landmarker_graph_options_cc_proto", - "//mediapipe/tasks/cc/vision/hand_landmarker/proto:hand_landmarks_detector_graph_options_cc_proto", "//mediapipe/tasks/metadata:metadata_schema_cc", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", @@ -91,10 +100,15 @@ cc_library( "//mediapipe/framework/formats:classification_cc_proto", "//mediapipe/framework/formats:image", "//mediapipe/framework/formats:landmark_cc_proto", + "//mediapipe/framework/formats:rect_cc_proto", + "//mediapipe/framework/port:status", "//mediapipe/tasks/cc:common", "//mediapipe/tasks/cc/components/processors/proto:classifier_options_cc_proto", + "//mediapipe/tasks/cc/core:model_asset_bundle_resources", + "//mediapipe/tasks/cc/core:model_resources_cache", "//mediapipe/tasks/cc/core:model_task_graph", "//mediapipe/tasks/cc/core:utils", + "//mediapipe/tasks/cc/metadata/utils:zip_utils", "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_recognizer_graph_options_cc_proto", "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:hand_gesture_recognizer_graph_options_cc_proto", "//mediapipe/tasks/cc/vision/hand_detector:hand_detector_graph", @@ -123,9 +137,11 @@ cc_library( "//mediapipe/framework/formats:classification_cc_proto", "//mediapipe/framework/formats:image", "//mediapipe/framework/formats:landmark_cc_proto", + "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/tasks/cc:common", "//mediapipe/tasks/cc/components:image_preprocessing", "//mediapipe/tasks/cc/components/containers:gesture_recognition_result", + "//mediapipe/tasks/cc/components/processors:classifier_options", "//mediapipe/tasks/cc/components/processors/proto:classifier_options_cc_proto", "//mediapipe/tasks/cc/core:base_options", "//mediapipe/tasks/cc/core:base_task_api", @@ -134,8 +150,10 @@ cc_library( "//mediapipe/tasks/cc/core:utils", "//mediapipe/tasks/cc/core/proto:inference_subgraph_cc_proto", "//mediapipe/tasks/cc/vision/core:base_vision_task_api", + "//mediapipe/tasks/cc/vision/core:image_processing_options", "//mediapipe/tasks/cc/vision/core:running_mode", "//mediapipe/tasks/cc/vision/core:vision_task_api_factory", + "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_classifier_graph_options_cc_proto", "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_recognizer_graph_options_cc_proto", "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:hand_gesture_recognizer_graph_options_cc_proto", "//mediapipe/tasks/cc/vision/hand_detector/proto:hand_detector_graph_options_cc_proto", diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/BUILD b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/BUILD index a6de4f95..8c2c2e59 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/BUILD +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/BUILD @@ -69,6 +69,7 @@ cc_library( "//mediapipe/framework:calculator_framework", "//mediapipe/framework/formats:landmark_cc_proto", "//mediapipe/framework/formats:matrix", + "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/framework/port:ret_check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", @@ -86,8 +87,52 @@ cc_test( "//mediapipe/framework:calculator_runner", "//mediapipe/framework/formats:landmark_cc_proto", "//mediapipe/framework/formats:matrix", + "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/framework/port:gtest_main", "//mediapipe/framework/port:parse_text_proto", "@com_google_absl//absl/strings", ], ) + +mediapipe_proto_library( + name = "combined_prediction_calculator_proto", + srcs = ["combined_prediction_calculator.proto"], + visibility = ["//visibility:public"], + deps = [ + "//mediapipe/framework:calculator_options_proto", + "//mediapipe/framework:calculator_proto", + ], +) + +cc_library( + name = "combined_prediction_calculator", + srcs = ["combined_prediction_calculator.cc"], + deps = [ + ":combined_prediction_calculator_cc_proto", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework:collection", + "//mediapipe/framework/api2:node", + "//mediapipe/framework/api2:packet", + "//mediapipe/framework/api2:port", + "//mediapipe/framework/formats:classification_cc_proto", + "@com_google_absl//absl/container:btree", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings:str_format", + ], + alwayslink = 1, +) + +cc_test( + name = "combined_prediction_calculator_test", + srcs = ["combined_prediction_calculator_test.cc"], + deps = [ + ":combined_prediction_calculator", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework:calculator_runner", + "//mediapipe/framework/formats:classification_cc_proto", + "//mediapipe/framework/port:gtest", + "//mediapipe/framework/port:gtest_main", + "@com_google_absl//absl/strings", + ], +) diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator.cc b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator.cc new file mode 100644 index 00000000..cb95091d --- /dev/null +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator.cc @@ -0,0 +1,186 @@ +/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include + +#include "absl/container/btree_map.h" +#include "absl/memory/memory.h" +#include "absl/status/status.h" +#include "absl/strings/str_format.h" +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/api2/packet.h" +#include "mediapipe/framework/api2/port.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/collection.h" +#include "mediapipe/framework/formats/classification.pb.h" +#include "mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator.pb.h" + +namespace mediapipe { +namespace api2 { +namespace { + +constexpr char kPredictionTag[] = "PREDICTION"; + +Classification GetMaxScoringClassification( + const ClassificationList& classifications) { + Classification max_classification; + max_classification.set_score(0); + for (const auto& input : classifications.classification()) { + if (max_classification.score() < input.score()) { + max_classification = input; + } + } + return max_classification; +} + +float GetScoreThreshold( + const std::string& input_label, + const absl::btree_map& classwise_thresholds, + const std::string& background_label, const float default_threshold) { + float threshold = default_threshold; + auto it = classwise_thresholds.find(input_label); + if (it != classwise_thresholds.end()) { + threshold = it->second; + } + return threshold; +} + +std::unique_ptr GetWinningPrediction( + const ClassificationList& classification_list, + const absl::btree_map& classwise_thresholds, + const std::string& background_label, const float default_threshold) { + auto prediction_list = std::make_unique(); + if (classification_list.classification().empty()) { + return prediction_list; + } + Classification& prediction = *prediction_list->add_classification(); + auto argmax_prediction = GetMaxScoringClassification(classification_list); + float argmax_prediction_thresh = + GetScoreThreshold(argmax_prediction.label(), classwise_thresholds, + background_label, default_threshold); + if (argmax_prediction.score() >= argmax_prediction_thresh) { + prediction.set_label(argmax_prediction.label()); + prediction.set_score(argmax_prediction.score()); + } else { + for (const auto& input : classification_list.classification()) { + if (input.label() == background_label) { + prediction.set_label(input.label()); + prediction.set_score(input.score()); + break; + } + } + } + return prediction_list; +} + +} // namespace + +// This calculator accepts multiple ClassificationList input streams. Each +// ClassificationList should contain classifications with labels and +// corresponding softmax scores. The calculator computes the best prediction for +// each ClassificationList input stream via argmax and thresholding. Thresholds +// for all classes can be specified in the +// `CombinedPredictionCalculatorOptions`, along with a default global +// threshold. +// Please note that for this calculator to work as designed, the class names +// other than the background class in the ClassificationList objects must be +// different, but the background class name has to be the same. This background +// label name can be set via `background_label` in +// `CombinedPredictionCalculatorOptions`. +// The ClassificationList in the PREDICTION output stream contains the label of +// the winning class and corresponding softmax score. If none of the +// ClassificationList objects has a non-background winning class, the output +// contains the background class and score of the background class in the first +// ClassificationList. If multiple ClassificationList objects have a +// non-background winning class, the output contains the winning prediction from +// the ClassificationList with the highest priority. Priority is in decreasing +// order of input streams to the graph node using this calculator. +// Input: +// At least one stream with ClassificationList. +// Output: +// PREDICTION - A ClassificationList with the winning label as the only item. +// +// Usage example: +// node { +// calculator: "CombinedPredictionCalculator" +// input_stream: "classification_list_0" +// input_stream: "classification_list_1" +// output_stream: "PREDICTION:prediction" +// options { +// [mediapipe.CombinedPredictionCalculatorOptions.ext] { +// class { +// label: "A" +// score_threshold: 0.7 +// } +// default_global_threshold: 0.1 +// background_label: "B" +// } +// } +// } + +class CombinedPredictionCalculator : public Node { + public: + static constexpr Input::Multiple kClassificationListIn{ + ""}; + static constexpr Output kPredictionOut{"PREDICTION"}; + MEDIAPIPE_NODE_CONTRACT(kClassificationListIn, kPredictionOut); + + absl::Status Open(CalculatorContext* cc) override { + options_ = cc->Options(); + for (const auto& input : options_.class_()) { + classwise_thresholds_[input.label()] = input.score_threshold(); + } + classwise_thresholds_[options_.background_label()] = 0; + return absl::OkStatus(); + } + + absl::Status Process(CalculatorContext* cc) override { + // After loop, if have winning prediction return. Otherwise empty packet. + std::unique_ptr first_winning_prediction = nullptr; + auto collection = kClassificationListIn(cc); + for (const auto& input : collection) { + if (input.IsEmpty() || input.Get().classification_size() == 0) { + continue; + } + auto prediction = GetWinningPrediction( + input.Get(), classwise_thresholds_, options_.background_label(), + options_.default_global_threshold()); + if (prediction->classification(0).label() != + options_.background_label()) { + kPredictionOut(cc).Send(std::move(prediction)); + return absl::OkStatus(); + } + if (first_winning_prediction == nullptr) { + first_winning_prediction = std::move(prediction); + } + } + if (first_winning_prediction != nullptr) { + kPredictionOut(cc).Send(std::move(first_winning_prediction)); + } + return absl::OkStatus(); + } + + private: + CombinedPredictionCalculatorOptions options_; + absl::btree_map classwise_thresholds_; +}; + +MEDIAPIPE_REGISTER_NODE(CombinedPredictionCalculator); + +} // namespace api2 +} // namespace mediapipe diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator.proto b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator.proto new file mode 100644 index 00000000..730e7dd7 --- /dev/null +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator.proto @@ -0,0 +1,41 @@ +/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +syntax = "proto2"; + +package mediapipe; + +import "mediapipe/framework/calculator.proto"; + +message CombinedPredictionCalculatorOptions { + extend mediapipe.CalculatorOptions { + optional CombinedPredictionCalculatorOptions ext = 483738635; + } + + message Class { + optional string label = 1; + optional float score_threshold = 2; + } + + // List of classes with score thresholds. + repeated Class class = 1; + + // Default score threshold applied to a label. + optional float default_global_threshold = 2 [default = 0]; + + // Name of the background class whose input scores will be ignored while + // thresholding. + optional string background_label = 3; +} diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator_test.cc b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator_test.cc new file mode 100644 index 00000000..ecf49795 --- /dev/null +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator_test.cc @@ -0,0 +1,315 @@ +/* Copyright 2022 The MediaPipe Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include + +#include "absl/strings/string_view.h" +#include "absl/strings/substitute.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/calculator_runner.h" +#include "mediapipe/framework/formats/classification.pb.h" +#include "mediapipe/framework/port/status_matchers.h" + +namespace mediapipe { + +namespace { + +constexpr char kPredictionTag[] = "PREDICTION"; + +std::unique_ptr BuildNodeRunnerWithOptions( + float drama_thresh, float llama_thresh, float bazinga_thresh, + float joy_thresh, float peace_thresh) { + constexpr absl::string_view kCalculatorProto = R"pb( + calculator: "CombinedPredictionCalculator" + input_stream: "custom_softmax_scores" + input_stream: "canned_softmax_scores" + output_stream: "PREDICTION:prediction" + options { + [mediapipe.CombinedPredictionCalculatorOptions.ext] { + class { label: "CustomDrama" score_threshold: $0 } + class { label: "CustomLlama" score_threshold: $1 } + class { label: "CannedBazinga" score_threshold: $2 } + class { label: "CannedJoy" score_threshold: $3 } + class { label: "CannedPeace" score_threshold: $4 } + background_label: "Negative" + } + } + )pb"; + auto runner = std::make_unique( + absl::Substitute(kCalculatorProto, drama_thresh, llama_thresh, + bazinga_thresh, joy_thresh, peace_thresh)); + return runner; +} + +std::unique_ptr BuildCustomScoreInput( + const float negative_score, const float drama_score, + const float llama_score) { + auto custom_scores = std::make_unique(); + auto custom_negative = custom_scores->add_classification(); + custom_negative->set_label("Negative"); + custom_negative->set_score(negative_score); + auto drama = custom_scores->add_classification(); + drama->set_label("CustomDrama"); + drama->set_score(drama_score); + auto llama = custom_scores->add_classification(); + llama->set_label("CustomLlama"); + llama->set_score(llama_score); + return custom_scores; +} + +std::unique_ptr BuildCannedScoreInput( + const float negative_score, const float bazinga_score, + const float joy_score, const float peace_score) { + auto canned_scores = std::make_unique(); + auto canned_negative = canned_scores->add_classification(); + canned_negative->set_label("Negative"); + canned_negative->set_score(negative_score); + auto bazinga = canned_scores->add_classification(); + bazinga->set_label("CannedBazinga"); + bazinga->set_score(bazinga_score); + auto joy = canned_scores->add_classification(); + joy->set_label("CannedJoy"); + joy->set_score(joy_score); + auto peace = canned_scores->add_classification(); + peace->set_label("CannedPeace"); + peace->set_score(peace_score); + return canned_scores; +} + +TEST(CombinedPredictionCalculatorPacketTest, + CustomEmpty_CannedEmpty_ResultIsEmpty) { + auto runner = BuildNodeRunnerWithOptions( + /*drama_thresh=*/0.0, /*llama_thresh=*/0.0, /*bazinga_thresh=*/0.0, + /*joy_thresh=*/0.0, /*peace_thresh=*/0.0); + MP_ASSERT_OK(runner->Run()) << "Calculator execution failed."; + EXPECT_THAT(runner->Outputs().Tag("PREDICTION").packets, testing::IsEmpty()); +} + +TEST(CombinedPredictionCalculatorPacketTest, + CustomEmpty_CannedNotEmpty_ResultIsCanned) { + auto runner = BuildNodeRunnerWithOptions( + /*drama_thresh=*/0.0, /*llama_thresh=*/0.0, /*bazinga_thresh=*/0.9, + /*joy_thresh=*/0.5, /*peace_thresh=*/0.8); + auto canned_scores = BuildCannedScoreInput( + /*negative_score=*/0.1, + /*bazinga_score=*/0.1, /*joy_score=*/0.6, /*peace_score=*/0.2); + runner->MutableInputs()->Index(1).packets.push_back( + Adopt(canned_scores.release()).At(Timestamp(1))); + MP_ASSERT_OK(runner->Run()) << "Calculator execution failed."; + + auto output_prediction_packets = + runner->Outputs().Tag(kPredictionTag).packets; + ASSERT_EQ(output_prediction_packets.size(), 1); + Classification output_prediction = + output_prediction_packets[0].Get().classification(0); + + EXPECT_EQ(output_prediction.label(), "CannedJoy"); + EXPECT_NEAR(output_prediction.score(), 0.6, 1e-4); +} + +TEST(CombinedPredictionCalculatorPacketTest, + CustomNotEmpty_CannedEmpty_ResultIsCustom) { + auto runner = BuildNodeRunnerWithOptions( + /*drama_thresh=*/0.3, /*llama_thresh=*/0.5, /*bazinga_thresh=*/0.0, + /*joy_thresh=*/0.0, /*peace_thresh=*/0.0); + auto custom_scores = + BuildCustomScoreInput(/*negative_score=*/0.1, + /*drama_score=*/0.2, /*llama_score=*/0.7); + runner->MutableInputs()->Index(0).packets.push_back( + Adopt(custom_scores.release()).At(Timestamp(1))); + MP_ASSERT_OK(runner->Run()) << "Calculator execution failed."; + + auto output_prediction_packets = + runner->Outputs().Tag(kPredictionTag).packets; + ASSERT_EQ(output_prediction_packets.size(), 1); + Classification output_prediction = + output_prediction_packets[0].Get().classification(0); + + EXPECT_EQ(output_prediction.label(), "CustomLlama"); + EXPECT_NEAR(output_prediction.score(), 0.7, 1e-4); +} + +struct CombinedPredictionCalculatorTestCase { + std::string test_name; + float custom_negative_score; + float drama_score; + float llama_score; + float drama_thresh; + float llama_thresh; + float canned_negative_score; + float bazinga_score; + float joy_score; + float peace_score; + float bazinga_thresh; + float joy_thresh; + float peace_thresh; + std::string max_scoring_label; + float max_score; +}; + +using CombinedPredictionCalculatorTest = + testing::TestWithParam; + +TEST_P(CombinedPredictionCalculatorTest, OutputsCorrectResult) { + const CombinedPredictionCalculatorTestCase& test_case = GetParam(); + + auto runner = BuildNodeRunnerWithOptions( + test_case.drama_thresh, test_case.llama_thresh, test_case.bazinga_thresh, + test_case.joy_thresh, test_case.peace_thresh); + + auto custom_scores = + BuildCustomScoreInput(test_case.custom_negative_score, + test_case.drama_score, test_case.llama_score); + + runner->MutableInputs()->Index(0).packets.push_back( + Adopt(custom_scores.release()).At(Timestamp(1))); + + auto canned_scores = BuildCannedScoreInput( + test_case.canned_negative_score, test_case.bazinga_score, + test_case.joy_score, test_case.peace_score); + runner->MutableInputs()->Index(1).packets.push_back( + Adopt(canned_scores.release()).At(Timestamp(1))); + + MP_ASSERT_OK(runner->Run()) << "Calculator execution failed."; + + auto output_prediction_packets = + runner->Outputs().Tag(kPredictionTag).packets; + ASSERT_EQ(output_prediction_packets.size(), 1); + Classification output_prediction = + output_prediction_packets[0].Get().classification(0); + + EXPECT_EQ(output_prediction.label(), test_case.max_scoring_label); + EXPECT_NEAR(output_prediction.score(), test_case.max_score, 1e-4); +} + +INSTANTIATE_TEST_CASE_P( + CombinedPredictionCalculatorTests, CombinedPredictionCalculatorTest, + testing::ValuesIn({ + { + .test_name = "TestCustomDramaWinnnerWith_HighCanned_Thresh", + .custom_negative_score = 0.1, + .drama_score = 0.5, + .llama_score = 0.3, + .drama_thresh = 0.25, + .llama_thresh = 0.7, + .canned_negative_score = 0.1, + .bazinga_score = 0.3, + .joy_score = 0.3, + .peace_score = 0.3, + .bazinga_thresh = 0.7, + .joy_thresh = 0.7, + .peace_thresh = 0.7, + .max_scoring_label = "CustomDrama", + .max_score = 0.5, + }, + { + .test_name = "TestCannedWinnerWith_HighCustom_ZeroCanned_Thresh", + .custom_negative_score = 0.1, + .drama_score = 0.3, + .llama_score = 0.6, + .drama_thresh = 0.4, + .llama_thresh = 0.8, + .canned_negative_score = 0.1, + .bazinga_score = 0.4, + .joy_score = 0.3, + .peace_score = 0.2, + .bazinga_thresh = 0.0, + .joy_thresh = 0.0, + .peace_thresh = 0.0, + .max_scoring_label = "CannedBazinga", + .max_score = 0.4, + }, + { + .test_name = "TestNegativeWinnerWith_LowCustom_HighCanned_Thresh", + .custom_negative_score = 0.5, + .drama_score = 0.1, + .llama_score = 0.4, + .drama_thresh = 0.1, + .llama_thresh = 0.05, + .canned_negative_score = 0.1, + .bazinga_score = 0.3, + .joy_score = 0.3, + .peace_score = 0.3, + .bazinga_thresh = 0.7, + .joy_thresh = 0.7, + .peace_thresh = 0.7, + .max_scoring_label = "Negative", + .max_score = 0.5, + }, + { + .test_name = "TestNegativeWinnerWith_HighCustom_HighCanned_Thresh", + .custom_negative_score = 0.8, + .drama_score = 0.1, + .llama_score = 0.1, + .drama_thresh = 0.25, + .llama_thresh = 0.7, + .canned_negative_score = 0.1, + .bazinga_score = 0.3, + .joy_score = 0.3, + .peace_score = 0.3, + .bazinga_thresh = 0.7, + .joy_thresh = 0.7, + .peace_thresh = 0.7, + .max_scoring_label = "Negative", + .max_score = 0.8, + }, + { + .test_name = "TestNegativeWinnerWith_HighCustom_HighCannedThresh2", + .custom_negative_score = 0.1, + .drama_score = 0.2, + .llama_score = 0.7, + .drama_thresh = 1.1, + .llama_thresh = 1.1, + .canned_negative_score = 0.1, + .bazinga_score = 0.3, + .joy_score = 0.3, + .peace_score = 0.3, + .bazinga_thresh = 0.7, + .joy_thresh = 0.7, + .peace_thresh = 0.7, + .max_scoring_label = "Negative", + .max_score = 0.1, + }, + { + .test_name = "TestNegativeWinnerWith_HighCustom_HighCanned_Thresh3", + .custom_negative_score = 0.1, + .drama_score = 0.3, + .llama_score = 0.6, + .drama_thresh = 0.4, + .llama_thresh = 0.8, + .canned_negative_score = 0.3, + .bazinga_score = 0.2, + .joy_score = 0.3, + .peace_score = 0.2, + .bazinga_thresh = 0.5, + .joy_thresh = 0.5, + .peace_thresh = 0.5, + .max_scoring_label = "Negative", + .max_score = 0.1, + }, + }), + [](const testing::TestParamInfo< + CombinedPredictionCalculatorTest::ParamType>& info) { + return info.param.test_name; + }); + +} // namespace + +} // namespace mediapipe diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator.cc b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator.cc index b70689ea..277bb170 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator.cc +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include +#include #include #include #include @@ -26,6 +27,7 @@ limitations under the License. #include "mediapipe/framework/calculator_framework.h" #include "mediapipe/framework/formats/landmark.pb.h" #include "mediapipe/framework/formats/matrix.h" +#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/port/ret_check.h" #include "mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator.pb.h" @@ -38,6 +40,7 @@ namespace { constexpr char kLandmarksTag[] = "LANDMARKS"; constexpr char kWorldLandmarksTag[] = "WORLD_LANDMARKS"; constexpr char kImageSizeTag[] = "IMAGE_SIZE"; +constexpr char kNormRectTag[] = "NORM_RECT"; constexpr char kLandmarksMatrixTag[] = "LANDMARKS_MATRIX"; constexpr int kFeaturesPerLandmark = 3; @@ -62,6 +65,25 @@ absl::StatusOr NormalizeLandmarkAspectRatio( return normalized_landmarks; } +template +absl::StatusOr RotateLandmarks(const LandmarkListT& landmarks, + float rotation) { + float cos = std::cos(rotation); + // Negate because Y-axis points down and not up. + float sin = std::sin(-rotation); + LandmarkListT rotated_landmarks; + for (int i = 0; i < landmarks.landmark_size(); ++i) { + const auto& old_landmark = landmarks.landmark(i); + float x = old_landmark.x() - 0.5; + float y = old_landmark.y() - 0.5; + auto* new_landmark = rotated_landmarks.add_landmark(); + new_landmark->set_x(x * cos - y * sin + 0.5); + new_landmark->set_y(y * cos + x * sin + 0.5); + new_landmark->set_z(old_landmark.z()); + } + return rotated_landmarks; +} + template absl::StatusOr NormalizeObject(const LandmarkListT& landmarks, int origin_offset) { @@ -134,6 +156,13 @@ absl::Status ProcessLandmarks(LandmarkListT landmarks, CalculatorContext* cc) { NormalizeLandmarkAspectRatio(landmarks, width, height)); } + if (cc->Inputs().HasTag(kNormRectTag)) { + RET_CHECK(!cc->Inputs().Tag(kNormRectTag).IsEmpty()); + const auto rotation = + cc->Inputs().Tag(kNormRectTag).Get().rotation(); + ASSIGN_OR_RETURN(landmarks, RotateLandmarks(landmarks, rotation)); + } + const auto& options = cc->Options(); if (options.object_normalization()) { ASSIGN_OR_RETURN( @@ -163,6 +192,8 @@ absl::Status ProcessLandmarks(LandmarkListT landmarks, CalculatorContext* cc) { // WORLD_LANDMARKS - World 3d landmarks of one object. Use *either* // LANDMARKS or WORLD_LANDMARKS. // IMAGE_SIZE - (width, height) of the image +// NORM_RECT - Optional NormalizedRect object whose 'rotation' field is used +// to rotate the landmarks. // Output: // LANDMARKS_MATRIX - Matrix for the landmarks. // @@ -185,6 +216,7 @@ class LandmarksToMatrixCalculator : public CalculatorBase { cc->Inputs().Tag(kLandmarksTag).Set().Optional(); cc->Inputs().Tag(kWorldLandmarksTag).Set().Optional(); cc->Inputs().Tag(kImageSizeTag).Set>().Optional(); + cc->Inputs().Tag(kNormRectTag).Set().Optional(); cc->Outputs().Tag(kLandmarksMatrixTag).Set(); return absl::OkStatus(); } diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator_test.cc b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator_test.cc index 8a68d8da..fe6f1162 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator_test.cc +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator_test.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include #include @@ -23,6 +24,7 @@ limitations under the License. #include "mediapipe/framework/calculator_runner.h" #include "mediapipe/framework/formats/landmark.pb.h" #include "mediapipe/framework/formats/matrix.h" +#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/port/gtest.h" #include "mediapipe/framework/port/parse_text_proto.h" #include "mediapipe/framework/port/status_matchers.h" @@ -35,6 +37,7 @@ constexpr char kLandmarksTag[] = "LANDMARKS"; constexpr char kWorldLandmarksTag[] = "WORLD_LANDMARKS"; constexpr char kImageSizeTag[] = "IMAGE_SIZE"; constexpr char kLandmarksMatrixTag[] = "LANDMARKS_MATRIX"; +constexpr char kNormRectTag[] = "NORM_RECT"; template LandmarkListT BuildPseudoLandmarks(int num_landmarks, int offset = 0) { @@ -54,6 +57,7 @@ struct Landmarks2dToMatrixCalculatorTestCase { int object_normalization_origin_offset = -1; float expected_cell_0_2; float expected_cell_1_5; + float rotation; }; using Landmarks2dToMatrixCalculatorTest = @@ -68,6 +72,7 @@ TEST_P(Landmarks2dToMatrixCalculatorTest, OutputsCorrectResult) { calculator: "LandmarksToMatrixCalculator" input_stream: "LANDMARKS:landmarks" input_stream: "IMAGE_SIZE:image_size" + input_stream: "NORM_RECT:norm_rect" output_stream: "LANDMARKS_MATRIX:landmarks_matrix" options { [mediapipe.LandmarksToMatrixCalculatorOptions.ext] { @@ -91,6 +96,11 @@ TEST_P(Landmarks2dToMatrixCalculatorTest, OutputsCorrectResult) { runner.MutableInputs() ->Tag(kImageSizeTag) .packets.push_back(Adopt(image_size.release()).At(Timestamp(0))); + auto norm_rect = std::make_unique(); + norm_rect->set_rotation(test_case.rotation); + runner.MutableInputs() + ->Tag(kNormRectTag) + .packets.push_back(Adopt(norm_rect.release()).At(Timestamp(0))); MP_ASSERT_OK(runner.Run()) << "Calculator execution failed."; @@ -109,12 +119,20 @@ INSTANTIATE_TEST_CASE_P( .base_offset = 0, .object_normalization_origin_offset = 0, .expected_cell_0_2 = 0.1f, - .expected_cell_1_5 = 0.1875f}, + .expected_cell_1_5 = 0.1875f, + .rotation = 0}, {.test_name = "TestWithOffset21", .base_offset = 21, .object_normalization_origin_offset = 0, .expected_cell_0_2 = 0.1f, - .expected_cell_1_5 = 0.1875f}}), + .expected_cell_1_5 = 0.1875f, + .rotation = 0}, + {.test_name = "TestWithRotation", + .base_offset = 0, + .object_normalization_origin_offset = 0, + .expected_cell_0_2 = 0.075f, + .expected_cell_1_5 = -0.25f, + .rotation = M_PI / 2.0}}), [](const testing::TestParamInfo< Landmarks2dToMatrixCalculatorTest::ParamType>& info) { return info.param.test_name; @@ -126,6 +144,7 @@ struct LandmarksWorld3dToMatrixCalculatorTestCase { int object_normalization_origin_offset = -1; float expected_cell_0_2; float expected_cell_1_5; + float rotation; }; using LandmarksWorld3dToMatrixCalculatorTest = @@ -140,6 +159,7 @@ TEST_P(LandmarksWorld3dToMatrixCalculatorTest, OutputsCorrectResult) { calculator: "LandmarksToMatrixCalculator" input_stream: "WORLD_LANDMARKS:landmarks" input_stream: "IMAGE_SIZE:image_size" + input_stream: "NORM_RECT:norm_rect" output_stream: "LANDMARKS_MATRIX:landmarks_matrix" options { [mediapipe.LandmarksToMatrixCalculatorOptions.ext] { @@ -162,6 +182,11 @@ TEST_P(LandmarksWorld3dToMatrixCalculatorTest, OutputsCorrectResult) { runner.MutableInputs() ->Tag(kImageSizeTag) .packets.push_back(Adopt(image_size.release()).At(Timestamp(0))); + auto norm_rect = std::make_unique(); + norm_rect->set_rotation(test_case.rotation); + runner.MutableInputs() + ->Tag(kNormRectTag) + .packets.push_back(Adopt(norm_rect.release()).At(Timestamp(0))); MP_ASSERT_OK(runner.Run()) << "Calculator execution failed."; @@ -180,17 +205,26 @@ INSTANTIATE_TEST_CASE_P( .base_offset = 0, .object_normalization_origin_offset = 0, .expected_cell_0_2 = 0.1f, - .expected_cell_1_5 = 0.25}, + .expected_cell_1_5 = 0.25, + .rotation = 0}, {.test_name = "TestWithOffset21", .base_offset = 21, .object_normalization_origin_offset = 0, .expected_cell_0_2 = 0.1f, - .expected_cell_1_5 = 0.25}, + .expected_cell_1_5 = 0.25, + .rotation = 0}, {.test_name = "NoObjectNormalization", .base_offset = 0, .object_normalization_origin_offset = -1, .expected_cell_0_2 = 0.021f, - .expected_cell_1_5 = 0.052f}}), + .expected_cell_1_5 = 0.052f, + .rotation = 0}, + {.test_name = "TestWithRotation", + .base_offset = 0, + .object_normalization_origin_offset = 0, + .expected_cell_0_2 = 0.1f, + .expected_cell_1_5 = -0.25f, + .rotation = M_PI / 2.0}}), [](const testing::TestParamInfo< LandmarksWorld3dToMatrixCalculatorTest::ParamType>& info) { return info.param.test_name; diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer.cc b/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer.cc index e0d1473c..38cb5169 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer.cc +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include "absl/memory/memory.h" @@ -27,6 +28,7 @@ limitations under the License. #include "mediapipe/framework/formats/classification.pb.h" #include "mediapipe/framework/formats/image.h" #include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/packet.h" #include "mediapipe/tasks/cc/common.h" #include "mediapipe/tasks/cc/components/image_preprocessing.h" @@ -37,7 +39,9 @@ limitations under the License. #include "mediapipe/tasks/cc/core/task_runner.h" #include "mediapipe/tasks/cc/core/utils.h" #include "mediapipe/tasks/cc/vision/core/base_vision_task_api.h" +#include "mediapipe/tasks/cc/vision/core/image_processing_options.h" #include "mediapipe/tasks/cc/vision/core/vision_task_api_factory.h" +#include "mediapipe/tasks/cc/vision/gesture_recognizer/proto/gesture_classifier_graph_options.pb.h" #include "mediapipe/tasks/cc/vision/gesture_recognizer/proto/gesture_recognizer_graph_options.pb.h" #include "mediapipe/tasks/cc/vision/gesture_recognizer/proto/hand_gesture_recognizer_graph_options.pb.h" #include "mediapipe/tasks/cc/vision/hand_detector/proto/hand_detector_graph_options.pb.h" @@ -62,6 +66,8 @@ constexpr char kHandGestureSubgraphTypeName[] = constexpr char kImageTag[] = "IMAGE"; constexpr char kImageInStreamName[] = "image_in"; constexpr char kImageOutStreamName[] = "image_out"; +constexpr char kNormRectTag[] = "NORM_RECT"; +constexpr char kNormRectStreamName[] = "norm_rect_in"; constexpr char kHandGesturesTag[] = "HAND_GESTURES"; constexpr char kHandGesturesStreamName[] = "hand_gestures"; constexpr char kHandednessTag[] = "HANDEDNESS"; @@ -83,6 +89,7 @@ CalculatorGraphConfig CreateGraphConfig( auto& subgraph = graph.AddNode(kHandGestureSubgraphTypeName); subgraph.GetOptions().Swap(options.get()); graph.In(kImageTag).SetName(kImageInStreamName); + graph.In(kNormRectTag).SetName(kNormRectStreamName); subgraph.Out(kHandGesturesTag).SetName(kHandGesturesStreamName) >> graph.Out(kHandGesturesTag); subgraph.Out(kHandednessTag).SetName(kHandednessStreamName) >> @@ -93,10 +100,11 @@ CalculatorGraphConfig CreateGraphConfig( graph.Out(kHandWorldLandmarksTag); subgraph.Out(kImageTag).SetName(kImageOutStreamName) >> graph.Out(kImageTag); if (enable_flow_limiting) { - return tasks::core::AddFlowLimiterCalculator(graph, subgraph, {kImageTag}, - kHandGesturesTag); + return tasks::core::AddFlowLimiterCalculator( + graph, subgraph, {kImageTag, kNormRectTag}, kHandGesturesTag); } graph.In(kImageTag) >> subgraph.In(kImageTag); + graph.In(kNormRectTag) >> subgraph.In(kNormRectTag); return graph.GetConfig(); } @@ -105,59 +113,50 @@ CalculatorGraphConfig CreateGraphConfig( std::unique_ptr ConvertGestureRecognizerGraphOptionsProto(GestureRecognizerOptions* options) { auto options_proto = std::make_unique(); + auto base_options_proto = std::make_unique( + tasks::core::ConvertBaseOptionsToProto(&(options->base_options))); + options_proto->mutable_base_options()->Swap(base_options_proto.get()); + options_proto->mutable_base_options()->set_use_stream_mode( + options->running_mode != core::RunningMode::IMAGE); - bool use_stream_mode = options->running_mode != core::RunningMode::IMAGE; - - // TODO remove these workarounds for base options of subgraphs. // Configure hand detector options. - auto base_options_proto_for_hand_detector = - std::make_unique( - tasks::core::ConvertBaseOptionsToProto( - &(options->base_options_for_hand_detector))); - base_options_proto_for_hand_detector->set_use_stream_mode(use_stream_mode); auto* hand_detector_graph_options = options_proto->mutable_hand_landmarker_graph_options() ->mutable_hand_detector_graph_options(); - hand_detector_graph_options->mutable_base_options()->Swap( - base_options_proto_for_hand_detector.get()); hand_detector_graph_options->set_num_hands(options->num_hands); hand_detector_graph_options->set_min_detection_confidence( options->min_hand_detection_confidence); // Configure hand landmark detector options. - auto base_options_proto_for_hand_landmarker = - std::make_unique( - tasks::core::ConvertBaseOptionsToProto( - &(options->base_options_for_hand_landmarker))); - base_options_proto_for_hand_landmarker->set_use_stream_mode(use_stream_mode); - auto* hand_landmarks_detector_graph_options = - options_proto->mutable_hand_landmarker_graph_options() - ->mutable_hand_landmarks_detector_graph_options(); - hand_landmarks_detector_graph_options->mutable_base_options()->Swap( - base_options_proto_for_hand_landmarker.get()); - hand_landmarks_detector_graph_options->set_min_detection_confidence( - options->min_hand_presence_confidence); - auto* hand_landmarker_graph_options = options_proto->mutable_hand_landmarker_graph_options(); hand_landmarker_graph_options->set_min_tracking_confidence( options->min_tracking_confidence); + auto* hand_landmarks_detector_graph_options = + hand_landmarker_graph_options + ->mutable_hand_landmarks_detector_graph_options(); + hand_landmarks_detector_graph_options->set_min_detection_confidence( + options->min_hand_presence_confidence); // Configure hand gesture recognizer options. - auto base_options_proto_for_gesture_recognizer = - std::make_unique( - tasks::core::ConvertBaseOptionsToProto( - &(options->base_options_for_gesture_recognizer))); - base_options_proto_for_gesture_recognizer->set_use_stream_mode( - use_stream_mode); auto* hand_gesture_recognizer_graph_options = options_proto->mutable_hand_gesture_recognizer_graph_options(); - hand_gesture_recognizer_graph_options->mutable_base_options()->Swap( - base_options_proto_for_gesture_recognizer.get()); - if (options->min_gesture_confidence >= 0) { - hand_gesture_recognizer_graph_options->mutable_classifier_options() - ->set_score_threshold(options->min_gesture_confidence); - } + auto canned_gestures_classifier_options_proto = + std::make_unique( + components::processors::ConvertClassifierOptionsToProto( + &(options->canned_gestures_classifier_options))); + hand_gesture_recognizer_graph_options + ->mutable_canned_gesture_classifier_graph_options() + ->mutable_classifier_options() + ->Swap(canned_gestures_classifier_options_proto.get()); + auto custom_gestures_classifier_options_proto = + std::make_unique( + components::processors::ConvertClassifierOptionsToProto( + &(options->canned_gestures_classifier_options))); + hand_gesture_recognizer_graph_options + ->mutable_custom_gesture_classifier_graph_options() + ->mutable_classifier_options() + ->Swap(canned_gestures_classifier_options_proto.get()); return options_proto; } @@ -216,16 +215,23 @@ absl::StatusOr> GestureRecognizer::Create( } absl::StatusOr GestureRecognizer::Recognize( - mediapipe::Image image) { + mediapipe::Image image, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, "GPU input images are currently not supported.", MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - ASSIGN_OR_RETURN(auto output_packets, - ProcessImageData({{kImageInStreamName, - MakePacket(std::move(image))}})); + ASSIGN_OR_RETURN( + NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false)); + ASSIGN_OR_RETURN( + auto output_packets, + ProcessImageData( + {{kImageInStreamName, MakePacket(std::move(image))}, + {kNormRectStreamName, + MakePacket(std::move(norm_rect))}})); if (output_packets[kHandGesturesStreamName].IsEmpty()) { return {{{}, {}, {}, {}}}; } @@ -245,18 +251,25 @@ absl::StatusOr GestureRecognizer::Recognize( } absl::StatusOr GestureRecognizer::RecognizeForVideo( - mediapipe::Image image, int64 timestamp_ms) { + mediapipe::Image image, int64 timestamp_ms, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, absl::StrCat("GPU input images are currently not supported."), MediaPipeTasksStatus::kRunnerUnexpectedInputError); } + ASSIGN_OR_RETURN( + NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false)); ASSIGN_OR_RETURN( auto output_packets, ProcessVideoData( {{kImageInStreamName, MakePacket(std::move(image)) + .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}, + {kNormRectStreamName, + MakePacket(std::move(norm_rect)) .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}})); if (output_packets[kHandGesturesStreamName].IsEmpty()) { return {{{}, {}, {}, {}}}; @@ -276,17 +289,24 @@ absl::StatusOr GestureRecognizer::RecognizeForVideo( }; } -absl::Status GestureRecognizer::RecognizeAsync(mediapipe::Image image, - int64 timestamp_ms) { +absl::Status GestureRecognizer::RecognizeAsync( + mediapipe::Image image, int64 timestamp_ms, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, absl::StrCat("GPU input images are currently not supported."), MediaPipeTasksStatus::kRunnerUnexpectedInputError); } + ASSIGN_OR_RETURN( + NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false)); return SendLiveStreamData( {{kImageInStreamName, MakePacket(std::move(image)) + .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}, + {kNormRectStreamName, + MakePacket(std::move(norm_rect)) .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}}); } diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer.h b/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer.h index 53b824e2..3f3d7acf 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer.h +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer.h @@ -17,14 +17,18 @@ limitations under the License. #define MEDIAPIPE_TASKS_CC_VISION_GESTURE_RECOGNIZRER_GESTURE_RECOGNIZER_H_ #include +#include +#include #include "absl/status/statusor.h" #include "mediapipe/framework/formats/classification.pb.h" #include "mediapipe/framework/formats/image.h" #include "mediapipe/framework/formats/landmark.pb.h" #include "mediapipe/tasks/cc/components/containers/gesture_recognition_result.h" +#include "mediapipe/tasks/cc/components/processors/classifier_options.h" #include "mediapipe/tasks/cc/core/base_options.h" #include "mediapipe/tasks/cc/vision/core/base_vision_task_api.h" +#include "mediapipe/tasks/cc/vision/core/image_processing_options.h" #include "mediapipe/tasks/cc/vision/core/running_mode.h" namespace mediapipe { @@ -37,12 +41,6 @@ struct GestureRecognizerOptions { // model file with metadata, accelerator options, op resolver, etc. tasks::core::BaseOptions base_options; - // TODO: remove these. Temporary solutions before bundle asset is - // ready. - tasks::core::BaseOptions base_options_for_hand_landmarker; - tasks::core::BaseOptions base_options_for_hand_detector; - tasks::core::BaseOptions base_options_for_gesture_recognizer; - // The running mode of the task. Default to the image mode. // GestureRecognizer has three running modes: // 1) The image mode for recognizing hand gestures on single image inputs. @@ -57,7 +55,7 @@ struct GestureRecognizerOptions { int num_hands = 1; // The minimum confidence score for the hand detection to be considered - // successfully. + // successful. float min_hand_detection_confidence = 0.5; // The minimum confidence score of hand presence score in the hand landmark @@ -65,15 +63,20 @@ struct GestureRecognizerOptions { float min_hand_presence_confidence = 0.5; // The minimum confidence score for the hand tracking to be considered - // successfully. + // successful. float min_tracking_confidence = 0.5; - // The minimum confidence score for the gestures to be considered - // successfully. If < 0, the gesture confidence thresholds in the model - // metadata are used. - // TODO Note this option is subject to change, after scoring - // merging calculator is implemented. - float min_gesture_confidence = -1; + // TODO Note this option is subject to change. + // Options for configuring the canned gestures classifier, such as score + // threshold, allow list and deny list of gestures. The categories for canned + // gesture classifiers are: ["None", "Closed_Fist", "Open_Palm", + // "Pointing_Up", "Thumb_Down", "Thumb_Up", "Victory", "ILoveYou"] + components::processors::ClassifierOptions canned_gestures_classifier_options; + + // TODO Note this option is subject to change. + // Options for configuring the custom gestures classifier, such as score + // threshold, allow list and deny list of gestures. + components::processors::ClassifierOptions custom_gestures_classifier_options; // The user-defined result callback for processing live stream data. // The result callback should only be specified when the running mode is set @@ -93,6 +96,13 @@ struct GestureRecognizerOptions { // Inputs: // Image // - The image that gesture recognition runs on. +// std::optional +// - If provided, can be used to specify the rotation to apply to the image +// before performing gesture recognition, by setting its 'rotation' field +// in radians (e.g. 'M_PI / 2' for a 90° anti-clockwise rotation). Note +// that specifying a region-of-interest using the 'x_center', 'y_center', +// 'width' and 'height' fields is NOT supported and will result in an +// invalid argument error being returned. // Outputs: // GestureRecognitionResult // - The hand gesture recognition results. @@ -120,24 +130,37 @@ class GestureRecognizer : tasks::vision::core::BaseVisionTaskApi { // Only use this method when the GestureRecognizer is created with the image // running mode. // - // image - mediapipe::Image - // Image to perform hand gesture recognition on. + // The optional 'image_processing_options' parameter can be used to specify + // the rotation to apply to the image before performing recognition, by + // setting its 'rotation_degrees' field. Note that specifying a + // region-of-interest using the 'region_of_interest' field is NOT supported + // and will result in an invalid argument error being returned. // // The image can be of any size with format RGB or RGBA. // TODO: Describes how the input image will be preprocessed // after the yuv support is implemented. absl::StatusOr Recognize( - Image image); + Image image, + std::optional image_processing_options = + std::nullopt); // Performs gesture recognition on the provided video frame. // Only use this method when the GestureRecognizer is created with the video // running mode. // + // The optional 'image_processing_options' parameter can be used to specify + // the rotation to apply to the image before performing recognition, by + // setting its 'rotation_degrees' field. Note that specifying a + // region-of-interest using the 'region_of_interest' field is NOT supported + // and will result in an invalid argument error being returned. + // // The image can be of any size with format RGB or RGBA. It's required to // provide the video frame's timestamp (in milliseconds). The input timestamps // must be monotonically increasing. absl::StatusOr - RecognizeForVideo(Image image, int64 timestamp_ms); + RecognizeForVideo(Image image, int64 timestamp_ms, + std::optional + image_processing_options = std::nullopt); // Sends live image data to perform gesture recognition, and the results will // be available via the "result_callback" provided in the @@ -149,6 +172,12 @@ class GestureRecognizer : tasks::vision::core::BaseVisionTaskApi { // sent to the gesture recognizer. The input timestamps must be monotonically // increasing. // + // The optional 'image_processing_options' parameter can be used to specify + // the rotation to apply to the image before performing recognition, by + // setting its 'rotation_degrees' field. Note that specifying a + // region-of-interest using the 'region_of_interest' field is NOT supported + // and will result in an invalid argument error being returned. + // // The "result_callback" provides // - A vector of GestureRecognitionResult, each is the recognized results // for a input frame. @@ -157,7 +186,9 @@ class GestureRecognizer : tasks::vision::core::BaseVisionTaskApi { // longer be valid when the callback returns. To access the image data // outside of the callback, callers need to make a copy of the image. // - The input timestamp in milliseconds. - absl::Status RecognizeAsync(Image image, int64 timestamp_ms); + absl::Status RecognizeAsync(Image image, int64 timestamp_ms, + std::optional + image_processing_options = std::nullopt); // Shuts down the GestureRecognizer when all works are done. absl::Status Close() { return runner_->Close(); } diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer_graph.cc b/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer_graph.cc index b4f2af4d..7ab4847d 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer_graph.cc +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/gesture_recognizer_graph.cc @@ -24,9 +24,14 @@ limitations under the License. #include "mediapipe/framework/formats/classification.pb.h" #include "mediapipe/framework/formats/image.h" #include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/formats/rect.pb.h" +#include "mediapipe/framework/port/status_macros.h" #include "mediapipe/tasks/cc/common.h" +#include "mediapipe/tasks/cc/core/model_asset_bundle_resources.h" +#include "mediapipe/tasks/cc/core/model_resources_cache.h" #include "mediapipe/tasks/cc/core/model_task_graph.h" #include "mediapipe/tasks/cc/core/utils.h" +#include "mediapipe/tasks/cc/metadata/utils/zip_utils.h" #include "mediapipe/tasks/cc/vision/gesture_recognizer/proto/gesture_recognizer_graph_options.pb.h" #include "mediapipe/tasks/cc/vision/gesture_recognizer/proto/hand_gesture_recognizer_graph_options.pb.h" #include "mediapipe/tasks/cc/vision/hand_detector/proto/hand_detector_graph_options.pb.h" @@ -45,6 +50,8 @@ using ::mediapipe::api2::Input; using ::mediapipe::api2::Output; using ::mediapipe::api2::builder::Graph; using ::mediapipe::api2::builder::Source; +using ::mediapipe::tasks::core::ModelAssetBundleResources; +using ::mediapipe::tasks::metadata::SetExternalFile; using ::mediapipe::tasks::vision::gesture_recognizer::proto:: GestureRecognizerGraphOptions; using ::mediapipe::tasks::vision::gesture_recognizer::proto:: @@ -53,12 +60,16 @@ using ::mediapipe::tasks::vision::hand_landmarker::proto:: HandLandmarkerGraphOptions; constexpr char kImageTag[] = "IMAGE"; +constexpr char kNormRectTag[] = "NORM_RECT"; constexpr char kLandmarksTag[] = "LANDMARKS"; constexpr char kWorldLandmarksTag[] = "WORLD_LANDMARKS"; constexpr char kHandednessTag[] = "HANDEDNESS"; constexpr char kImageSizeTag[] = "IMAGE_SIZE"; constexpr char kHandGesturesTag[] = "HAND_GESTURES"; constexpr char kHandTrackingIdsTag[] = "HAND_TRACKING_IDS"; +constexpr char kHandLandmarkerBundleAssetName[] = "hand_landmarker.task"; +constexpr char kHandGestureRecognizerBundleAssetName[] = + "hand_gesture_recognizer.task"; struct GestureRecognizerOutputs { Source> gesture; @@ -68,6 +79,53 @@ struct GestureRecognizerOutputs { Source image; }; +// Sets the base options in the sub tasks. +absl::Status SetSubTaskBaseOptions(const ModelAssetBundleResources& resources, + GestureRecognizerGraphOptions* options, + bool is_copy) { + ASSIGN_OR_RETURN(const auto hand_landmarker_file, + resources.GetModelFile(kHandLandmarkerBundleAssetName)); + auto* hand_landmarker_graph_options = + options->mutable_hand_landmarker_graph_options(); + SetExternalFile(hand_landmarker_file, + hand_landmarker_graph_options->mutable_base_options() + ->mutable_model_asset(), + is_copy); + hand_landmarker_graph_options->mutable_base_options() + ->mutable_acceleration() + ->CopyFrom(options->base_options().acceleration()); + hand_landmarker_graph_options->mutable_base_options()->set_use_stream_mode( + options->base_options().use_stream_mode()); + + ASSIGN_OR_RETURN( + const auto hand_gesture_recognizer_file, + resources.GetModelFile(kHandGestureRecognizerBundleAssetName)); + auto* hand_gesture_recognizer_graph_options = + options->mutable_hand_gesture_recognizer_graph_options(); + SetExternalFile(hand_gesture_recognizer_file, + hand_gesture_recognizer_graph_options->mutable_base_options() + ->mutable_model_asset(), + is_copy); + hand_gesture_recognizer_graph_options->mutable_base_options() + ->mutable_acceleration() + ->CopyFrom(options->base_options().acceleration()); + if (!hand_gesture_recognizer_graph_options->base_options() + .acceleration() + .has_xnnpack() && + !hand_gesture_recognizer_graph_options->base_options() + .acceleration() + .has_tflite()) { + hand_gesture_recognizer_graph_options->mutable_base_options() + ->mutable_acceleration() + ->mutable_xnnpack(); + LOG(WARNING) << "Hand Gesture Recognizer contains CPU only ops. Sets " + << "HandGestureRecognizerGraph acceleartion to Xnnpack."; + } + hand_gesture_recognizer_graph_options->mutable_base_options() + ->set_use_stream_mode(options->base_options().use_stream_mode()); + return absl::OkStatus(); +} + } // namespace // A "mediapipe.tasks.vision.gesture_recognizer.GestureRecognizerGraph" performs @@ -76,6 +134,9 @@ struct GestureRecognizerOutputs { // Inputs: // IMAGE - Image // Image to perform hand gesture recognition on. +// NORM_RECT - NormalizedRect +// Describes image rotation and region of image to perform landmarks +// detection on. // // Outputs: // HAND_GESTURES - std::vector @@ -93,13 +154,15 @@ struct GestureRecognizerOutputs { // IMAGE - mediapipe::Image // The image that gesture recognizer runs on and has the pixel data stored // on the target storage (CPU vs GPU). -// +// All returned coordinates are in the unrotated and uncropped input image +// coordinates system. // // Example: // node { // calculator: // "mediapipe.tasks.vision.gesture_recognizer.GestureRecognizerGraph" // input_stream: "IMAGE:image_in" +// input_stream: "NORM_RECT:norm_rect" // output_stream: "HAND_GESTURES:hand_gestures" // output_stream: "LANDMARKS:hand_landmarks" // output_stream: "WORLD_LANDMARKS:world_hand_landmarks" @@ -129,10 +192,26 @@ class GestureRecognizerGraph : public core::ModelTaskGraph { absl::StatusOr GetConfig( SubgraphContext* sc) override { Graph graph; + if (sc->Options() + .base_options() + .has_model_asset()) { + ASSIGN_OR_RETURN( + const auto* model_asset_bundle_resources, + CreateModelAssetBundleResources(sc)); + // When the model resources cache service is available, filling in + // the file pointer meta in the subtasks' base options. Otherwise, + // providing the file contents instead. + MP_RETURN_IF_ERROR(SetSubTaskBaseOptions( + *model_asset_bundle_resources, + sc->MutableOptions(), + !sc->Service(::mediapipe::tasks::core::kModelResourcesCacheService) + .IsAvailable())); + } ASSIGN_OR_RETURN(auto hand_gesture_recognition_output, BuildGestureRecognizerGraph( *sc->MutableOptions(), - graph[Input(kImageTag)], graph)); + graph[Input(kImageTag)], + graph[Input(kNormRectTag)], graph)); hand_gesture_recognition_output.gesture >> graph[Output>(kHandGesturesTag)]; hand_gesture_recognition_output.handedness >> @@ -148,7 +227,7 @@ class GestureRecognizerGraph : public core::ModelTaskGraph { private: absl::StatusOr BuildGestureRecognizerGraph( GestureRecognizerGraphOptions& graph_options, Source image_in, - Graph& graph) { + Source norm_rect_in, Graph& graph) { auto& image_property = graph.AddNode("ImagePropertiesCalculator"); image_in >> image_property.In("IMAGE"); auto image_size = image_property.Out("SIZE"); @@ -162,6 +241,7 @@ class GestureRecognizerGraph : public core::ModelTaskGraph { graph_options.mutable_hand_landmarker_graph_options()); image_in >> hand_landmarker_graph.In(kImageTag); + norm_rect_in >> hand_landmarker_graph.In(kNormRectTag); auto hand_landmarks = hand_landmarker_graph[Output>( kLandmarksTag)]; @@ -187,6 +267,7 @@ class GestureRecognizerGraph : public core::ModelTaskGraph { hand_world_landmarks >> hand_gesture_subgraph.In(kWorldLandmarksTag); handedness >> hand_gesture_subgraph.In(kHandednessTag); image_size >> hand_gesture_subgraph.In(kImageSizeTag); + norm_rect_in >> hand_gesture_subgraph.In(kNormRectTag); hand_landmarks_id >> hand_gesture_subgraph.In(kHandTrackingIdsTag); auto hand_gestures = hand_gesture_subgraph[Output>( diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/hand_gesture_recognizer_graph.cc b/mediapipe/tasks/cc/vision/gesture_recognizer/hand_gesture_recognizer_graph.cc index 8d7e0bc0..7b6a8c79 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/hand_gesture_recognizer_graph.cc +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/hand_gesture_recognizer_graph.cc @@ -25,15 +25,25 @@ limitations under the License. #include "mediapipe/framework/formats/classification.pb.h" #include "mediapipe/framework/formats/landmark.pb.h" #include "mediapipe/framework/formats/matrix.h" +#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/formats/tensor.h" #include "mediapipe/tasks/cc/common.h" #include "mediapipe/tasks/cc/components/processors/classification_postprocessing_graph.h" #include "mediapipe/tasks/cc/components/processors/proto/classification_postprocessing_graph_options.pb.h" +#include "mediapipe/tasks/cc/components/processors/proto/classifier_options.pb.h" +#include "mediapipe/tasks/cc/core/model_asset_bundle_resources.h" #include "mediapipe/tasks/cc/core/model_resources.h" +#include "mediapipe/tasks/cc/core/model_resources_cache.h" #include "mediapipe/tasks/cc/core/model_task_graph.h" +#include "mediapipe/tasks/cc/core/proto/base_options.pb.h" +#include "mediapipe/tasks/cc/core/proto/external_file.pb.h" #include "mediapipe/tasks/cc/core/proto/inference_subgraph.pb.h" #include "mediapipe/tasks/cc/core/utils.h" +#include "mediapipe/tasks/cc/metadata/utils/zip_utils.h" +#include "mediapipe/tasks/cc/vision/gesture_recognizer/calculators/combined_prediction_calculator.pb.h" #include "mediapipe/tasks/cc/vision/gesture_recognizer/calculators/landmarks_to_matrix_calculator.pb.h" +#include "mediapipe/tasks/cc/vision/gesture_recognizer/proto/gesture_classifier_graph_options.pb.h" +#include "mediapipe/tasks/cc/vision/gesture_recognizer/proto/gesture_embedder_graph_options.pb.h" #include "mediapipe/tasks/cc/vision/gesture_recognizer/proto/hand_gesture_recognizer_graph_options.pb.h" #include "mediapipe/tasks/metadata/metadata_schema_generated.h" @@ -50,6 +60,9 @@ using ::mediapipe::api2::builder::Graph; using ::mediapipe::api2::builder::Source; using ::mediapipe::tasks::components::processors:: ConfigureTensorsToClassificationCalculator; +using ::mediapipe::tasks::core::ModelAssetBundleResources; +using ::mediapipe::tasks::core::proto::BaseOptions; +using ::mediapipe::tasks::metadata::SetExternalFile; using ::mediapipe::tasks::vision::gesture_recognizer::proto:: HandGestureRecognizerGraphOptions; @@ -57,6 +70,7 @@ constexpr char kHandednessTag[] = "HANDEDNESS"; constexpr char kLandmarksTag[] = "LANDMARKS"; constexpr char kWorldLandmarksTag[] = "WORLD_LANDMARKS"; constexpr char kImageSizeTag[] = "IMAGE_SIZE"; +constexpr char kNormRectTag[] = "NORM_RECT"; constexpr char kHandTrackingIdsTag[] = "HAND_TRACKING_IDS"; constexpr char kHandGesturesTag[] = "HAND_GESTURES"; constexpr char kLandmarksMatrixTag[] = "LANDMARKS_MATRIX"; @@ -68,6 +82,21 @@ constexpr char kVectorTag[] = "VECTOR"; constexpr char kIndexTag[] = "INDEX"; constexpr char kIterableTag[] = "ITERABLE"; constexpr char kBatchEndTag[] = "BATCH_END"; +constexpr char kPredictionTag[] = "PREDICTION"; +constexpr char kBackgroundLabel[] = "None"; +constexpr char kGestureEmbedderTFLiteName[] = "gesture_embedder.tflite"; +constexpr char kCannedGestureClassifierTFLiteName[] = + "canned_gesture_classifier.tflite"; +constexpr char kCustomGestureClassifierTFLiteName[] = + "custom_gesture_classifier.tflite"; + +struct SubTaskModelResources { + const core::ModelResources* gesture_embedder_model_resource = nullptr; + const core::ModelResources* canned_gesture_classifier_model_resource = + nullptr; + const core::ModelResources* custom_gesture_classifier_model_resource = + nullptr; +}; Source> ConvertMatrixToTensor(Source matrix, Graph& graph) { @@ -76,6 +105,21 @@ Source> ConvertMatrixToTensor(Source matrix, return node[Output>{"TENSORS"}]; } +absl::Status ConfigureCombinedPredictionCalculator( + CombinedPredictionCalculatorOptions* options) { + options->set_background_label(kBackgroundLabel); + return absl::OkStatus(); +} + +void PopulateAccelerationAndUseStreamMode( + const BaseOptions& parent_base_options, + BaseOptions* sub_task_base_options) { + sub_task_base_options->mutable_acceleration()->CopyFrom( + parent_base_options.acceleration()); + sub_task_base_options->set_use_stream_mode( + parent_base_options.use_stream_mode()); +} + } // namespace // A @@ -92,6 +136,9 @@ Source> ConvertMatrixToTensor(Source matrix, // Detected hand landmarks in world coordinates. // IMAGE_SIZE - std::pair // The size of image from which the landmarks detected from. +// NORM_RECT - NormalizedRect +// NormalizedRect whose 'rotation' field is used to rotate the +// landmarks before processing them. // // Outputs: // HAND_GESTURES - ClassificationList @@ -106,6 +153,7 @@ Source> ConvertMatrixToTensor(Source matrix, // input_stream: "LANDMARKS:landmarks" // input_stream: "WORLD_LANDMARKS:world_landmarks" // input_stream: "IMAGE_SIZE:image_size" +// input_stream: "NORM_RECT:norm_rect" // output_stream: "HAND_GESTURES:hand_gestures" // options { // [mediapipe.tasks.vision.gesture_recognizer.proto.HandGestureRecognizerGraphOptions.ext] @@ -122,30 +170,138 @@ class SingleHandGestureRecognizerGraph : public core::ModelTaskGraph { public: absl::StatusOr GetConfig( SubgraphContext* sc) override { - ASSIGN_OR_RETURN( - const auto* model_resources, - CreateModelResources(sc)); + if (sc->Options() + .base_options() + .has_model_asset()) { + ASSIGN_OR_RETURN( + const auto* model_asset_bundle_resources, + CreateModelAssetBundleResources( + sc)); + // When the model resources cache service is available, filling in + // the file pointer meta in the subtasks' base options. Otherwise, + // providing the file contents instead. + MP_RETURN_IF_ERROR(SetSubTaskBaseOptions( + *model_asset_bundle_resources, + sc->MutableOptions(), + !sc->Service(::mediapipe::tasks::core::kModelResourcesCacheService) + .IsAvailable())); + } + ASSIGN_OR_RETURN(const auto sub_task_model_resources, + CreateSubTaskModelResources(sc)); Graph graph; - ASSIGN_OR_RETURN( - auto hand_gestures, - BuildGestureRecognizerGraph( - sc->Options(), *model_resources, - graph[Input(kHandednessTag)], - graph[Input(kLandmarksTag)], - graph[Input(kWorldLandmarksTag)], - graph[Input>(kImageSizeTag)], graph)); + ASSIGN_OR_RETURN(auto hand_gestures, + BuildGestureRecognizerGraph( + sc->Options(), + sub_task_model_resources, + graph[Input(kHandednessTag)], + graph[Input(kLandmarksTag)], + graph[Input(kWorldLandmarksTag)], + graph[Input>(kImageSizeTag)], + graph[Input(kNormRectTag)], graph)); hand_gestures >> graph[Output(kHandGesturesTag)]; return graph.GetConfig(); } private: + // Sets the base options in the sub tasks. + absl::Status SetSubTaskBaseOptions(const ModelAssetBundleResources& resources, + HandGestureRecognizerGraphOptions* options, + bool is_copy) { + ASSIGN_OR_RETURN(const auto gesture_embedder_file, + resources.GetModelFile(kGestureEmbedderTFLiteName)); + auto* gesture_embedder_graph_options = + options->mutable_gesture_embedder_graph_options(); + SetExternalFile(gesture_embedder_file, + gesture_embedder_graph_options->mutable_base_options() + ->mutable_model_asset(), + is_copy); + PopulateAccelerationAndUseStreamMode( + options->base_options(), + gesture_embedder_graph_options->mutable_base_options()); + + ASSIGN_OR_RETURN( + const auto canned_gesture_classifier_file, + resources.GetModelFile(kCannedGestureClassifierTFLiteName)); + auto* canned_gesture_classifier_graph_options = + options->mutable_canned_gesture_classifier_graph_options(); + SetExternalFile( + canned_gesture_classifier_file, + canned_gesture_classifier_graph_options->mutable_base_options() + ->mutable_model_asset(), + is_copy); + PopulateAccelerationAndUseStreamMode( + options->base_options(), + canned_gesture_classifier_graph_options->mutable_base_options()); + + const auto custom_gesture_classifier_file = + resources.GetModelFile(kCustomGestureClassifierTFLiteName); + if (custom_gesture_classifier_file.ok()) { + has_custom_gesture_classifier = true; + auto* custom_gesture_classifier_graph_options = + options->mutable_custom_gesture_classifier_graph_options(); + SetExternalFile( + custom_gesture_classifier_file.value(), + custom_gesture_classifier_graph_options->mutable_base_options() + ->mutable_model_asset(), + is_copy); + PopulateAccelerationAndUseStreamMode( + options->base_options(), + custom_gesture_classifier_graph_options->mutable_base_options()); + } else { + LOG(INFO) << "Custom gesture classifier is not defined."; + } + return absl::OkStatus(); + } + + absl::StatusOr CreateSubTaskModelResources( + SubgraphContext* sc) { + auto* options = sc->MutableOptions(); + SubTaskModelResources sub_task_model_resources; + auto& gesture_embedder_model_asset = + *options->mutable_gesture_embedder_graph_options() + ->mutable_base_options() + ->mutable_model_asset(); + ASSIGN_OR_RETURN( + sub_task_model_resources.gesture_embedder_model_resource, + CreateModelResources(sc, + std::make_unique( + std::move(gesture_embedder_model_asset)), + "_gesture_embedder")); + auto& canned_gesture_classifier_model_asset = + *options->mutable_canned_gesture_classifier_graph_options() + ->mutable_base_options() + ->mutable_model_asset(); + ASSIGN_OR_RETURN( + sub_task_model_resources.canned_gesture_classifier_model_resource, + CreateModelResources( + sc, + std::make_unique( + std::move(canned_gesture_classifier_model_asset)), + "_canned_gesture_classifier")); + if (has_custom_gesture_classifier) { + auto& custom_gesture_classifier_model_asset = + *options->mutable_custom_gesture_classifier_graph_options() + ->mutable_base_options() + ->mutable_model_asset(); + ASSIGN_OR_RETURN( + sub_task_model_resources.custom_gesture_classifier_model_resource, + CreateModelResources( + sc, + std::make_unique( + std::move(custom_gesture_classifier_model_asset)), + "_custom_gesture_classifier")); + } + return sub_task_model_resources; + } + absl::StatusOr> BuildGestureRecognizerGraph( const HandGestureRecognizerGraphOptions& graph_options, - const core::ModelResources& model_resources, + const SubTaskModelResources& sub_task_model_resources, Source handedness, Source hand_landmarks, Source hand_world_landmarks, - Source> image_size, Graph& graph) { + Source> image_size, Source norm_rect, + Graph& graph) { // Converts the ClassificationList to a matrix. auto& handedness_to_matrix = graph.AddNode("HandednessToMatrixCalculator"); handedness >> handedness_to_matrix.In(kHandednessTag); @@ -166,6 +322,7 @@ class SingleHandGestureRecognizerGraph : public core::ModelTaskGraph { landmarks_options; hand_landmarks >> hand_landmarks_to_matrix.In(kLandmarksTag); image_size >> hand_landmarks_to_matrix.In(kImageSizeTag); + norm_rect >> hand_landmarks_to_matrix.In(kNormRectTag); auto hand_landmarks_matrix = hand_landmarks_to_matrix[Output(kLandmarksMatrixTag)]; @@ -181,6 +338,7 @@ class SingleHandGestureRecognizerGraph : public core::ModelTaskGraph { hand_world_landmarks >> hand_world_landmarks_to_matrix.In(kWorldLandmarksTag); image_size >> hand_world_landmarks_to_matrix.In(kImageSizeTag); + norm_rect >> hand_world_landmarks_to_matrix.In(kNormRectTag); auto hand_world_landmarks_matrix = hand_world_landmarks_to_matrix[Output(kLandmarksMatrixTag)]; @@ -198,26 +356,71 @@ class SingleHandGestureRecognizerGraph : public core::ModelTaskGraph { hand_world_landmarks_tensor >> concatenate_tensor_vector.In(2); auto concatenated_tensors = concatenate_tensor_vector.Out(""); - // Inference for static hand gesture recognition. - // TODO add embedding step. - auto& inference = AddInference( - model_resources, graph_options.base_options().acceleration(), graph); - concatenated_tensors >> inference.In(kTensorsTag); - auto inference_output_tensors = inference.Out(kTensorsTag); + // Inference for gesture embedder. + auto& gesture_embedder_inference = + AddInference(*sub_task_model_resources.gesture_embedder_model_resource, + graph_options.gesture_embedder_graph_options() + .base_options() + .acceleration(), + graph); + concatenated_tensors >> gesture_embedder_inference.In(kTensorsTag); + auto embedding_tensors = + gesture_embedder_inference.Out(kTensorsTag).Cast(); + auto& combine_predictions = graph.AddNode("CombinedPredictionCalculator"); + MP_RETURN_IF_ERROR(ConfigureCombinedPredictionCalculator( + &combine_predictions + .GetOptions())); + + int classifier_nums = 0; + // Inference for custom gesture classifier if it exists. + if (has_custom_gesture_classifier) { + ASSIGN_OR_RETURN( + auto gesture_clasification_list, + GetGestureClassificationList( + sub_task_model_resources.custom_gesture_classifier_model_resource, + graph_options.custom_gesture_classifier_graph_options(), + embedding_tensors, graph)); + gesture_clasification_list >> combine_predictions.In(classifier_nums++); + } + + // Inference for canned gesture classifier. + ASSIGN_OR_RETURN( + auto gesture_clasification_list, + GetGestureClassificationList( + sub_task_model_resources.canned_gesture_classifier_model_resource, + graph_options.canned_gesture_classifier_graph_options(), + embedding_tensors, graph)); + gesture_clasification_list >> combine_predictions.In(classifier_nums++); + + auto combined_classification_list = + combine_predictions.Out(kPredictionTag).Cast(); + + return combined_classification_list; + } + + absl::StatusOr> GetGestureClassificationList( + const core::ModelResources* model_resources, + const proto::GestureClassifierGraphOptions& options, + Source& embedding_tensors, Graph& graph) { + auto& gesture_classifier_inference = AddInference( + *model_resources, options.base_options().acceleration(), graph); + embedding_tensors >> gesture_classifier_inference.In(kTensorsTag); + auto gesture_inference_out_tensors = + gesture_classifier_inference.Out(kTensorsTag); auto& tensors_to_classification = graph.AddNode("TensorsToClassificationCalculator"); MP_RETURN_IF_ERROR(ConfigureTensorsToClassificationCalculator( - graph_options.classifier_options(), - *model_resources.GetMetadataExtractor(), 0, + options.classifier_options(), *model_resources->GetMetadataExtractor(), + 0, &tensors_to_classification.GetOptions< mediapipe::TensorsToClassificationCalculatorOptions>())); - inference_output_tensors >> tensors_to_classification.In(kTensorsTag); - auto classification_list = - tensors_to_classification[Output( - "CLASSIFICATIONS")]; - return classification_list; + gesture_inference_out_tensors >> tensors_to_classification.In(kTensorsTag); + return tensors_to_classification.Out("CLASSIFICATIONS") + .Cast(); } + + bool has_custom_gesture_classifier = false; }; // clang-format off @@ -239,6 +442,9 @@ REGISTER_MEDIAPIPE_GRAPH( // A vector hand landmarks in world coordinates. // IMAGE_SIZE - std::pair // The size of image from which the landmarks detected from. +// NORM_RECT - NormalizedRect +// NormalizedRect whose 'rotation' field is used to rotate the +// landmarks before processing them. // HAND_TRACKING_IDS - std::vector // A vector of the tracking ids of the hands. The tracking id is the vector // index corresponding to the same hand if the graph runs multiple times. @@ -257,6 +463,7 @@ REGISTER_MEDIAPIPE_GRAPH( // input_stream: "LANDMARKS:landmarks" // input_stream: "WORLD_LANDMARKS:world_landmarks" // input_stream: "IMAGE_SIZE:image_size" +// input_stream: "NORM_RECT:norm_rect" // input_stream: "HAND_TRACKING_IDS:hand_tracking_ids" // output_stream: "HAND_GESTURES:hand_gestures" // options { @@ -283,6 +490,7 @@ class MultipleHandGestureRecognizerGraph : public core::ModelTaskGraph { graph[Input>(kLandmarksTag)], graph[Input>(kWorldLandmarksTag)], graph[Input>(kImageSizeTag)], + graph[Input(kNormRectTag)], graph[Input>(kHandTrackingIdsTag)], graph)); multi_hand_gestures >> graph[Output>(kHandGesturesTag)]; @@ -296,18 +504,20 @@ class MultipleHandGestureRecognizerGraph : public core::ModelTaskGraph { Source> multi_handedness, Source> multi_hand_landmarks, Source> multi_hand_world_landmarks, - Source> image_size, + Source> image_size, Source norm_rect, Source> multi_hand_tracking_ids, Graph& graph) { auto& begin_loop_int = graph.AddNode("BeginLoopIntCalculator"); image_size >> begin_loop_int.In(kCloneTag)[0]; - multi_handedness >> begin_loop_int.In(kCloneTag)[1]; - multi_hand_landmarks >> begin_loop_int.In(kCloneTag)[2]; - multi_hand_world_landmarks >> begin_loop_int.In(kCloneTag)[3]; + norm_rect >> begin_loop_int.In(kCloneTag)[1]; + multi_handedness >> begin_loop_int.In(kCloneTag)[2]; + multi_hand_landmarks >> begin_loop_int.In(kCloneTag)[3]; + multi_hand_world_landmarks >> begin_loop_int.In(kCloneTag)[4]; multi_hand_tracking_ids >> begin_loop_int.In(kIterableTag); auto image_size_clone = begin_loop_int.Out(kCloneTag)[0]; - auto multi_handedness_clone = begin_loop_int.Out(kCloneTag)[1]; - auto multi_hand_landmarks_clone = begin_loop_int.Out(kCloneTag)[2]; - auto multi_hand_world_landmarks_clone = begin_loop_int.Out(kCloneTag)[3]; + auto norm_rect_clone = begin_loop_int.Out(kCloneTag)[1]; + auto multi_handedness_clone = begin_loop_int.Out(kCloneTag)[2]; + auto multi_hand_landmarks_clone = begin_loop_int.Out(kCloneTag)[3]; + auto multi_hand_world_landmarks_clone = begin_loop_int.Out(kCloneTag)[4]; auto hand_tracking_id = begin_loop_int.Out(kItemTag); auto batch_end = begin_loop_int.Out(kBatchEndTag); @@ -341,6 +551,7 @@ class MultipleHandGestureRecognizerGraph : public core::ModelTaskGraph { hand_world_landmarks >> hand_gesture_recognizer_graph.In(kWorldLandmarksTag); image_size_clone >> hand_gesture_recognizer_graph.In(kImageSizeTag); + norm_rect_clone >> hand_gesture_recognizer_graph.In(kNormRectTag); auto hand_gestures = hand_gesture_recognizer_graph.Out(kHandGesturesTag); auto& end_loop_classification_lists = diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/proto/BUILD b/mediapipe/tasks/cc/vision/gesture_recognizer/proto/BUILD index 3b73bf2b..0db47da7 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/proto/BUILD +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/proto/BUILD @@ -49,7 +49,6 @@ mediapipe_proto_library( ":gesture_embedder_graph_options_proto", "//mediapipe/framework:calculator_options_proto", "//mediapipe/framework:calculator_proto", - "//mediapipe/tasks/cc/components/processors/proto:classifier_options_proto", "//mediapipe/tasks/cc/core/proto:base_options_proto", ], ) diff --git a/mediapipe/tasks/cc/vision/gesture_recognizer/proto/hand_gesture_recognizer_graph_options.proto b/mediapipe/tasks/cc/vision/gesture_recognizer/proto/hand_gesture_recognizer_graph_options.proto index a3281702..7df2fed3 100644 --- a/mediapipe/tasks/cc/vision/gesture_recognizer/proto/hand_gesture_recognizer_graph_options.proto +++ b/mediapipe/tasks/cc/vision/gesture_recognizer/proto/hand_gesture_recognizer_graph_options.proto @@ -18,7 +18,6 @@ syntax = "proto2"; package mediapipe.tasks.vision.gesture_recognizer.proto; import "mediapipe/framework/calculator.proto"; -import "mediapipe/tasks/cc/components/processors/proto/classifier_options.proto"; import "mediapipe/tasks/cc/core/proto/base_options.proto"; import "mediapipe/tasks/cc/vision/gesture_recognizer/proto/gesture_classifier_graph_options.proto"; import "mediapipe/tasks/cc/vision/gesture_recognizer/proto/gesture_embedder_graph_options.proto"; @@ -37,15 +36,11 @@ message HandGestureRecognizerGraphOptions { // Options for GestureEmbedder. optional GestureEmbedderGraphOptions gesture_embedder_graph_options = 2; - // Options for GestureClassifier of default gestures. + // Options for GestureClassifier of canned gestures. optional GestureClassifierGraphOptions canned_gesture_classifier_graph_options = 3; // Options for GestureClassifier of custom gestures. optional GestureClassifierGraphOptions custom_gesture_classifier_graph_options = 4; - - // TODO: remove these. Temporary solutions before bundle asset is - // ready. - optional components.processors.proto.ClassifierOptions classifier_options = 5; } diff --git a/mediapipe/tasks/cc/vision/hand_detector/BUILD b/mediapipe/tasks/cc/vision/hand_detector/BUILD index 433a3047..71cef627 100644 --- a/mediapipe/tasks/cc/vision/hand_detector/BUILD +++ b/mediapipe/tasks/cc/vision/hand_detector/BUILD @@ -32,7 +32,7 @@ cc_library( "//mediapipe/calculators/tflite:ssd_anchors_calculator_cc_proto", "//mediapipe/calculators/util:detection_label_id_to_text_calculator", "//mediapipe/calculators/util:detection_label_id_to_text_calculator_cc_proto", - "//mediapipe/calculators/util:detection_letterbox_removal_calculator", + "//mediapipe/calculators/util:detection_projection_calculator", "//mediapipe/calculators/util:detections_to_rects_calculator", "//mediapipe/calculators/util:detections_to_rects_calculator_cc_proto", "//mediapipe/calculators/util:non_max_suppression_calculator", diff --git a/mediapipe/tasks/cc/vision/hand_detector/hand_detector_graph.cc b/mediapipe/tasks/cc/vision/hand_detector/hand_detector_graph.cc index 8573d718..06bb2e54 100644 --- a/mediapipe/tasks/cc/vision/hand_detector/hand_detector_graph.cc +++ b/mediapipe/tasks/cc/vision/hand_detector/hand_detector_graph.cc @@ -58,6 +58,7 @@ using ::mediapipe::tasks::vision::hand_detector::proto:: HandDetectorGraphOptions; constexpr char kImageTag[] = "IMAGE"; +constexpr char kNormRectTag[] = "NORM_RECT"; constexpr char kPalmDetectionsTag[] = "PALM_DETECTIONS"; constexpr char kHandRectsTag[] = "HAND_RECTS"; constexpr char kPalmRectsTag[] = "PALM_RECTS"; @@ -148,6 +149,9 @@ void ConfigureRectTransformationCalculator( // Inputs: // IMAGE - Image // Image to perform detection on. +// NORM_RECT - NormalizedRect +// Describes image rotation and region of image to perform detection +// on. // // Outputs: // PALM_DETECTIONS - std::vector @@ -159,11 +163,14 @@ void ConfigureRectTransformationCalculator( // IMAGE - Image // The input image that the hand detector runs on and has the pixel data // stored on the target storage (CPU vs GPU). +// All returned coordinates are in the unrotated and uncropped input image +// coordinates system. // // Example: // node { // calculator: "mediapipe.tasks.vision.hand_detector.HandDetectorGraph" // input_stream: "IMAGE:image" +// input_stream: "NORM_RECT:norm_rect" // output_stream: "PALM_DETECTIONS:palm_detections" // output_stream: "HAND_RECTS:hand_rects_from_palm_detections" // output_stream: "PALM_RECTS:palm_rects" @@ -189,11 +196,11 @@ class HandDetectorGraph : public core::ModelTaskGraph { ASSIGN_OR_RETURN(const auto* model_resources, CreateModelResources(sc)); Graph graph; - ASSIGN_OR_RETURN( - auto hand_detection_outs, - BuildHandDetectionSubgraph(sc->Options(), - *model_resources, - graph[Input(kImageTag)], graph)); + ASSIGN_OR_RETURN(auto hand_detection_outs, + BuildHandDetectionSubgraph( + sc->Options(), + *model_resources, graph[Input(kImageTag)], + graph[Input(kNormRectTag)], graph)); hand_detection_outs.palm_detections >> graph[Output>(kPalmDetectionsTag)]; hand_detection_outs.hand_rects >> @@ -216,7 +223,7 @@ class HandDetectorGraph : public core::ModelTaskGraph { absl::StatusOr BuildHandDetectionSubgraph( const HandDetectorGraphOptions& subgraph_options, const core::ModelResources& model_resources, Source image_in, - Graph& graph) { + Source norm_rect_in, Graph& graph) { // Add image preprocessing subgraph. The model expects aspect ratio // unchanged. auto& preprocessing = @@ -228,13 +235,16 @@ class HandDetectorGraph : public core::ModelTaskGraph { image_to_tensor_options.set_keep_aspect_ratio(true); image_to_tensor_options.set_border_mode( mediapipe::ImageToTensorCalculatorOptions::BORDER_ZERO); + bool use_gpu = components::DetermineImagePreprocessingGpuBackend( + subgraph_options.base_options().acceleration()); MP_RETURN_IF_ERROR(ConfigureImagePreprocessing( - model_resources, + model_resources, use_gpu, &preprocessing .GetOptions())); image_in >> preprocessing.In("IMAGE"); + norm_rect_in >> preprocessing.In("NORM_RECT"); auto preprocessed_tensors = preprocessing.Out("TENSORS"); - auto letterbox_padding = preprocessing.Out("LETTERBOX_PADDING"); + auto matrix = preprocessing.Out("MATRIX"); auto image_size = preprocessing.Out("IMAGE_SIZE"); // Adds SSD palm detection model. @@ -278,17 +288,12 @@ class HandDetectorGraph : public core::ModelTaskGraph { nms_detections >> detection_label_id_to_text.In(""); auto detections_with_text = detection_label_id_to_text.Out(""); - // Adjusts detection locations (already normalized to [0.f, 1.f]) on the - // letterboxed image (after image transformation with the FIT scale mode) to - // the corresponding locations on the same image with the letterbox removed - // (the input image to the graph before image transformation). - auto& detection_letterbox_removal = - graph.AddNode("DetectionLetterboxRemovalCalculator"); - detections_with_text >> detection_letterbox_removal.In("DETECTIONS"); - letterbox_padding >> detection_letterbox_removal.In("LETTERBOX_PADDING"); + // Projects detections back into the input image coordinates system. + auto& detection_projection = graph.AddNode("DetectionProjectionCalculator"); + detections_with_text >> detection_projection.In("DETECTIONS"); + matrix >> detection_projection.In("PROJECTION_MATRIX"); auto palm_detections = - detection_letterbox_removal[Output>( - "DETECTIONS")]; + detection_projection[Output>("DETECTIONS")]; // Converts each palm detection into a rectangle (normalized by image size) // that encloses the palm and is rotated such that the line connecting diff --git a/mediapipe/tasks/cc/vision/hand_detector/hand_detector_graph_test.cc b/mediapipe/tasks/cc/vision/hand_detector/hand_detector_graph_test.cc index 11cfc302..cbbc0e19 100644 --- a/mediapipe/tasks/cc/vision/hand_detector/hand_detector_graph_test.cc +++ b/mediapipe/tasks/cc/vision/hand_detector/hand_detector_graph_test.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include #include @@ -75,13 +76,18 @@ using ::testing::proto::Partially; constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/vision/"; constexpr char kPalmDetectionModel[] = "palm_detection_full.tflite"; constexpr char kTestRightHandsImage[] = "right_hands.jpg"; +constexpr char kTestRightHandsRotatedImage[] = "right_hands_rotated.jpg"; constexpr char kTestModelResourcesTag[] = "test_model_resources"; constexpr char kOneHandResultFile[] = "hand_detector_result_one_hand.pbtxt"; +constexpr char kOneHandRotatedResultFile[] = + "hand_detector_result_one_hand_rotated.pbtxt"; constexpr char kTwoHandsResultFile[] = "hand_detector_result_two_hands.pbtxt"; constexpr char kImageTag[] = "IMAGE"; constexpr char kImageName[] = "image"; +constexpr char kNormRectTag[] = "NORM_RECT"; +constexpr char kNormRectName[] = "norm_rect"; constexpr char kPalmDetectionsTag[] = "PALM_DETECTIONS"; constexpr char kPalmDetectionsName[] = "palm_detections"; constexpr char kHandRectsTag[] = "HAND_RECTS"; @@ -117,6 +123,8 @@ absl::StatusOr> CreateTaskRunner( graph[Input(kImageTag)].SetName(kImageName) >> hand_detection.In(kImageTag); + graph[Input(kNormRectTag)].SetName(kNormRectName) >> + hand_detection.In(kNormRectTag); hand_detection.Out(kPalmDetectionsTag).SetName(kPalmDetectionsName) >> graph[Output>(kPalmDetectionsTag)]; @@ -142,6 +150,9 @@ struct TestParams { std::string hand_detection_model_name; // The filename of test image. std::string test_image_name; + // The rotation to apply to the test image before processing, in radians + // counter-clockwise. + float rotation; // The number of maximum detected hands. int num_hands; // The expected hand detector result. @@ -154,14 +165,22 @@ TEST_P(HandDetectionTest, DetectTwoHands) { MP_ASSERT_OK_AND_ASSIGN( Image image, DecodeImageFromFile(JoinPath("./", kTestDataDirectory, GetParam().test_image_name))); + NormalizedRect input_norm_rect; + input_norm_rect.set_rotation(GetParam().rotation); + input_norm_rect.set_x_center(0.5); + input_norm_rect.set_y_center(0.5); + input_norm_rect.set_width(1.0); + input_norm_rect.set_height(1.0); MP_ASSERT_OK_AND_ASSIGN( auto model_resources, CreateModelResourcesForModel(GetParam().hand_detection_model_name)); MP_ASSERT_OK_AND_ASSIGN( auto task_runner, CreateTaskRunner(*model_resources, kPalmDetectionModel, GetParam().num_hands)); - auto output_packets = - task_runner->Process({{kImageName, MakePacket(std::move(image))}}); + auto output_packets = task_runner->Process( + {{kImageName, MakePacket(std::move(image))}, + {kNormRectName, + MakePacket(std::move(input_norm_rect))}}); MP_ASSERT_OK(output_packets); const std::vector& palm_detections = (*output_packets)[kPalmDetectionsName].Get>(); @@ -188,15 +207,24 @@ INSTANTIATE_TEST_SUITE_P( Values(TestParams{.test_name = "DetectOneHand", .hand_detection_model_name = kPalmDetectionModel, .test_image_name = kTestRightHandsImage, + .rotation = 0, .num_hands = 1, .expected_result = GetExpectedHandDetectorResult(kOneHandResultFile)}, TestParams{.test_name = "DetectTwoHands", .hand_detection_model_name = kPalmDetectionModel, .test_image_name = kTestRightHandsImage, + .rotation = 0, .num_hands = 2, .expected_result = - GetExpectedHandDetectorResult(kTwoHandsResultFile)}), + GetExpectedHandDetectorResult(kTwoHandsResultFile)}, + TestParams{.test_name = "DetectOneHandWithRotation", + .hand_detection_model_name = kPalmDetectionModel, + .test_image_name = kTestRightHandsRotatedImage, + .rotation = M_PI / 2.0f, + .num_hands = 1, + .expected_result = GetExpectedHandDetectorResult( + kOneHandRotatedResultFile)}), [](const TestParamInfo& info) { return info.param.test_name; }); diff --git a/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarker_graph.cc b/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarker_graph.cc index 7e199348..e610a412 100644 --- a/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarker_graph.cc +++ b/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarker_graph.cc @@ -64,6 +64,7 @@ using ::mediapipe::tasks::vision::hand_landmarker::proto:: HandLandmarksDetectorGraphOptions; constexpr char kImageTag[] = "IMAGE"; +constexpr char kNormRectTag[] = "NORM_RECT"; constexpr char kLandmarksTag[] = "LANDMARKS"; constexpr char kWorldLandmarksTag[] = "WORLD_LANDMARKS"; constexpr char kHandRectNextFrameTag[] = "HAND_RECT_NEXT_FRAME"; @@ -91,18 +92,30 @@ absl::Status SetSubTaskBaseOptions(const ModelAssetBundleResources& resources, bool is_copy) { ASSIGN_OR_RETURN(const auto hand_detector_file, resources.GetModelFile(kHandDetectorTFLiteName)); + auto* hand_detector_graph_options = + options->mutable_hand_detector_graph_options(); SetExternalFile(hand_detector_file, - options->mutable_hand_detector_graph_options() - ->mutable_base_options() + hand_detector_graph_options->mutable_base_options() ->mutable_model_asset(), is_copy); + hand_detector_graph_options->mutable_base_options() + ->mutable_acceleration() + ->CopyFrom(options->base_options().acceleration()); + hand_detector_graph_options->mutable_base_options()->set_use_stream_mode( + options->base_options().use_stream_mode()); ASSIGN_OR_RETURN(const auto hand_landmarks_detector_file, resources.GetModelFile(kHandLandmarksDetectorTFLiteName)); + auto* hand_landmarks_detector_graph_options = + options->mutable_hand_landmarks_detector_graph_options(); SetExternalFile(hand_landmarks_detector_file, - options->mutable_hand_landmarks_detector_graph_options() - ->mutable_base_options() + hand_landmarks_detector_graph_options->mutable_base_options() ->mutable_model_asset(), is_copy); + hand_landmarks_detector_graph_options->mutable_base_options() + ->mutable_acceleration() + ->CopyFrom(options->base_options().acceleration()); + hand_landmarks_detector_graph_options->mutable_base_options() + ->set_use_stream_mode(options->base_options().use_stream_mode()); return absl::OkStatus(); } @@ -122,6 +135,9 @@ absl::Status SetSubTaskBaseOptions(const ModelAssetBundleResources& resources, // Inputs: // IMAGE - Image // Image to perform hand landmarks detection on. +// NORM_RECT - NormalizedRect +// Describes image rotation and region of image to perform landmarks +// detection on. // // Outputs: // LANDMARKS: - std::vector @@ -140,11 +156,14 @@ absl::Status SetSubTaskBaseOptions(const ModelAssetBundleResources& resources, // IMAGE - Image // The input image that the hand landmarker runs on and has the pixel data // stored on the target storage (CPU vs GPU). +// All returned coordinates are in the unrotated and uncropped input image +// coordinates system. // // Example: // node { // calculator: "mediapipe.tasks.vision.hand_landmarker.HandLandmarkerGraph" // input_stream: "IMAGE:image_in" +// input_stream: "NORM_RECT:norm_rect" // output_stream: "LANDMARKS:hand_landmarks" // output_stream: "WORLD_LANDMARKS:world_hand_landmarks" // output_stream: "HAND_RECT_NEXT_FRAME:hand_rect_next_frame" @@ -198,10 +217,11 @@ class HandLandmarkerGraph : public core::ModelTaskGraph { !sc->Service(::mediapipe::tasks::core::kModelResourcesCacheService) .IsAvailable())); } - ASSIGN_OR_RETURN( - auto hand_landmarker_outputs, - BuildHandLandmarkerGraph(sc->Options(), - graph[Input(kImageTag)], graph)); + ASSIGN_OR_RETURN(auto hand_landmarker_outputs, + BuildHandLandmarkerGraph( + sc->Options(), + graph[Input(kImageTag)], + graph[Input(kNormRectTag)], graph)); hand_landmarker_outputs.landmark_lists >> graph[Output>(kLandmarksTag)]; hand_landmarker_outputs.world_landmark_lists >> @@ -240,7 +260,7 @@ class HandLandmarkerGraph : public core::ModelTaskGraph { // graph: the mediapipe graph instance to be updated. absl::StatusOr BuildHandLandmarkerGraph( const HandLandmarkerGraphOptions& tasks_options, Source image_in, - Graph& graph) { + Source norm_rect_in, Graph& graph) { const int max_num_hands = tasks_options.hand_detector_graph_options().num_hands(); @@ -258,12 +278,15 @@ class HandLandmarkerGraph : public core::ModelTaskGraph { auto image_for_hand_detector = DisallowIf(image_in, has_enough_hands, graph); + auto norm_rect_in_for_hand_detector = + DisallowIf(norm_rect_in, has_enough_hands, graph); auto& hand_detector = graph.AddNode("mediapipe.tasks.vision.hand_detector.HandDetectorGraph"); hand_detector.GetOptions().CopyFrom( tasks_options.hand_detector_graph_options()); image_for_hand_detector >> hand_detector.In("IMAGE"); + norm_rect_in_for_hand_detector >> hand_detector.In("NORM_RECT"); auto hand_rects_from_hand_detector = hand_detector.Out("HAND_RECTS"); auto& hand_association = graph.AddNode("HandAssociationCalculator"); diff --git a/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarker_graph_test.cc b/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarker_graph_test.cc index 604f37d5..f275486f 100644 --- a/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarker_graph_test.cc +++ b/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarker_graph_test.cc @@ -13,10 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include #include #include +#include #include "absl/flags/flag.h" #include "absl/status/statusor.h" @@ -65,11 +67,14 @@ using ::testing::proto::Approximately; using ::testing::proto::Partially; constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/vision/"; -constexpr char kHandLandmarkerModelBundle[] = "hand_landmark.task"; +constexpr char kHandLandmarkerModelBundle[] = "hand_landmarker.task"; constexpr char kLeftHandsImage[] = "left_hands.jpg"; +constexpr char kLeftHandsRotatedImage[] = "left_hands_rotated.jpg"; constexpr char kImageTag[] = "IMAGE"; constexpr char kImageName[] = "image_in"; +constexpr char kNormRectTag[] = "NORM_RECT"; +constexpr char kNormRectName[] = "norm_rect_in"; constexpr char kLandmarksTag[] = "LANDMARKS"; constexpr char kLandmarksName[] = "landmarks"; constexpr char kWorldLandmarksTag[] = "WORLD_LANDMARKS"; @@ -84,6 +89,11 @@ constexpr char kExpectedLeftUpHandLandmarksFilename[] = "expected_left_up_hand_landmarks.prototxt"; constexpr char kExpectedLeftDownHandLandmarksFilename[] = "expected_left_down_hand_landmarks.prototxt"; +// Same but for the rotated image. +constexpr char kExpectedLeftUpHandRotatedLandmarksFilename[] = + "expected_left_up_hand_rotated_landmarks.prototxt"; +constexpr char kExpectedLeftDownHandRotatedLandmarksFilename[] = + "expected_left_down_hand_rotated_landmarks.prototxt"; constexpr float kFullModelFractionDiff = 0.03; // percentage constexpr float kAbsMargin = 0.03; @@ -111,6 +121,8 @@ absl::StatusOr> CreateTaskRunner() { graph[Input(kImageTag)].SetName(kImageName) >> hand_landmarker_graph.In(kImageTag); + graph[Input(kNormRectTag)].SetName(kNormRectName) >> + hand_landmarker_graph.In(kNormRectTag); hand_landmarker_graph.Out(kLandmarksTag).SetName(kLandmarksName) >> graph[Output>(kLandmarksTag)]; hand_landmarker_graph.Out(kWorldLandmarksTag).SetName(kWorldLandmarksName) >> @@ -130,9 +142,16 @@ TEST_F(HandLandmarkerTest, Succeeds) { MP_ASSERT_OK_AND_ASSIGN( Image image, DecodeImageFromFile(JoinPath("./", kTestDataDirectory, kLeftHandsImage))); + NormalizedRect input_norm_rect; + input_norm_rect.set_x_center(0.5); + input_norm_rect.set_y_center(0.5); + input_norm_rect.set_width(1.0); + input_norm_rect.set_height(1.0); MP_ASSERT_OK_AND_ASSIGN(auto task_runner, CreateTaskRunner()); - auto output_packets = - task_runner->Process({{kImageName, MakePacket(std::move(image))}}); + auto output_packets = task_runner->Process( + {{kImageName, MakePacket(std::move(image))}, + {kNormRectName, + MakePacket(std::move(input_norm_rect))}}); const auto& landmarks = (*output_packets)[kLandmarksName] .Get>(); ASSERT_EQ(landmarks.size(), kMaxNumHands); @@ -150,6 +169,38 @@ TEST_F(HandLandmarkerTest, Succeeds) { /*fraction=*/kFullModelFractionDiff)); } +TEST_F(HandLandmarkerTest, SucceedsWithRotation) { + MP_ASSERT_OK_AND_ASSIGN( + Image image, DecodeImageFromFile(JoinPath("./", kTestDataDirectory, + kLeftHandsRotatedImage))); + NormalizedRect input_norm_rect; + input_norm_rect.set_x_center(0.5); + input_norm_rect.set_y_center(0.5); + input_norm_rect.set_width(1.0); + input_norm_rect.set_height(1.0); + input_norm_rect.set_rotation(M_PI / 2.0); + MP_ASSERT_OK_AND_ASSIGN(auto task_runner, CreateTaskRunner()); + auto output_packets = task_runner->Process( + {{kImageName, MakePacket(std::move(image))}, + {kNormRectName, + MakePacket(std::move(input_norm_rect))}}); + const auto& landmarks = (*output_packets)[kLandmarksName] + .Get>(); + ASSERT_EQ(landmarks.size(), kMaxNumHands); + std::vector expected_landmarks = { + GetExpectedLandmarkList(kExpectedLeftUpHandRotatedLandmarksFilename), + GetExpectedLandmarkList(kExpectedLeftDownHandRotatedLandmarksFilename)}; + + EXPECT_THAT(landmarks[0], + Approximately(Partially(EqualsProto(expected_landmarks[0])), + /*margin=*/kAbsMargin, + /*fraction=*/kFullModelFractionDiff)); + EXPECT_THAT(landmarks[1], + Approximately(Partially(EqualsProto(expected_landmarks[1])), + /*margin=*/kAbsMargin, + /*fraction=*/kFullModelFractionDiff)); +} + } // namespace } // namespace hand_landmarker diff --git a/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarks_detector_graph.cc b/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarks_detector_graph.cc index 23521790..1f127deb 100644 --- a/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarks_detector_graph.cc +++ b/mediapipe/tasks/cc/vision/hand_landmarker/hand_landmarks_detector_graph.cc @@ -283,8 +283,10 @@ class SingleHandLandmarksDetectorGraph : public core::ModelTaskGraph { auto& preprocessing = graph.AddNode("mediapipe.tasks.components.ImagePreprocessingSubgraph"); + bool use_gpu = components::DetermineImagePreprocessingGpuBackend( + subgraph_options.base_options().acceleration()); MP_RETURN_IF_ERROR(ConfigureImagePreprocessing( - model_resources, + model_resources, use_gpu, &preprocessing .GetOptions())); image_in >> preprocessing.In("IMAGE"); diff --git a/mediapipe/tasks/cc/vision/image_classifier/BUILD b/mediapipe/tasks/cc/vision/image_classifier/BUILD index dfa77cb9..b59d8d68 100644 --- a/mediapipe/tasks/cc/vision/image_classifier/BUILD +++ b/mediapipe/tasks/cc/vision/image_classifier/BUILD @@ -50,6 +50,7 @@ cc_library( "//mediapipe/framework/api2:builder", "//mediapipe/framework/formats:image", "//mediapipe/framework/formats:rect_cc_proto", + "//mediapipe/tasks/cc/components/containers:classification_result", "//mediapipe/tasks/cc/components/containers/proto:classifications_cc_proto", "//mediapipe/tasks/cc/components/processors:classifier_options", "//mediapipe/tasks/cc/components/processors/proto:classifier_options_cc_proto", @@ -59,6 +60,7 @@ cc_library( "//mediapipe/tasks/cc/core/proto:base_options_cc_proto", "//mediapipe/tasks/cc/core/proto:inference_subgraph_cc_proto", "//mediapipe/tasks/cc/vision/core:base_vision_task_api", + "//mediapipe/tasks/cc/vision/core:image_processing_options", "//mediapipe/tasks/cc/vision/core:running_mode", "//mediapipe/tasks/cc/vision/core:vision_task_api_factory", "//mediapipe/tasks/cc/vision/image_classifier/proto:image_classifier_graph_options_cc_proto", diff --git a/mediapipe/tasks/cc/vision/image_classifier/image_classifier.cc b/mediapipe/tasks/cc/vision/image_classifier/image_classifier.cc index f3dcdd07..60f8f7ed 100644 --- a/mediapipe/tasks/cc/vision/image_classifier/image_classifier.cc +++ b/mediapipe/tasks/cc/vision/image_classifier/image_classifier.cc @@ -26,6 +26,7 @@ limitations under the License. #include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/packet.h" #include "mediapipe/framework/timestamp.h" +#include "mediapipe/tasks/cc/components/containers/classification_result.h" #include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h" #include "mediapipe/tasks/cc/components/processors/classifier_options.h" #include "mediapipe/tasks/cc/components/processors/proto/classifier_options.pb.h" @@ -34,6 +35,7 @@ limitations under the License. #include "mediapipe/tasks/cc/core/proto/inference_subgraph.pb.h" #include "mediapipe/tasks/cc/core/task_runner.h" #include "mediapipe/tasks/cc/core/utils.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/core/vision_task_api_factory.h" #include "mediapipe/tasks/cc/vision/image_classifier/proto/image_classifier_graph_options.pb.h" @@ -45,8 +47,8 @@ namespace image_classifier { namespace { -constexpr char kClassificationResultStreamName[] = "classification_result_out"; -constexpr char kClassificationResultTag[] = "CLASSIFICATION_RESULT"; +constexpr char kClassificationsStreamName[] = "classifications_out"; +constexpr char kClassificationsTag[] = "CLASSIFICATIONS"; constexpr char kImageInStreamName[] = "image_in"; constexpr char kImageOutStreamName[] = "image_out"; constexpr char kImageTag[] = "IMAGE"; @@ -56,29 +58,10 @@ constexpr char kSubgraphTypeName[] = "mediapipe.tasks.vision.image_classifier.ImageClassifierGraph"; constexpr int kMicroSecondsPerMilliSecond = 1000; +using ::mediapipe::tasks::components::containers::ConvertToClassificationResult; using ::mediapipe::tasks::components::containers::proto::ClassificationResult; using ::mediapipe::tasks::core::PacketMap; -// Returns a NormalizedRect covering the full image if input is not present. -// Otherwise, makes sure the x_center, y_center, width and height are set in -// case only a rotation was provided in the input. -NormalizedRect FillNormalizedRect( - std::optional normalized_rect) { - NormalizedRect result; - if (normalized_rect.has_value()) { - result = *normalized_rect; - } - bool has_coordinates = result.has_x_center() || result.has_y_center() || - result.has_width() || result.has_height(); - if (!has_coordinates) { - result.set_x_center(0.5); - result.set_y_center(0.5); - result.set_width(1); - result.set_height(1); - } - return result; -} - // Creates a MediaPipe graph config that contains a subgraph node of // type "ImageClassifierGraph". If the task is running in the live stream mode, // a "FlowLimiterCalculator" will be added to limit the number of frames in @@ -92,15 +75,13 @@ CalculatorGraphConfig CreateGraphConfig( auto& task_subgraph = graph.AddNode(kSubgraphTypeName); task_subgraph.GetOptions().Swap( options_proto.get()); - task_subgraph.Out(kClassificationResultTag) - .SetName(kClassificationResultStreamName) >> - graph.Out(kClassificationResultTag); + task_subgraph.Out(kClassificationsTag).SetName(kClassificationsStreamName) >> + graph.Out(kClassificationsTag); task_subgraph.Out(kImageTag).SetName(kImageOutStreamName) >> graph.Out(kImageTag); if (enable_flow_limiting) { - return tasks::core::AddFlowLimiterCalculator(graph, task_subgraph, - {kImageTag, kNormRectTag}, - kClassificationResultTag); + return tasks::core::AddFlowLimiterCalculator( + graph, task_subgraph, {kImageTag, kNormRectTag}, kClassificationsTag); } graph.In(kImageTag) >> task_subgraph.In(kImageTag); graph.In(kNormRectTag) >> task_subgraph.In(kNormRectTag); @@ -144,13 +125,14 @@ absl::StatusOr> ImageClassifier::Create( if (status_or_packets.value()[kImageOutStreamName].IsEmpty()) { return; } - Packet classification_result_packet = - status_or_packets.value()[kClassificationResultStreamName]; + Packet classifications_packet = + status_or_packets.value()[kClassificationsStreamName]; Packet image_packet = status_or_packets.value()[kImageOutStreamName]; result_callback( - classification_result_packet.Get(), + ConvertToClassificationResult( + classifications_packet.Get()), image_packet.Get(), - classification_result_packet.Timestamp().Value() / + classifications_packet.Timestamp().Value() / kMicroSecondsPerMilliSecond); }; } @@ -163,34 +145,37 @@ absl::StatusOr> ImageClassifier::Create( std::move(packets_callback)); } -absl::StatusOr ImageClassifier::Classify( - Image image, std::optional image_processing_options) { +absl::StatusOr ImageClassifier::Classify( + Image image, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, "GPU input images are currently not supported.", MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - NormalizedRect norm_rect = FillNormalizedRect(image_processing_options); + ASSIGN_OR_RETURN(NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options)); ASSIGN_OR_RETURN( auto output_packets, ProcessImageData( {{kImageInStreamName, MakePacket(std::move(image))}, {kNormRectName, MakePacket(std::move(norm_rect))}})); - return output_packets[kClassificationResultStreamName] - .Get(); + return ConvertToClassificationResult( + output_packets[kClassificationsStreamName].Get()); } -absl::StatusOr ImageClassifier::ClassifyForVideo( +absl::StatusOr ImageClassifier::ClassifyForVideo( Image image, int64 timestamp_ms, - std::optional image_processing_options) { + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, "GPU input images are currently not supported.", MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - NormalizedRect norm_rect = FillNormalizedRect(image_processing_options); + ASSIGN_OR_RETURN(NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options)); ASSIGN_OR_RETURN( auto output_packets, ProcessVideoData( @@ -200,20 +185,21 @@ absl::StatusOr ImageClassifier::ClassifyForVideo( {kNormRectName, MakePacket(std::move(norm_rect)) .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}})); - return output_packets[kClassificationResultStreamName] - .Get(); + return ConvertToClassificationResult( + output_packets[kClassificationsStreamName].Get()); } absl::Status ImageClassifier::ClassifyAsync( Image image, int64 timestamp_ms, - std::optional image_processing_options) { + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, "GPU input images are currently not supported.", MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - NormalizedRect norm_rect = FillNormalizedRect(image_processing_options); + ASSIGN_OR_RETURN(NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options)); return SendLiveStreamData( {{kImageInStreamName, MakePacket(std::move(image)) diff --git a/mediapipe/tasks/cc/vision/image_classifier/image_classifier.h b/mediapipe/tasks/cc/vision/image_classifier/image_classifier.h index 5dff06cc..9b0c376a 100644 --- a/mediapipe/tasks/cc/vision/image_classifier/image_classifier.h +++ b/mediapipe/tasks/cc/vision/image_classifier/image_classifier.h @@ -22,11 +22,11 @@ limitations under the License. #include "absl/status/statusor.h" #include "mediapipe/framework/formats/image.h" -#include "mediapipe/framework/formats/rect.pb.h" -#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h" +#include "mediapipe/tasks/cc/components/containers/classification_result.h" #include "mediapipe/tasks/cc/components/processors/classifier_options.h" #include "mediapipe/tasks/cc/core/base_options.h" #include "mediapipe/tasks/cc/vision/core/base_vision_task_api.h" +#include "mediapipe/tasks/cc/vision/core/image_processing_options.h" #include "mediapipe/tasks/cc/vision/core/running_mode.h" namespace mediapipe { @@ -34,6 +34,10 @@ namespace tasks { namespace vision { namespace image_classifier { +// Alias the shared ClassificationResult struct as result type. +using ImageClassifierResult = + ::mediapipe::tasks::components::containers::ClassificationResult; + // The options for configuring a Mediapipe image classifier task. struct ImageClassifierOptions { // Base options for configuring MediaPipe Tasks, such as specifying the model @@ -56,9 +60,8 @@ struct ImageClassifierOptions { // The user-defined result callback for processing live stream data. // The result callback should only be specified when the running mode is set // to RunningMode::LIVE_STREAM. - std::function, - const Image&, int64)> + std::function, const Image&, + int64)> result_callback = nullptr; }; @@ -109,12 +112,10 @@ class ImageClassifier : tasks::vision::core::BaseVisionTaskApi { // // The optional 'image_processing_options' parameter can be used to specify: // - the rotation to apply to the image before performing classification, by - // setting its 'rotation' field in radians (e.g. 'M_PI / 2' for a 90° - // anti-clockwise rotation). + // setting its 'rotation_degrees' field. // and/or // - the region-of-interest on which to perform classification, by setting its - // 'x_center', 'y_center', 'width' and 'height' fields. If none of these is - // set, they will automatically be set to cover the full image. + // 'region_of_interest' field. If not specified, the full image is used. // If both are specified, the crop around the region-of-interest is extracted // first, then the specified rotation is applied to the crop. // @@ -124,21 +125,19 @@ class ImageClassifier : tasks::vision::core::BaseVisionTaskApi { // The image can be of any size with format RGB or RGBA. // TODO: describe exact preprocessing steps once // YUVToImageCalculator is integrated. - absl::StatusOr Classify( + absl::StatusOr Classify( mediapipe::Image image, - std::optional image_processing_options = + std::optional image_processing_options = std::nullopt); // Performs image classification on the provided video frame. // // The optional 'image_processing_options' parameter can be used to specify: // - the rotation to apply to the image before performing classification, by - // setting its 'rotation' field in radians (e.g. 'M_PI / 2' for a 90° - // anti-clockwise rotation). + // setting its 'rotation_degrees' field. // and/or // - the region-of-interest on which to perform classification, by setting its - // 'x_center', 'y_center', 'width' and 'height' fields. If none of these is - // set, they will automatically be set to cover the full image. + // 'region_of_interest' field. If not specified, the full image is used. // If both are specified, the crop around the region-of-interest is extracted // first, then the specified rotation is applied to the crop. // @@ -148,22 +147,20 @@ class ImageClassifier : tasks::vision::core::BaseVisionTaskApi { // The image can be of any size with format RGB or RGBA. It's required to // provide the video frame's timestamp (in milliseconds). The input timestamps // must be monotonically increasing. - absl::StatusOr - ClassifyForVideo(mediapipe::Image image, int64 timestamp_ms, - std::optional - image_processing_options = std::nullopt); + absl::StatusOr ClassifyForVideo( + mediapipe::Image image, int64 timestamp_ms, + std::optional image_processing_options = + std::nullopt); // Sends live image data to image classification, and the results will be // available via the "result_callback" provided in the ImageClassifierOptions. // // The optional 'image_processing_options' parameter can be used to specify: // - the rotation to apply to the image before performing classification, by - // setting its 'rotation' field in radians (e.g. 'M_PI / 2' for a 90° - // anti-clockwise rotation). + // setting its 'rotation_degrees' field. // and/or // - the region-of-interest on which to perform classification, by setting its - // 'x_center', 'y_center', 'width' and 'height' fields. If none of these is - // set, they will automatically be set to cover the full image. + // 'region_of_interest' field. If not specified, the full image is used. // If both are specified, the crop around the region-of-interest is extracted // first, then the specified rotation is applied to the crop. // @@ -175,20 +172,17 @@ class ImageClassifier : tasks::vision::core::BaseVisionTaskApi { // sent to the object detector. The input timestamps must be monotonically // increasing. // - // The "result_callback" prvoides - // - The classification results as a ClassificationResult object. + // The "result_callback" provides: + // - The classification results as an ImageClassifierResult object. // - The const reference to the corresponding input image that the image // classifier runs on. Note that the const reference to the image will no // longer be valid when the callback returns. To access the image data // outside of the callback, callers need to make a copy of the image. // - The input timestamp in milliseconds. absl::Status ClassifyAsync(mediapipe::Image image, int64 timestamp_ms, - std::optional + std::optional image_processing_options = std::nullopt); - // TODO: add Classify() variants taking a region of interest as - // additional argument. - // Shuts down the ImageClassifier when all works are done. absl::Status Close() { return runner_->Close(); } }; diff --git a/mediapipe/tasks/cc/vision/image_classifier/image_classifier_graph.cc b/mediapipe/tasks/cc/vision/image_classifier/image_classifier_graph.cc index 9a0078c5..8fa1a0d2 100644 --- a/mediapipe/tasks/cc/vision/image_classifier/image_classifier_graph.cc +++ b/mediapipe/tasks/cc/vision/image_classifier/image_classifier_graph.cc @@ -48,6 +48,7 @@ using ::mediapipe::tasks::components::containers::proto::ClassificationResult; constexpr float kDefaultScoreThreshold = std::numeric_limits::lowest(); constexpr char kClassificationResultTag[] = "CLASSIFICATION_RESULT"; +constexpr char kClassificationsTag[] = "CLASSIFICATIONS"; constexpr char kImageTag[] = "IMAGE"; constexpr char kNormRectTag[] = "NORM_RECT"; constexpr char kTensorsTag[] = "TENSORS"; @@ -56,6 +57,7 @@ constexpr char kTensorsTag[] = "TENSORS"; // subgraph. struct ImageClassifierOutputStreams { Source classification_result; + Source classifications; Source image; }; @@ -71,17 +73,19 @@ struct ImageClassifierOutputStreams { // Describes region of image to perform classification on. // @Optional: rect covering the whole image is used if not specified. // Outputs: -// CLASSIFICATION_RESULT - ClassificationResult -// The aggregated classification result object has two dimensions: -// (classification head, classification category) +// CLASSIFICATIONS - ClassificationResult @Optional +// The classification results aggregated by classifier head. // IMAGE - Image // The image that object detection runs on. +// TODO: remove this output once Java API migration is over. +// CLASSIFICATION_RESULT - (DEPRECATED) ClassificationResult @Optional +// The aggregated classification result. // // Example: // node { // calculator: "mediapipe.tasks.vision.image_classifier.ImageClassifierGraph" // input_stream: "IMAGE:image_in" -// output_stream: "CLASSIFICATION_RESULT:classification_result_out" +// output_stream: "CLASSIFICATIONS:classifications_out" // output_stream: "IMAGE:image_out" // options { // [mediapipe.tasks.vision.image_classifier.proto.ImageClassifierGraphOptions.ext] @@ -115,6 +119,8 @@ class ImageClassifierGraph : public core::ModelTaskGraph { graph[Input::Optional(kNormRectTag)], graph)); output_streams.classification_result >> graph[Output(kClassificationResultTag)]; + output_streams.classifications >> + graph[Output(kClassificationsTag)]; output_streams.image >> graph[Output(kImageTag)]; return graph.GetConfig(); } @@ -138,8 +144,10 @@ class ImageClassifierGraph : public core::ModelTaskGraph { // stream. auto& preprocessing = graph.AddNode("mediapipe.tasks.components.ImagePreprocessingSubgraph"); + bool use_gpu = components::DetermineImagePreprocessingGpuBackend( + task_options.base_options().acceleration()); MP_RETURN_IF_ERROR(ConfigureImagePreprocessing( - model_resources, + model_resources, use_gpu, &preprocessing .GetOptions())); image_in >> preprocessing.In(kImageTag); @@ -168,6 +176,8 @@ class ImageClassifierGraph : public core::ModelTaskGraph { return ImageClassifierOutputStreams{ /*classification_result=*/postprocessing[Output( kClassificationResultTag)], + /*classifications=*/ + postprocessing[Output(kClassificationsTag)], /*image=*/preprocessing[Output(kImageTag)]}; } }; diff --git a/mediapipe/tasks/cc/vision/image_classifier/image_classifier_test.cc b/mediapipe/tasks/cc/vision/image_classifier/image_classifier_test.cc index 55830e52..1144e903 100644 --- a/mediapipe/tasks/cc/vision/image_classifier/image_classifier_test.cc +++ b/mediapipe/tasks/cc/vision/image_classifier/image_classifier_test.cc @@ -27,14 +27,15 @@ limitations under the License. #include "absl/strings/str_format.h" #include "mediapipe/framework/deps/file_path.h" #include "mediapipe/framework/formats/image.h" -#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/port/gmock.h" #include "mediapipe/framework/port/gtest.h" #include "mediapipe/framework/port/parse_text_proto.h" #include "mediapipe/framework/port/status_matchers.h" #include "mediapipe/tasks/cc/common.h" -#include "mediapipe/tasks/cc/components/containers/proto/category.pb.h" -#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h" +#include "mediapipe/tasks/cc/components/containers/category.h" +#include "mediapipe/tasks/cc/components/containers/classification_result.h" +#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_utils.h" #include "tensorflow/lite/core/api/op_resolver.h" @@ -49,9 +50,10 @@ namespace image_classifier { namespace { using ::mediapipe::file::JoinPath; -using ::mediapipe::tasks::components::containers::proto::ClassificationEntry; -using ::mediapipe::tasks::components::containers::proto::ClassificationResult; -using ::mediapipe::tasks::components::containers::proto::Classifications; +using ::mediapipe::tasks::components::containers::Category; +using ::mediapipe::tasks::components::containers::Classifications; +using ::mediapipe::tasks::components::containers::Rect; +using ::mediapipe::tasks::vision::core::ImageProcessingOptions; using ::testing::HasSubstr; using ::testing::Optional; @@ -62,83 +64,56 @@ constexpr char kMobileNetQuantizedWithMetadata[] = constexpr char kMobileNetQuantizedWithDummyScoreCalibration[] = "mobilenet_v1_0.25_224_quant_with_dummy_score_calibration.tflite"; -// Checks that the two provided `ClassificationResult` are equal, with a +// Checks that the two provided `ImageClassifierResult` are equal, with a // tolerancy on floating-point score to account for numerical instabilities. -void ExpectApproximatelyEqual(const ClassificationResult& actual, - const ClassificationResult& expected) { +void ExpectApproximatelyEqual(const ImageClassifierResult& actual, + const ImageClassifierResult& expected) { const float kPrecision = 1e-6; - ASSERT_EQ(actual.classifications_size(), expected.classifications_size()); - for (int i = 0; i < actual.classifications_size(); ++i) { - const Classifications& a = actual.classifications(i); - const Classifications& b = expected.classifications(i); - EXPECT_EQ(a.head_index(), b.head_index()); - EXPECT_EQ(a.head_name(), b.head_name()); - EXPECT_EQ(a.entries_size(), b.entries_size()); - for (int j = 0; j < a.entries_size(); ++j) { - const ClassificationEntry& x = a.entries(j); - const ClassificationEntry& y = b.entries(j); - EXPECT_EQ(x.timestamp_ms(), y.timestamp_ms()); - EXPECT_EQ(x.categories_size(), y.categories_size()); - for (int k = 0; k < x.categories_size(); ++k) { - EXPECT_EQ(x.categories(k).index(), y.categories(k).index()); - EXPECT_EQ(x.categories(k).category_name(), - y.categories(k).category_name()); - EXPECT_EQ(x.categories(k).display_name(), - y.categories(k).display_name()); - EXPECT_NEAR(x.categories(k).score(), y.categories(k).score(), - kPrecision); - } + ASSERT_EQ(actual.classifications.size(), expected.classifications.size()); + for (int i = 0; i < actual.classifications.size(); ++i) { + const Classifications& a = actual.classifications[i]; + const Classifications& b = expected.classifications[i]; + EXPECT_EQ(a.head_index, b.head_index); + EXPECT_EQ(a.head_name, b.head_name); + EXPECT_EQ(a.categories.size(), b.categories.size()); + for (int j = 0; j < a.categories.size(); ++j) { + const Category& x = a.categories[j]; + const Category& y = b.categories[j]; + EXPECT_EQ(x.index, y.index); + EXPECT_NEAR(x.score, y.score, kPrecision); + EXPECT_EQ(x.category_name, y.category_name); + EXPECT_EQ(x.display_name, y.display_name); } } } // Generates expected results for "burger.jpg" using kMobileNetFloatWithMetadata // with max_results set to 3. -ClassificationResult GenerateBurgerResults(int64 timestamp) { - return ParseTextProtoOrDie( - absl::StrFormat(R"pb(classifications { - entries { - categories { - index: 934 - score: 0.7939592 - category_name: "cheeseburger" - } - categories { - index: 932 - score: 0.027392805 - category_name: "bagel" - } - categories { - index: 925 - score: 0.019340655 - category_name: "guacamole" - } - timestamp_ms: %d - } - head_index: 0 - head_name: "probability" - })pb", - timestamp)); +ImageClassifierResult GenerateBurgerResults() { + ImageClassifierResult result; + result.classifications.emplace_back(Classifications{ + /*categories=*/{ + {/*index=*/934, /*score=*/0.793959200, + /*category_name=*/"cheeseburger"}, + {/*index=*/932, /*score=*/0.027392805, /*category_name=*/"bagel"}, + {/*index=*/925, /*score=*/0.019340655, + /*category_name=*/"guacamole"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + return result; } // Generates expected results for "multi_objects.jpg" using // kMobileNetFloatWithMetadata with max_results set to 1 and the right bounding // box set around the soccer ball. -ClassificationResult GenerateSoccerBallResults(int64 timestamp) { - return ParseTextProtoOrDie( - absl::StrFormat(R"pb(classifications { - entries { - categories { - index: 806 - score: 0.996527493 - category_name: "soccer ball" - } - timestamp_ms: %d - } - head_index: 0 - head_name: "probability" - })pb", - timestamp)); +ImageClassifierResult GenerateSoccerBallResults() { + ImageClassifierResult result; + result.classifications.emplace_back( + Classifications{/*categories=*/{{/*index=*/806, /*score=*/0.996527493, + /*category_name=*/"soccer ball"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + return result; } // A custom OpResolver only containing the Ops required by the test model. @@ -257,7 +232,7 @@ TEST_F(CreateTest, FailsWithIllegalCallbackInImageOrVideoMode) { options->base_options.model_asset_path = JoinPath("./", kTestDataDirectory, kMobileNetQuantizedWithMetadata); options->running_mode = running_mode; - options->result_callback = [](absl::StatusOr, + options->result_callback = [](absl::StatusOr, const Image& image, int64 timestamp_ms) {}; auto image_classifier = ImageClassifier::Create(std::move(options)); @@ -333,7 +308,7 @@ TEST_F(ImageModeTest, SucceedsWithFloatModel) { MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify(image)); - ExpectApproximatelyEqual(results, GenerateBurgerResults(0)); + ExpectApproximatelyEqual(results, GenerateBurgerResults()); } TEST_F(ImageModeTest, SucceedsWithQuantizedModel) { @@ -352,19 +327,13 @@ TEST_F(ImageModeTest, SucceedsWithQuantizedModel) { MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify(image)); - ExpectApproximatelyEqual(results, ParseTextProtoOrDie( - R"pb(classifications { - entries { - categories { - index: 934 - score: 0.97265625 - category_name: "cheeseburger" - } - timestamp_ms: 0 - } - head_index: 0 - head_name: "probability" - })pb")); + ImageClassifierResult expected; + expected.classifications.emplace_back( + Classifications{/*categories=*/{{/*index=*/934, /*score=*/0.97265625, + /*category_name=*/"cheeseburger"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(results, expected); } TEST_F(ImageModeTest, SucceedsWithMaxResultsOption) { @@ -380,19 +349,13 @@ TEST_F(ImageModeTest, SucceedsWithMaxResultsOption) { MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify(image)); - ExpectApproximatelyEqual(results, ParseTextProtoOrDie( - R"pb(classifications { - entries { - categories { - index: 934 - score: 0.7939592 - category_name: "cheeseburger" - } - timestamp_ms: 0 - } - head_index: 0 - head_name: "probability" - })pb")); + ImageClassifierResult expected; + expected.classifications.emplace_back( + Classifications{/*categories=*/{{/*index=*/934, /*score=*/0.7939592, + /*category_name=*/"cheeseburger"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(results, expected); } TEST_F(ImageModeTest, SucceedsWithScoreThresholdOption) { @@ -408,24 +371,15 @@ TEST_F(ImageModeTest, SucceedsWithScoreThresholdOption) { MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify(image)); - ExpectApproximatelyEqual(results, ParseTextProtoOrDie( - R"pb(classifications { - entries { - categories { - index: 934 - score: 0.7939592 - category_name: "cheeseburger" - } - categories { - index: 932 - score: 0.027392805 - category_name: "bagel" - } - timestamp_ms: 0 - } - head_index: 0 - head_name: "probability" - })pb")); + ImageClassifierResult expected; + expected.classifications.emplace_back(Classifications{ + /*categories=*/{ + {/*index=*/934, /*score=*/0.7939592, + /*category_name=*/"cheeseburger"}, + {/*index=*/932, /*score=*/0.027392805, /*category_name=*/"bagel"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(results, expected); } TEST_F(ImageModeTest, SucceedsWithAllowlistOption) { @@ -442,29 +396,17 @@ TEST_F(ImageModeTest, SucceedsWithAllowlistOption) { MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify(image)); - ExpectApproximatelyEqual(results, ParseTextProtoOrDie( - R"pb(classifications { - entries { - categories { - index: 934 - score: 0.7939592 - category_name: "cheeseburger" - } - categories { - index: 925 - score: 0.019340655 - category_name: "guacamole" - } - categories { - index: 963 - score: 0.0063278517 - category_name: "meat loaf" - } - timestamp_ms: 0 - } - head_index: 0 - head_name: "probability" - })pb")); + ImageClassifierResult expected; + expected.classifications.emplace_back(Classifications{ + /*categories=*/{ + {/*index=*/934, /*score=*/0.7939592, + /*category_name=*/"cheeseburger"}, + {/*index=*/925, /*score=*/0.019340655, /*category_name=*/"guacamole"}, + {/*index=*/963, /*score=*/0.0063278517, + /*category_name=*/"meat loaf"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(results, expected); } TEST_F(ImageModeTest, SucceedsWithDenylistOption) { @@ -481,29 +423,17 @@ TEST_F(ImageModeTest, SucceedsWithDenylistOption) { MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify(image)); - ExpectApproximatelyEqual(results, ParseTextProtoOrDie( - R"pb(classifications { - entries { - categories { - index: 934 - score: 0.7939592 - category_name: "cheeseburger" - } - categories { - index: 925 - score: 0.019340655 - category_name: "guacamole" - } - categories { - index: 963 - score: 0.0063278517 - category_name: "meat loaf" - } - timestamp_ms: 0 - } - head_index: 0 - head_name: "probability" - })pb")); + ImageClassifierResult expected; + expected.classifications.emplace_back(Classifications{ + /*categories=*/{ + {/*index=*/934, /*score=*/0.7939592, + /*category_name=*/"cheeseburger"}, + {/*index=*/925, /*score=*/0.019340655, /*category_name=*/"guacamole"}, + {/*index=*/963, /*score=*/0.0063278517, + /*category_name=*/"meat loaf"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(results, expected); } TEST_F(ImageModeTest, SucceedsWithScoreCalibration) { @@ -522,19 +452,13 @@ TEST_F(ImageModeTest, SucceedsWithScoreCalibration) { MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify(image)); - ExpectApproximatelyEqual(results, ParseTextProtoOrDie( - R"pb(classifications { - entries { - categories { - index: 934 - score: 0.725648628 - category_name: "cheeseburger" - } - timestamp_ms: 0 - } - head_index: 0 - head_name: "probability" - })pb")); + ImageClassifierResult expected; + expected.classifications.emplace_back( + Classifications{/*categories=*/{{/*index=*/934, /*score=*/0.725648628, + /*category_name=*/"cheeseburger"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(results, expected); } TEST_F(ImageModeTest, SucceedsWithRegionOfInterest) { @@ -547,17 +471,14 @@ TEST_F(ImageModeTest, SucceedsWithRegionOfInterest) { options->classifier_options.max_results = 1; MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr image_classifier, ImageClassifier::Create(std::move(options))); - // Crop around the soccer ball. - NormalizedRect image_processing_options; - image_processing_options.set_x_center(0.532); - image_processing_options.set_y_center(0.521); - image_processing_options.set_width(0.164); - image_processing_options.set_height(0.427); + // Region-of-interest around the soccer ball. + Rect roi{/*left=*/0.45, /*top=*/0.3075, /*right=*/0.614, /*bottom=*/0.7345}; + ImageProcessingOptions image_processing_options{roi, /*rotation_degrees=*/0}; MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify( image, image_processing_options)); - ExpectApproximatelyEqual(results, GenerateSoccerBallResults(0)); + ExpectApproximatelyEqual(results, GenerateSoccerBallResults()); } TEST_F(ImageModeTest, SucceedsWithRotation) { @@ -572,8 +493,8 @@ TEST_F(ImageModeTest, SucceedsWithRotation) { ImageClassifier::Create(std::move(options))); // Specify a 90° anti-clockwise rotation. - NormalizedRect image_processing_options; - image_processing_options.set_rotation(M_PI / 2.0); + ImageProcessingOptions image_processing_options; + image_processing_options.rotation_degrees = -90; MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify( image, image_processing_options)); @@ -581,29 +502,17 @@ TEST_F(ImageModeTest, SucceedsWithRotation) { // Results differ slightly from the non-rotated image, but that's expected // as models are very sensitive to the slightest numerical differences // introduced by the rotation and JPG encoding. - ExpectApproximatelyEqual(results, ParseTextProtoOrDie( - R"pb(classifications { - entries { - categories { - index: 934 - score: 0.6371766 - category_name: "cheeseburger" - } - categories { - index: 963 - score: 0.049443405 - category_name: "meat loaf" - } - categories { - index: 925 - score: 0.047918003 - category_name: "guacamole" - } - timestamp_ms: 0 - } - head_index: 0 - head_name: "probability" - })pb")); + ImageClassifierResult expected; + expected.classifications.emplace_back(Classifications{ + /*categories=*/{ + {/*index=*/934, /*score=*/0.6371766, + /*category_name=*/"cheeseburger"}, + {/*index=*/963, /*score=*/0.049443405, /*category_name=*/"meat loaf"}, + {/*index=*/925, /*score=*/0.047918003, + /*category_name=*/"guacamole"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(results, expected); } TEST_F(ImageModeTest, SucceedsWithRegionOfInterestAndRotation) { @@ -616,31 +525,84 @@ TEST_F(ImageModeTest, SucceedsWithRegionOfInterestAndRotation) { options->classifier_options.max_results = 1; MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr image_classifier, ImageClassifier::Create(std::move(options))); - // Crop around the chair, with 90° anti-clockwise rotation. - NormalizedRect image_processing_options; - image_processing_options.set_x_center(0.2821); - image_processing_options.set_y_center(0.2406); - image_processing_options.set_width(0.5642); - image_processing_options.set_height(0.1286); - image_processing_options.set_rotation(M_PI / 2.0); + // Region-of-interest around the chair, with 90° anti-clockwise rotation. + Rect roi{/*left=*/0.006, /*top=*/0.1763, /*right=*/0.5702, /*bottom=*/0.3049}; + ImageProcessingOptions image_processing_options{roi, + /*rotation_degrees=*/-90}; MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->Classify( image, image_processing_options)); - ExpectApproximatelyEqual(results, - ParseTextProtoOrDie( - R"pb(classifications { - entries { - categories { - index: 560 - score: 0.6800408 - category_name: "folding chair" - } - timestamp_ms: 0 - } - head_index: 0 - head_name: "probability" - })pb")); + ImageClassifierResult expected; + expected.classifications.emplace_back( + Classifications{/*categories=*/{{/*index=*/560, /*score=*/0.6522213, + /*category_name=*/"folding chair"}}, + /*head_index=*/0, + /*head_name=*/"probability"}); + ExpectApproximatelyEqual(results, expected); +} + +// Testing all these once with ImageClassifier. +TEST_F(ImageModeTest, FailsWithInvalidImageProcessingOptions) { + MP_ASSERT_OK_AND_ASSIGN(Image image, + DecodeImageFromFile(JoinPath("./", kTestDataDirectory, + "multi_objects.jpg"))); + auto options = std::make_unique(); + options->base_options.model_asset_path = + JoinPath("./", kTestDataDirectory, kMobileNetFloatWithMetadata); + MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr image_classifier, + ImageClassifier::Create(std::move(options))); + + // Invalid: left > right. + Rect roi{/*left=*/0.9, /*top=*/0, /*right=*/0.1, /*bottom=*/1}; + ImageProcessingOptions image_processing_options{roi, + /*rotation_degrees=*/0}; + auto results = image_classifier->Classify(image, image_processing_options); + EXPECT_EQ(results.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(results.status().message(), + HasSubstr("Expected Rect with left < right and top < bottom")); + EXPECT_THAT( + results.status().GetPayload(kMediaPipeTasksPayload), + Optional(absl::Cord(absl::StrCat( + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError)))); + + // Invalid: top > bottom. + roi = {/*left=*/0, /*top=*/0.9, /*right=*/1, /*bottom=*/0.1}; + image_processing_options = {roi, + /*rotation_degrees=*/0}; + results = image_classifier->Classify(image, image_processing_options); + EXPECT_EQ(results.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(results.status().message(), + HasSubstr("Expected Rect with left < right and top < bottom")); + EXPECT_THAT( + results.status().GetPayload(kMediaPipeTasksPayload), + Optional(absl::Cord(absl::StrCat( + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError)))); + + // Invalid: coordinates out of [0,1] range. + roi = {/*left=*/-0.1, /*top=*/0, /*right=*/1, /*bottom=*/1}; + image_processing_options = {roi, + /*rotation_degrees=*/0}; + results = image_classifier->Classify(image, image_processing_options); + EXPECT_EQ(results.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(results.status().message(), + HasSubstr("Expected Rect values to be in [0,1]")); + EXPECT_THAT( + results.status().GetPayload(kMediaPipeTasksPayload), + Optional(absl::Cord(absl::StrCat( + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError)))); + + // Invalid: rotation not a multiple of 90°. + image_processing_options = {/*region_of_interest=*/std::nullopt, + /*rotation_degrees=*/1}; + results = image_classifier->Classify(image, image_processing_options); + EXPECT_EQ(results.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(results.status().message(), + HasSubstr("Expected rotation to be a multiple of 90°")); + EXPECT_THAT( + results.status().GetPayload(kMediaPipeTasksPayload), + Optional(absl::Cord(absl::StrCat( + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError)))); } class VideoModeTest : public tflite_shims::testing::Test {}; @@ -714,7 +676,7 @@ TEST_F(VideoModeTest, Succeeds) { for (int i = 0; i < iterations; ++i) { MP_ASSERT_OK_AND_ASSIGN(auto results, image_classifier->ClassifyForVideo(image, i)); - ExpectApproximatelyEqual(results, GenerateBurgerResults(i)); + ExpectApproximatelyEqual(results, GenerateBurgerResults()); } MP_ASSERT_OK(image_classifier->Close()); } @@ -732,17 +694,15 @@ TEST_F(VideoModeTest, SucceedsWithRegionOfInterest) { MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr image_classifier, ImageClassifier::Create(std::move(options))); // Crop around the soccer ball. - NormalizedRect image_processing_options; - image_processing_options.set_x_center(0.532); - image_processing_options.set_y_center(0.521); - image_processing_options.set_width(0.164); - image_processing_options.set_height(0.427); + // Region-of-interest around the soccer ball. + Rect roi{/*left=*/0.45, /*top=*/0.3075, /*right=*/0.614, /*bottom=*/0.7345}; + ImageProcessingOptions image_processing_options{roi, /*rotation_degrees=*/0}; for (int i = 0; i < iterations; ++i) { MP_ASSERT_OK_AND_ASSIGN( auto results, image_classifier->ClassifyForVideo(image, i, image_processing_options)); - ExpectApproximatelyEqual(results, GenerateSoccerBallResults(i)); + ExpectApproximatelyEqual(results, GenerateSoccerBallResults()); } MP_ASSERT_OK(image_classifier->Close()); } @@ -757,7 +717,7 @@ TEST_F(LiveStreamModeTest, FailsWithCallingWrongMethod) { options->base_options.model_asset_path = JoinPath("./", kTestDataDirectory, kMobileNetFloatWithMetadata); options->running_mode = core::RunningMode::LIVE_STREAM; - options->result_callback = [](absl::StatusOr, + options->result_callback = [](absl::StatusOr, const Image& image, int64 timestamp_ms) {}; MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr image_classifier, ImageClassifier::Create(std::move(options))); @@ -788,7 +748,7 @@ TEST_F(LiveStreamModeTest, FailsWithOutOfOrderInputTimestamps) { options->base_options.model_asset_path = JoinPath("./", kTestDataDirectory, kMobileNetFloatWithMetadata); options->running_mode = core::RunningMode::LIVE_STREAM; - options->result_callback = [](absl::StatusOr, + options->result_callback = [](absl::StatusOr, const Image& image, int64 timestamp_ms) {}; MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr image_classifier, ImageClassifier::Create(std::move(options))); @@ -806,7 +766,7 @@ TEST_F(LiveStreamModeTest, FailsWithOutOfOrderInputTimestamps) { } struct LiveStreamModeResults { - ClassificationResult classification_result; + ImageClassifierResult classification_result; std::pair image_size; int64 timestamp_ms; }; @@ -823,7 +783,7 @@ TEST_F(LiveStreamModeTest, Succeeds) { options->running_mode = core::RunningMode::LIVE_STREAM; options->classifier_options.max_results = 3; options->result_callback = - [&results](absl::StatusOr classification_result, + [&results](absl::StatusOr classification_result, const Image& image, int64 timestamp_ms) { MP_ASSERT_OK(classification_result.status()); results.push_back( @@ -850,7 +810,7 @@ TEST_F(LiveStreamModeTest, Succeeds) { EXPECT_EQ(result.image_size.first, image.width()); EXPECT_EQ(result.image_size.second, image.height()); ExpectApproximatelyEqual(result.classification_result, - GenerateBurgerResults(timestamp_ms)); + GenerateBurgerResults()); } } @@ -866,7 +826,7 @@ TEST_F(LiveStreamModeTest, SucceedsWithRegionOfInterest) { options->running_mode = core::RunningMode::LIVE_STREAM; options->classifier_options.max_results = 1; options->result_callback = - [&results](absl::StatusOr classification_result, + [&results](absl::StatusOr classification_result, const Image& image, int64 timestamp_ms) { MP_ASSERT_OK(classification_result.status()); results.push_back( @@ -877,11 +837,8 @@ TEST_F(LiveStreamModeTest, SucceedsWithRegionOfInterest) { MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr image_classifier, ImageClassifier::Create(std::move(options))); // Crop around the soccer ball. - NormalizedRect image_processing_options; - image_processing_options.set_x_center(0.532); - image_processing_options.set_y_center(0.521); - image_processing_options.set_width(0.164); - image_processing_options.set_height(0.427); + Rect roi{/*left=*/0.45, /*top=*/0.3075, /*right=*/0.614, /*bottom=*/0.7345}; + ImageProcessingOptions image_processing_options{roi, /*rotation_degrees=*/0}; for (int i = 0; i < iterations; ++i) { MP_ASSERT_OK( @@ -900,7 +857,7 @@ TEST_F(LiveStreamModeTest, SucceedsWithRegionOfInterest) { EXPECT_EQ(result.image_size.first, image.width()); EXPECT_EQ(result.image_size.second, image.height()); ExpectApproximatelyEqual(result.classification_result, - GenerateSoccerBallResults(timestamp_ms)); + GenerateSoccerBallResults()); } } diff --git a/mediapipe/tasks/cc/vision/image_embedder/BUILD b/mediapipe/tasks/cc/vision/image_embedder/BUILD index e619b8d1..0f63f87e 100644 --- a/mediapipe/tasks/cc/vision/image_embedder/BUILD +++ b/mediapipe/tasks/cc/vision/image_embedder/BUILD @@ -58,6 +58,7 @@ cc_library( "//mediapipe/tasks/cc/core:utils", "//mediapipe/tasks/cc/core/proto:base_options_cc_proto", "//mediapipe/tasks/cc/vision/core:base_vision_task_api", + "//mediapipe/tasks/cc/vision/core:image_processing_options", "//mediapipe/tasks/cc/vision/core:running_mode", "//mediapipe/tasks/cc/vision/core:vision_task_api_factory", "//mediapipe/tasks/cc/vision/image_embedder/proto:image_embedder_graph_options_cc_proto", diff --git a/mediapipe/tasks/cc/vision/image_embedder/image_embedder.cc b/mediapipe/tasks/cc/vision/image_embedder/image_embedder.cc index 24fd2862..1dc31630 100644 --- a/mediapipe/tasks/cc/vision/image_embedder/image_embedder.cc +++ b/mediapipe/tasks/cc/vision/image_embedder/image_embedder.cc @@ -29,6 +29,7 @@ limitations under the License. #include "mediapipe/tasks/cc/core/proto/base_options.pb.h" #include "mediapipe/tasks/cc/core/task_runner.h" #include "mediapipe/tasks/cc/core/utils.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/core/vision_task_api_factory.h" #include "mediapipe/tasks/cc/vision/image_embedder/proto/image_embedder_graph_options.pb.h" @@ -58,16 +59,6 @@ using ::mediapipe::tasks::core::PacketMap; using ::mediapipe::tasks::vision::image_embedder::proto:: ImageEmbedderGraphOptions; -// Builds a NormalizedRect covering the entire image. -NormalizedRect BuildFullImageNormRect() { - NormalizedRect norm_rect; - norm_rect.set_x_center(0.5); - norm_rect.set_y_center(0.5); - norm_rect.set_width(1); - norm_rect.set_height(1); - return norm_rect; -} - // Creates a MediaPipe graph config that contains a single node of type // "mediapipe.tasks.vision.image_embedder.ImageEmbedderGraph". If the task is // running in the live stream mode, a "FlowLimiterCalculator" will be added to @@ -148,15 +139,16 @@ absl::StatusOr> ImageEmbedder::Create( } absl::StatusOr ImageEmbedder::Embed( - Image image, std::optional roi) { + Image image, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, "GPU input images are currently not supported.", MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - NormalizedRect norm_rect = - roi.has_value() ? roi.value() : BuildFullImageNormRect(); + ASSIGN_OR_RETURN(NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options)); ASSIGN_OR_RETURN( auto output_packets, ProcessImageData( @@ -167,15 +159,16 @@ absl::StatusOr ImageEmbedder::Embed( } absl::StatusOr ImageEmbedder::EmbedForVideo( - Image image, int64 timestamp_ms, std::optional roi) { + Image image, int64 timestamp_ms, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, "GPU input images are currently not supported.", MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - NormalizedRect norm_rect = - roi.has_value() ? roi.value() : BuildFullImageNormRect(); + ASSIGN_OR_RETURN(NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options)); ASSIGN_OR_RETURN( auto output_packets, ProcessVideoData( @@ -188,16 +181,17 @@ absl::StatusOr ImageEmbedder::EmbedForVideo( return output_packets[kEmbeddingResultStreamName].Get(); } -absl::Status ImageEmbedder::EmbedAsync(Image image, int64 timestamp_ms, - std::optional roi) { +absl::Status ImageEmbedder::EmbedAsync( + Image image, int64 timestamp_ms, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, "GPU input images are currently not supported.", MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - NormalizedRect norm_rect = - roi.has_value() ? roi.value() : BuildFullImageNormRect(); + ASSIGN_OR_RETURN(NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options)); return SendLiveStreamData( {{kImageInStreamName, MakePacket(std::move(image)) diff --git a/mediapipe/tasks/cc/vision/image_embedder/image_embedder.h b/mediapipe/tasks/cc/vision/image_embedder/image_embedder.h index 13f4702d..3a2a1dbe 100644 --- a/mediapipe/tasks/cc/vision/image_embedder/image_embedder.h +++ b/mediapipe/tasks/cc/vision/image_embedder/image_embedder.h @@ -21,11 +21,11 @@ limitations under the License. #include "absl/status/statusor.h" #include "mediapipe/framework/formats/image.h" -#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h" #include "mediapipe/tasks/cc/components/embedder_options.h" #include "mediapipe/tasks/cc/core/base_options.h" #include "mediapipe/tasks/cc/vision/core/base_vision_task_api.h" +#include "mediapipe/tasks/cc/vision/core/image_processing_options.h" #include "mediapipe/tasks/cc/vision/core/running_mode.h" namespace mediapipe { @@ -88,9 +88,17 @@ class ImageEmbedder : core::BaseVisionTaskApi { static absl::StatusOr> Create( std::unique_ptr options); - // Performs embedding extraction on the provided single image. Extraction - // is performed on the region of interest specified by the `roi` argument if - // provided, or on the entire image otherwise. + // Performs embedding extraction on the provided single image. + // + // The optional 'image_processing_options' parameter can be used to specify: + // - the rotation to apply to the image before performing embedding + // extraction, by setting its 'rotation_degrees' field. + // and/or + // - the region-of-interest on which to perform embedding extraction, by + // setting its 'region_of_interest' field. If not specified, the full image + // is used. + // If both are specified, the crop around the region-of-interest is extracted + // first, then the specified rotation is applied to the crop. // // Only use this method when the ImageEmbedder is created with the image // running mode. @@ -98,11 +106,20 @@ class ImageEmbedder : core::BaseVisionTaskApi { // The image can be of any size with format RGB or RGBA. absl::StatusOr Embed( mediapipe::Image image, - std::optional roi = std::nullopt); + std::optional image_processing_options = + std::nullopt); - // Performs embedding extraction on the provided video frame. Extraction - // is performed on the region of interested specified by the `roi` argument if - // provided, or on the entire image otherwise. + // Performs embedding extraction on the provided video frame. + // + // The optional 'image_processing_options' parameter can be used to specify: + // - the rotation to apply to the image before performing embedding + // extraction, by setting its 'rotation_degrees' field. + // and/or + // - the region-of-interest on which to perform embedding extraction, by + // setting its 'region_of_interest' field. If not specified, the full image + // is used. + // If both are specified, the crop around the region-of-interest is extracted + // first, then the specified rotation is applied to the crop. // // Only use this method when the ImageEmbedder is created with the video // running mode. @@ -112,12 +129,21 @@ class ImageEmbedder : core::BaseVisionTaskApi { // must be monotonically increasing. absl::StatusOr EmbedForVideo( mediapipe::Image image, int64 timestamp_ms, - std::optional roi = std::nullopt); + std::optional image_processing_options = + std::nullopt); // Sends live image data to embedder, and the results will be available via - // the "result_callback" provided in the ImageEmbedderOptions. Embedding - // extraction is performed on the region of interested specified by the `roi` - // argument if provided, or on the entire image otherwise. + // the "result_callback" provided in the ImageEmbedderOptions. + // + // The optional 'image_processing_options' parameter can be used to specify: + // - the rotation to apply to the image before performing embedding + // extraction, by setting its 'rotation_degrees' field. + // and/or + // - the region-of-interest on which to perform embedding extraction, by + // setting its 'region_of_interest' field. If not specified, the full image + // is used. + // If both are specified, the crop around the region-of-interest is extracted + // first, then the specified rotation is applied to the crop. // // Only use this method when the ImageEmbedder is created with the live // stream running mode. @@ -135,9 +161,9 @@ class ImageEmbedder : core::BaseVisionTaskApi { // longer be valid when the callback returns. To access the image data // outside of the callback, callers need to make a copy of the image. // - The input timestamp in milliseconds. - absl::Status EmbedAsync( - mediapipe::Image image, int64 timestamp_ms, - std::optional roi = std::nullopt); + absl::Status EmbedAsync(mediapipe::Image image, int64 timestamp_ms, + std::optional + image_processing_options = std::nullopt); // Shuts down the ImageEmbedder when all works are done. absl::Status Close() { return runner_->Close(); } diff --git a/mediapipe/tasks/cc/vision/image_embedder/image_embedder_graph.cc b/mediapipe/tasks/cc/vision/image_embedder/image_embedder_graph.cc index fff0f436..f0f44098 100644 --- a/mediapipe/tasks/cc/vision/image_embedder/image_embedder_graph.cc +++ b/mediapipe/tasks/cc/vision/image_embedder/image_embedder_graph.cc @@ -134,8 +134,10 @@ class ImageEmbedderGraph : public core::ModelTaskGraph { // stream. auto& preprocessing = graph.AddNode("mediapipe.tasks.components.ImagePreprocessingSubgraph"); + bool use_gpu = components::DetermineImagePreprocessingGpuBackend( + task_options.base_options().acceleration()); MP_RETURN_IF_ERROR(ConfigureImagePreprocessing( - model_resources, + model_resources, use_gpu, &preprocessing .GetOptions())); image_in >> preprocessing.In(kImageTag); diff --git a/mediapipe/tasks/cc/vision/image_embedder/image_embedder_test.cc b/mediapipe/tasks/cc/vision/image_embedder/image_embedder_test.cc index db1019b3..386b6c8e 100644 --- a/mediapipe/tasks/cc/vision/image_embedder/image_embedder_test.cc +++ b/mediapipe/tasks/cc/vision/image_embedder/image_embedder_test.cc @@ -23,7 +23,6 @@ limitations under the License. #include "absl/status/statusor.h" #include "mediapipe/framework/deps/file_path.h" #include "mediapipe/framework/formats/image.h" -#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/port/gmock.h" #include "mediapipe/framework/port/gtest.h" #include "mediapipe/framework/port/status_matchers.h" @@ -42,7 +41,9 @@ namespace image_embedder { namespace { using ::mediapipe::file::JoinPath; +using ::mediapipe::tasks::components::containers::Rect; using ::mediapipe::tasks::components::containers::proto::EmbeddingResult; +using ::mediapipe::tasks::vision::core::ImageProcessingOptions; using ::testing::HasSubstr; using ::testing::Optional; @@ -326,16 +327,14 @@ TEST_F(ImageModeTest, SucceedsWithRegionOfInterest) { MP_ASSERT_OK_AND_ASSIGN( Image crop, DecodeImageFromFile( JoinPath("./", kTestDataDirectory, "burger_crop.jpg"))); - // Bounding box in "burger.jpg" corresponding to "burger_crop.jpg". - NormalizedRect roi; - roi.set_x_center(200.0 / 480); - roi.set_y_center(0.5); - roi.set_width(400.0 / 480); - roi.set_height(1.0f); + // Region-of-interest in "burger.jpg" corresponding to "burger_crop.jpg". + Rect roi{/*left=*/0, /*top=*/0, /*right=*/0.833333, /*bottom=*/1}; + ImageProcessingOptions image_processing_options{roi, /*rotation_degrees=*/0}; // Extract both embeddings. - MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& image_result, - image_embedder->Embed(image, roi)); + MP_ASSERT_OK_AND_ASSIGN( + const EmbeddingResult& image_result, + image_embedder->Embed(image, image_processing_options)); MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& crop_result, image_embedder->Embed(crop)); @@ -351,6 +350,77 @@ TEST_F(ImageModeTest, SucceedsWithRegionOfInterest) { EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy); } +TEST_F(ImageModeTest, SucceedsWithRotation) { + auto options = std::make_unique(); + options->base_options.model_asset_path = + JoinPath("./", kTestDataDirectory, kMobileNetV3Embedder); + MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr image_embedder, + ImageEmbedder::Create(std::move(options))); + // Load images: one is a rotated version of the other. + MP_ASSERT_OK_AND_ASSIGN( + Image image, + DecodeImageFromFile(JoinPath("./", kTestDataDirectory, "burger.jpg"))); + MP_ASSERT_OK_AND_ASSIGN(Image rotated, + DecodeImageFromFile(JoinPath("./", kTestDataDirectory, + "burger_rotated.jpg"))); + ImageProcessingOptions image_processing_options; + image_processing_options.rotation_degrees = -90; + + // Extract both embeddings. + MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& image_result, + image_embedder->Embed(image)); + MP_ASSERT_OK_AND_ASSIGN( + const EmbeddingResult& rotated_result, + image_embedder->Embed(rotated, image_processing_options)); + + // Check results. + CheckMobileNetV3Result(image_result, false); + CheckMobileNetV3Result(rotated_result, false); + // CheckCosineSimilarity. + MP_ASSERT_OK_AND_ASSIGN( + double similarity, + ImageEmbedder::CosineSimilarity(image_result.embeddings(0).entries(0), + rotated_result.embeddings(0).entries(0))); + double expected_similarity = 0.572265; + EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy); +} + +TEST_F(ImageModeTest, SucceedsWithRegionOfInterestAndRotation) { + auto options = std::make_unique(); + options->base_options.model_asset_path = + JoinPath("./", kTestDataDirectory, kMobileNetV3Embedder); + MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr image_embedder, + ImageEmbedder::Create(std::move(options))); + MP_ASSERT_OK_AND_ASSIGN( + Image crop, DecodeImageFromFile( + JoinPath("./", kTestDataDirectory, "burger_crop.jpg"))); + MP_ASSERT_OK_AND_ASSIGN(Image rotated, + DecodeImageFromFile(JoinPath("./", kTestDataDirectory, + "burger_rotated.jpg"))); + // Region-of-interest corresponding to burger_crop.jpg. + Rect roi{/*left=*/0, /*top=*/0, /*right=*/1, /*bottom=*/0.8333333}; + ImageProcessingOptions image_processing_options{roi, + /*rotation_degrees=*/-90}; + + // Extract both embeddings. + MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& crop_result, + image_embedder->Embed(crop)); + MP_ASSERT_OK_AND_ASSIGN( + const EmbeddingResult& rotated_result, + image_embedder->Embed(rotated, image_processing_options)); + + // Check results. + CheckMobileNetV3Result(crop_result, false); + CheckMobileNetV3Result(rotated_result, false); + // CheckCosineSimilarity. + MP_ASSERT_OK_AND_ASSIGN( + double similarity, + ImageEmbedder::CosineSimilarity(crop_result.embeddings(0).entries(0), + rotated_result.embeddings(0).entries(0))); + double expected_similarity = 0.62838; + EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy); +} + class VideoModeTest : public tflite_shims::testing::Test {}; TEST_F(VideoModeTest, FailsWithCallingWrongMethod) { diff --git a/mediapipe/tasks/cc/vision/image_segmenter/BUILD b/mediapipe/tasks/cc/vision/image_segmenter/BUILD index 6bdbf41d..81cd43e3 100644 --- a/mediapipe/tasks/cc/vision/image_segmenter/BUILD +++ b/mediapipe/tasks/cc/vision/image_segmenter/BUILD @@ -24,10 +24,12 @@ cc_library( ":image_segmenter_graph", "//mediapipe/framework/api2:builder", "//mediapipe/framework/formats:image", + "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/tasks/cc/components/proto:segmenter_options_cc_proto", "//mediapipe/tasks/cc/core:base_options", "//mediapipe/tasks/cc/core:utils", "//mediapipe/tasks/cc/vision/core:base_vision_task_api", + "//mediapipe/tasks/cc/vision/core:image_processing_options", "//mediapipe/tasks/cc/vision/core:running_mode", "//mediapipe/tasks/cc/vision/core:vision_task_api_factory", "//mediapipe/tasks/cc/vision/image_segmenter/proto:image_segmenter_options_cc_proto", @@ -48,6 +50,7 @@ cc_library( "//mediapipe/framework/api2:builder", "//mediapipe/framework/api2:port", "//mediapipe/framework/formats:image", + "//mediapipe/framework/formats:rect_cc_proto", "//mediapipe/framework/port:status", "//mediapipe/tasks/cc:common", "//mediapipe/tasks/cc/components:image_preprocessing", diff --git a/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter.cc b/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter.cc index 84ceea88..209ee0df 100644 --- a/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter.cc +++ b/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter.cc @@ -17,8 +17,10 @@ limitations under the License. #include "mediapipe/framework/api2/builder.h" #include "mediapipe/framework/formats/image.h" +#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/tasks/cc/components/proto/segmenter_options.pb.h" #include "mediapipe/tasks/cc/core/utils.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/core/vision_task_api_factory.h" @@ -32,6 +34,8 @@ constexpr char kGroupedSegmentationTag[] = "GROUPED_SEGMENTATION"; constexpr char kImageInStreamName[] = "image_in"; constexpr char kImageOutStreamName[] = "image_out"; constexpr char kImageTag[] = "IMAGE"; +constexpr char kNormRectStreamName[] = "norm_rect_in"; +constexpr char kNormRectTag[] = "NORM_RECT"; constexpr char kSubgraphTypeName[] = "mediapipe.tasks.vision.ImageSegmenterGraph"; constexpr int kMicroSecondsPerMilliSecond = 1000; @@ -51,15 +55,18 @@ CalculatorGraphConfig CreateGraphConfig( auto& task_subgraph = graph.AddNode(kSubgraphTypeName); task_subgraph.GetOptions().Swap(options.get()); graph.In(kImageTag).SetName(kImageInStreamName); + graph.In(kNormRectTag).SetName(kNormRectStreamName); task_subgraph.Out(kGroupedSegmentationTag).SetName(kSegmentationStreamName) >> graph.Out(kGroupedSegmentationTag); task_subgraph.Out(kImageTag).SetName(kImageOutStreamName) >> graph.Out(kImageTag); if (enable_flow_limiting) { - return tasks::core::AddFlowLimiterCalculator( - graph, task_subgraph, {kImageTag}, kGroupedSegmentationTag); + return tasks::core::AddFlowLimiterCalculator(graph, task_subgraph, + {kImageTag, kNormRectTag}, + kGroupedSegmentationTag); } graph.In(kImageTag) >> task_subgraph.In(kImageTag); + graph.In(kNormRectTag) >> task_subgraph.In(kNormRectTag); return graph.GetConfig(); } @@ -139,47 +146,68 @@ absl::StatusOr> ImageSegmenter::Create( } absl::StatusOr> ImageSegmenter::Segment( - mediapipe::Image image) { + mediapipe::Image image, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, absl::StrCat("GPU input images are currently not supported."), MediaPipeTasksStatus::kRunnerUnexpectedInputError); } + ASSIGN_OR_RETURN( + NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false)); ASSIGN_OR_RETURN( auto output_packets, - ProcessImageData({{kImageInStreamName, - mediapipe::MakePacket(std::move(image))}})); + ProcessImageData( + {{kImageInStreamName, mediapipe::MakePacket(std::move(image))}, + {kNormRectStreamName, + MakePacket(std::move(norm_rect))}})); return output_packets[kSegmentationStreamName].Get>(); } absl::StatusOr> ImageSegmenter::SegmentForVideo( - mediapipe::Image image, int64 timestamp_ms) { + mediapipe::Image image, int64 timestamp_ms, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, absl::StrCat("GPU input images are currently not supported."), MediaPipeTasksStatus::kRunnerUnexpectedInputError); } + ASSIGN_OR_RETURN( + NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false)); ASSIGN_OR_RETURN( auto output_packets, ProcessVideoData( {{kImageInStreamName, MakePacket(std::move(image)) + .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}, + {kNormRectStreamName, + MakePacket(std::move(norm_rect)) .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}})); return output_packets[kSegmentationStreamName].Get>(); } -absl::Status ImageSegmenter::SegmentAsync(Image image, int64 timestamp_ms) { +absl::Status ImageSegmenter::SegmentAsync( + Image image, int64 timestamp_ms, + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, absl::StrCat("GPU input images are currently not supported."), MediaPipeTasksStatus::kRunnerUnexpectedInputError); } + ASSIGN_OR_RETURN( + NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false)); return SendLiveStreamData( {{kImageInStreamName, MakePacket(std::move(image)) + .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}, + {kNormRectStreamName, + MakePacket(std::move(norm_rect)) .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}}); } diff --git a/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter.h b/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter.h index e2734c4e..54269ec0 100644 --- a/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter.h +++ b/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter.h @@ -25,6 +25,7 @@ limitations under the License. #include "mediapipe/framework/formats/image.h" #include "mediapipe/tasks/cc/core/base_options.h" #include "mediapipe/tasks/cc/vision/core/base_vision_task_api.h" +#include "mediapipe/tasks/cc/vision/core/image_processing_options.h" #include "mediapipe/tasks/cc/vision/image_segmenter/proto/image_segmenter_options.pb.h" #include "tensorflow/lite/kernels/register.h" @@ -116,14 +117,21 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi { // running mode. // // The image can be of any size with format RGB or RGBA. - // TODO: Describes how the input image will be preprocessed - // after the yuv support is implemented. + // + // The optional 'image_processing_options' parameter can be used to specify + // the rotation to apply to the image before performing segmentation, by + // setting its 'rotation_degrees' field. Note that specifying a + // region-of-interest using the 'region_of_interest' field is NOT supported + // and will result in an invalid argument error being returned. // // If the output_type is CATEGORY_MASK, the returned vector of images is // per-category segmented image mask. // If the output_type is CONFIDENCE_MASK, the returned vector of images // contains only one confidence image mask. - absl::StatusOr> Segment(mediapipe::Image image); + absl::StatusOr> Segment( + mediapipe::Image image, + std::optional image_processing_options = + std::nullopt); // Performs image segmentation on the provided video frame. // Only use this method when the ImageSegmenter is created with the video @@ -133,12 +141,20 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi { // provide the video frame's timestamp (in milliseconds). The input timestamps // must be monotonically increasing. // + // The optional 'image_processing_options' parameter can be used to specify + // the rotation to apply to the image before performing segmentation, by + // setting its 'rotation_degrees' field. Note that specifying a + // region-of-interest using the 'region_of_interest' field is NOT supported + // and will result in an invalid argument error being returned. + // // If the output_type is CATEGORY_MASK, the returned vector of images is // per-category segmented image mask. // If the output_type is CONFIDENCE_MASK, the returned vector of images // contains only one confidence image mask. absl::StatusOr> SegmentForVideo( - mediapipe::Image image, int64 timestamp_ms); + mediapipe::Image image, int64 timestamp_ms, + std::optional image_processing_options = + std::nullopt); // Sends live image data to perform image segmentation, and the results will // be available via the "result_callback" provided in the @@ -150,6 +166,12 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi { // sent to the image segmenter. The input timestamps must be monotonically // increasing. // + // The optional 'image_processing_options' parameter can be used to specify + // the rotation to apply to the image before performing segmentation, by + // setting its 'rotation_degrees' field. Note that specifying a + // region-of-interest using the 'region_of_interest' field is NOT supported + // and will result in an invalid argument error being returned. + // // The "result_callback" prvoides // - A vector of segmented image masks. // If the output_type is CATEGORY_MASK, the returned vector of images is @@ -161,7 +183,9 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi { // no longer be valid when the callback returns. To access the image data // outside of the callback, callers need to make a copy of the image. // - The input timestamp in milliseconds. - absl::Status SegmentAsync(mediapipe::Image image, int64 timestamp_ms); + absl::Status SegmentAsync(mediapipe::Image image, int64 timestamp_ms, + std::optional + image_processing_options = std::nullopt); // Shuts down the ImageSegmenter when all works are done. absl::Status Close() { return runner_->Close(); } diff --git a/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_graph.cc b/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_graph.cc index 1678dd08..d3e522d9 100644 --- a/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_graph.cc +++ b/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_graph.cc @@ -23,6 +23,7 @@ limitations under the License. #include "mediapipe/framework/api2/builder.h" #include "mediapipe/framework/api2/port.h" #include "mediapipe/framework/formats/image.h" +#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/port/status_macros.h" #include "mediapipe/tasks/cc/common.h" #include "mediapipe/tasks/cc/components/calculators/tensor/tensors_to_segmentation_calculator.pb.h" @@ -62,6 +63,7 @@ using LabelItems = mediapipe::proto_ns::Map; constexpr char kSegmentationTag[] = "SEGMENTATION"; constexpr char kGroupedSegmentationTag[] = "GROUPED_SEGMENTATION"; constexpr char kImageTag[] = "IMAGE"; +constexpr char kNormRectTag[] = "NORM_RECT"; constexpr char kTensorsTag[] = "TENSORS"; constexpr char kOutputSizeTag[] = "OUTPUT_SIZE"; @@ -159,6 +161,10 @@ absl::StatusOr GetOutputTensor( // Inputs: // IMAGE - Image // Image to perform segmentation on. +// NORM_RECT - NormalizedRect @Optional +// Describes image rotation and region of image to perform detection +// on. +// @Optional: rect covering the whole image is used if not specified. // // Outputs: // SEGMENTATION - mediapipe::Image @Multiple @@ -196,10 +202,12 @@ class ImageSegmenterGraph : public core::ModelTaskGraph { ASSIGN_OR_RETURN(const auto* model_resources, CreateModelResources(sc)); Graph graph; - ASSIGN_OR_RETURN(auto output_streams, - BuildSegmentationTask( - sc->Options(), *model_resources, - graph[Input(kImageTag)], graph)); + ASSIGN_OR_RETURN( + auto output_streams, + BuildSegmentationTask( + sc->Options(), *model_resources, + graph[Input(kImageTag)], + graph[Input::Optional(kNormRectTag)], graph)); auto& merge_images_to_vector = graph.AddNode("MergeImagesToVectorCalculator"); @@ -228,18 +236,21 @@ class ImageSegmenterGraph : public core::ModelTaskGraph { absl::StatusOr BuildSegmentationTask( const ImageSegmenterOptions& task_options, const core::ModelResources& model_resources, Source image_in, - Graph& graph) { + Source norm_rect_in, Graph& graph) { MP_RETURN_IF_ERROR(SanityCheckOptions(task_options)); // Adds preprocessing calculators and connects them to the graph input image // stream. auto& preprocessing = graph.AddNode("mediapipe.tasks.components.ImagePreprocessingSubgraph"); + bool use_gpu = components::DetermineImagePreprocessingGpuBackend( + task_options.base_options().acceleration()); MP_RETURN_IF_ERROR(ConfigureImagePreprocessing( - model_resources, + model_resources, use_gpu, &preprocessing .GetOptions())); image_in >> preprocessing.In(kImageTag); + norm_rect_in >> preprocessing.In(kNormRectTag); // Adds inference subgraph and connects its input stream to the output // tensors produced by the ImageToTensorCalculator. diff --git a/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_test.cc b/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_test.cc index ab23a725..07235563 100644 --- a/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_test.cc +++ b/mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_test.cc @@ -29,8 +29,10 @@ limitations under the License. #include "mediapipe/framework/port/opencv_imgcodecs_inc.h" #include "mediapipe/framework/port/status_matchers.h" #include "mediapipe/tasks/cc/components/calculators/tensor/tensors_to_segmentation_calculator.pb.h" +#include "mediapipe/tasks/cc/components/containers/rect.h" #include "mediapipe/tasks/cc/core/proto/base_options.pb.h" #include "mediapipe/tasks/cc/core/proto/external_file.pb.h" +#include "mediapipe/tasks/cc/vision/core/image_processing_options.h" #include "mediapipe/tasks/cc/vision/image_segmenter/proto/image_segmenter_options.pb.h" #include "mediapipe/tasks/cc/vision/utils/image_utils.h" #include "tensorflow/lite/core/shims/cc/shims_test_util.h" @@ -44,6 +46,8 @@ namespace { using ::mediapipe::Image; using ::mediapipe::file::JoinPath; +using ::mediapipe::tasks::components::containers::Rect; +using ::mediapipe::tasks::vision::core::ImageProcessingOptions; using ::testing::HasSubstr; using ::testing::Optional; @@ -237,7 +241,6 @@ TEST_F(ImageModeTest, SucceedsWithConfidenceMask) { MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr segmenter, ImageSegmenter::Create(std::move(options))); - MP_ASSERT_OK_AND_ASSIGN(auto results, segmenter->Segment(image)); MP_ASSERT_OK_AND_ASSIGN(auto confidence_masks, segmenter->Segment(image)); EXPECT_EQ(confidence_masks.size(), 21); @@ -253,6 +256,61 @@ TEST_F(ImageModeTest, SucceedsWithConfidenceMask) { SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity)); } +TEST_F(ImageModeTest, SucceedsWithRotation) { + MP_ASSERT_OK_AND_ASSIGN( + Image image, DecodeImageFromFile( + JoinPath("./", kTestDataDirectory, "cat_rotated.jpg"))); + auto options = std::make_unique(); + options->base_options.model_asset_path = + JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata); + options->output_type = ImageSegmenterOptions::OutputType::CONFIDENCE_MASK; + options->activation = ImageSegmenterOptions::Activation::SOFTMAX; + + MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr segmenter, + ImageSegmenter::Create(std::move(options))); + ImageProcessingOptions image_processing_options; + image_processing_options.rotation_degrees = -90; + MP_ASSERT_OK_AND_ASSIGN(auto confidence_masks, segmenter->Segment(image)); + EXPECT_EQ(confidence_masks.size(), 21); + + cv::Mat expected_mask = + cv::imread(JoinPath("./", kTestDataDirectory, "cat_rotated_mask.jpg"), + cv::IMREAD_GRAYSCALE); + cv::Mat expected_mask_float; + expected_mask.convertTo(expected_mask_float, CV_32FC1, 1 / 255.f); + + // Cat category index 8. + cv::Mat cat_mask = mediapipe::formats::MatView( + confidence_masks[8].GetImageFrameSharedPtr().get()); + EXPECT_THAT(cat_mask, + SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity)); +} + +TEST_F(ImageModeTest, FailsWithRegionOfInterest) { + MP_ASSERT_OK_AND_ASSIGN( + Image image, + DecodeImageFromFile(JoinPath("./", kTestDataDirectory, "cat.jpg"))); + auto options = std::make_unique(); + options->base_options.model_asset_path = + JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata); + options->output_type = ImageSegmenterOptions::OutputType::CONFIDENCE_MASK; + options->activation = ImageSegmenterOptions::Activation::SOFTMAX; + + MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr segmenter, + ImageSegmenter::Create(std::move(options))); + Rect roi{/*left=*/0.1, /*top=*/0, /*right=*/0.9, /*bottom=*/1}; + ImageProcessingOptions image_processing_options{roi, /*rotation_degrees=*/0}; + + auto results = segmenter->Segment(image, image_processing_options); + EXPECT_EQ(results.status().code(), absl::StatusCode::kInvalidArgument); + EXPECT_THAT(results.status().message(), + HasSubstr("This task doesn't support region-of-interest")); + EXPECT_THAT( + results.status().GetPayload(kMediaPipeTasksPayload), + Optional(absl::Cord(absl::StrCat( + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError)))); +} + TEST_F(ImageModeTest, SucceedsSelfie128x128Segmentation) { Image image = GetSRGBImage(JoinPath("./", kTestDataDirectory, "mozart_square.jpg")); diff --git a/mediapipe/tasks/cc/vision/object_detector/BUILD b/mediapipe/tasks/cc/vision/object_detector/BUILD index 18690950..8220d8b7 100644 --- a/mediapipe/tasks/cc/vision/object_detector/BUILD +++ b/mediapipe/tasks/cc/vision/object_detector/BUILD @@ -75,6 +75,7 @@ cc_library( "//mediapipe/tasks/cc/core/proto:base_options_cc_proto", "//mediapipe/tasks/cc/core/proto:inference_subgraph_cc_proto", "//mediapipe/tasks/cc/vision/core:base_vision_task_api", + "//mediapipe/tasks/cc/vision/core:image_processing_options", "//mediapipe/tasks/cc/vision/core:running_mode", "//mediapipe/tasks/cc/vision/core:vision_task_api_factory", "//mediapipe/tasks/cc/vision/object_detector/proto:object_detector_options_cc_proto", diff --git a/mediapipe/tasks/cc/vision/object_detector/object_detector.cc b/mediapipe/tasks/cc/vision/object_detector/object_detector.cc index 9149a3cb..dd19237f 100644 --- a/mediapipe/tasks/cc/vision/object_detector/object_detector.cc +++ b/mediapipe/tasks/cc/vision/object_detector/object_detector.cc @@ -34,6 +34,7 @@ limitations under the License. #include "mediapipe/tasks/cc/core/proto/base_options.pb.h" #include "mediapipe/tasks/cc/core/proto/inference_subgraph.pb.h" #include "mediapipe/tasks/cc/core/utils.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/core/vision_task_api_factory.h" #include "mediapipe/tasks/cc/vision/object_detector/proto/object_detector_options.pb.h" @@ -58,31 +59,6 @@ constexpr int kMicroSecondsPerMilliSecond = 1000; using ObjectDetectorOptionsProto = object_detector::proto::ObjectDetectorOptions; -// Returns a NormalizedRect filling the whole image. If input is present, its -// rotation is set in the returned NormalizedRect and a check is performed to -// make sure no region-of-interest was provided. Otherwise, rotation is set to -// 0. -absl::StatusOr FillNormalizedRect( - std::optional normalized_rect) { - NormalizedRect result; - if (normalized_rect.has_value()) { - result = *normalized_rect; - } - bool has_coordinates = result.has_x_center() || result.has_y_center() || - result.has_width() || result.has_height(); - if (has_coordinates) { - return CreateStatusWithPayload( - absl::StatusCode::kInvalidArgument, - "ObjectDetector does not support region-of-interest.", - MediaPipeTasksStatus::kInvalidArgumentError); - } - result.set_x_center(0.5); - result.set_y_center(0.5); - result.set_width(1); - result.set_height(1); - return result; -} - // Creates a MediaPipe graph config that contains a subgraph node of // "mediapipe.tasks.vision.ObjectDetectorGraph". If the task is running in the // live stream mode, a "FlowLimiterCalculator" will be added to limit the @@ -170,15 +146,16 @@ absl::StatusOr> ObjectDetector::Create( absl::StatusOr> ObjectDetector::Detect( mediapipe::Image image, - std::optional image_processing_options) { + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, absl::StrCat("GPU input images are currently not supported."), MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - ASSIGN_OR_RETURN(NormalizedRect norm_rect, - FillNormalizedRect(image_processing_options)); + ASSIGN_OR_RETURN( + NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false)); ASSIGN_OR_RETURN( auto output_packets, ProcessImageData( @@ -189,15 +166,16 @@ absl::StatusOr> ObjectDetector::Detect( absl::StatusOr> ObjectDetector::DetectForVideo( mediapipe::Image image, int64 timestamp_ms, - std::optional image_processing_options) { + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, absl::StrCat("GPU input images are currently not supported."), MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - ASSIGN_OR_RETURN(NormalizedRect norm_rect, - FillNormalizedRect(image_processing_options)); + ASSIGN_OR_RETURN( + NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false)); ASSIGN_OR_RETURN( auto output_packets, ProcessVideoData( @@ -212,15 +190,16 @@ absl::StatusOr> ObjectDetector::DetectForVideo( absl::Status ObjectDetector::DetectAsync( Image image, int64 timestamp_ms, - std::optional image_processing_options) { + std::optional image_processing_options) { if (image.UsesGpu()) { return CreateStatusWithPayload( absl::StatusCode::kInvalidArgument, absl::StrCat("GPU input images are currently not supported."), MediaPipeTasksStatus::kRunnerUnexpectedInputError); } - ASSIGN_OR_RETURN(NormalizedRect norm_rect, - FillNormalizedRect(image_processing_options)); + ASSIGN_OR_RETURN( + NormalizedRect norm_rect, + ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false)); return SendLiveStreamData( {{kImageInStreamName, MakePacket(std::move(image)) diff --git a/mediapipe/tasks/cc/vision/object_detector/object_detector.h b/mediapipe/tasks/cc/vision/object_detector/object_detector.h index 2e5ed7b8..44ce68ed 100644 --- a/mediapipe/tasks/cc/vision/object_detector/object_detector.h +++ b/mediapipe/tasks/cc/vision/object_detector/object_detector.h @@ -27,9 +27,9 @@ limitations under the License. #include "absl/status/statusor.h" #include "mediapipe/framework/formats/detection.pb.h" #include "mediapipe/framework/formats/image.h" -#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/tasks/cc/core/base_options.h" #include "mediapipe/tasks/cc/vision/core/base_vision_task_api.h" +#include "mediapipe/tasks/cc/vision/core/image_processing_options.h" #include "mediapipe/tasks/cc/vision/core/running_mode.h" namespace mediapipe { @@ -154,10 +154,9 @@ class ObjectDetector : tasks::vision::core::BaseVisionTaskApi { // after the yuv support is implemented. // // The optional 'image_processing_options' parameter can be used to specify - // the rotation to apply to the image before performing classification, by - // setting its 'rotation' field in radians (e.g. 'M_PI / 2' for a 90° - // anti-clockwise rotation). Note that specifying a region-of-interest using - // the 'x_center', 'y_center', 'width' and 'height' fields is NOT supported + // the rotation to apply to the image before performing detection, by + // setting its 'rotation_degrees' field. Note that specifying a + // region-of-interest using the 'region_of_interest' field is NOT supported // and will result in an invalid argument error being returned. // // For CPU images, the returned bounding boxes are expressed in the @@ -168,7 +167,7 @@ class ObjectDetector : tasks::vision::core::BaseVisionTaskApi { // images after enabling the gpu support in MediaPipe Tasks. absl::StatusOr> Detect( mediapipe::Image image, - std::optional image_processing_options = + std::optional image_processing_options = std::nullopt); // Performs object detection on the provided video frame. @@ -180,10 +179,9 @@ class ObjectDetector : tasks::vision::core::BaseVisionTaskApi { // must be monotonically increasing. // // The optional 'image_processing_options' parameter can be used to specify - // the rotation to apply to the image before performing classification, by - // setting its 'rotation' field in radians (e.g. 'M_PI / 2' for a 90° - // anti-clockwise rotation). Note that specifying a region-of-interest using - // the 'x_center', 'y_center', 'width' and 'height' fields is NOT supported + // the rotation to apply to the image before performing detection, by + // setting its 'rotation_degrees' field. Note that specifying a + // region-of-interest using the 'region_of_interest' field is NOT supported // and will result in an invalid argument error being returned. // // For CPU images, the returned bounding boxes are expressed in the @@ -192,7 +190,7 @@ class ObjectDetector : tasks::vision::core::BaseVisionTaskApi { // underlying image data. absl::StatusOr> DetectForVideo( mediapipe::Image image, int64 timestamp_ms, - std::optional image_processing_options = + std::optional image_processing_options = std::nullopt); // Sends live image data to perform object detection, and the results will be @@ -206,10 +204,9 @@ class ObjectDetector : tasks::vision::core::BaseVisionTaskApi { // increasing. // // The optional 'image_processing_options' parameter can be used to specify - // the rotation to apply to the image before performing classification, by - // setting its 'rotation' field in radians (e.g. 'M_PI / 2' for a 90° - // anti-clockwise rotation). Note that specifying a region-of-interest using - // the 'x_center', 'y_center', 'width' and 'height' fields is NOT supported + // the rotation to apply to the image before performing detection, by + // setting its 'rotation_degrees' field. Note that specifying a + // region-of-interest using the 'region_of_interest' field is NOT supported // and will result in an invalid argument error being returned. // // The "result_callback" provides @@ -223,7 +220,7 @@ class ObjectDetector : tasks::vision::core::BaseVisionTaskApi { // outside of the callback, callers need to make a copy of the image. // - The input timestamp in milliseconds. absl::Status DetectAsync(mediapipe::Image image, int64 timestamp_ms, - std::optional + std::optional image_processing_options = std::nullopt); // Shuts down the ObjectDetector when all works are done. diff --git a/mediapipe/tasks/cc/vision/object_detector/object_detector_graph.cc b/mediapipe/tasks/cc/vision/object_detector/object_detector_graph.cc index 07e912cf..b149cea0 100644 --- a/mediapipe/tasks/cc/vision/object_detector/object_detector_graph.cc +++ b/mediapipe/tasks/cc/vision/object_detector/object_detector_graph.cc @@ -563,8 +563,10 @@ class ObjectDetectorGraph : public core::ModelTaskGraph { // stream. auto& preprocessing = graph.AddNode("mediapipe.tasks.components.ImagePreprocessingSubgraph"); + bool use_gpu = components::DetermineImagePreprocessingGpuBackend( + task_options.base_options().acceleration()); MP_RETURN_IF_ERROR(ConfigureImagePreprocessing( - model_resources, + model_resources, use_gpu, &preprocessing .GetOptions())); image_in >> preprocessing.In(kImageTag); 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 8db3fa76..1747685d 100644 --- a/mediapipe/tasks/cc/vision/object_detector/object_detector_test.cc +++ b/mediapipe/tasks/cc/vision/object_detector/object_detector_test.cc @@ -31,11 +31,12 @@ limitations under the License. #include "mediapipe/framework/deps/file_path.h" #include "mediapipe/framework/formats/image.h" #include "mediapipe/framework/formats/location_data.pb.h" -#include "mediapipe/framework/formats/rect.pb.h" #include "mediapipe/framework/port/gmock.h" #include "mediapipe/framework/port/gtest.h" #include "mediapipe/framework/port/parse_text_proto.h" #include "mediapipe/framework/port/status_matchers.h" +#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_utils.h" #include "tensorflow/lite/c/common.h" @@ -64,6 +65,8 @@ namespace vision { namespace { using ::mediapipe::file::JoinPath; +using ::mediapipe::tasks::components::containers::Rect; +using ::mediapipe::tasks::vision::core::ImageProcessingOptions; using ::testing::HasSubstr; using ::testing::Optional; @@ -532,8 +535,8 @@ TEST_F(ImageModeTest, SucceedsWithRotation) { JoinPath("./", kTestDataDirectory, kMobileSsdWithMetadata); MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr object_detector, ObjectDetector::Create(std::move(options))); - NormalizedRect image_processing_options; - image_processing_options.set_rotation(M_PI / 2.0); + ImageProcessingOptions image_processing_options; + image_processing_options.rotation_degrees = -90; MP_ASSERT_OK_AND_ASSIGN( auto results, object_detector->Detect(image, image_processing_options)); MP_ASSERT_OK(object_detector->Close()); @@ -557,16 +560,17 @@ TEST_F(ImageModeTest, FailsWithRegionOfInterest) { JoinPath("./", kTestDataDirectory, kMobileSsdWithMetadata); MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr object_detector, ObjectDetector::Create(std::move(options))); - NormalizedRect image_processing_options; - image_processing_options.set_x_center(0.5); - image_processing_options.set_y_center(0.5); - image_processing_options.set_width(1.0); - image_processing_options.set_height(1.0); + Rect roi{/*left=*/0.1, /*top=*/0, /*right=*/0.9, /*bottom=*/1}; + ImageProcessingOptions image_processing_options{roi, /*rotation_degrees=*/0}; auto results = object_detector->Detect(image, image_processing_options); EXPECT_EQ(results.status().code(), absl::StatusCode::kInvalidArgument); EXPECT_THAT(results.status().message(), - HasSubstr("ObjectDetector does not support region-of-interest")); + HasSubstr("This task doesn't support region-of-interest")); + EXPECT_THAT( + results.status().GetPayload(kMediaPipeTasksPayload), + Optional(absl::Cord(absl::StrCat( + MediaPipeTasksStatus::kImageProcessingInvalidArgumentError)))); } class VideoModeTest : public tflite_shims::testing::Test {}; diff --git a/mediapipe/tasks/examples/android/objectdetector/src/main/BUILD b/mediapipe/tasks/examples/android/objectdetector/src/main/BUILD index acbdbd6e..89c1edcb 100644 --- a/mediapipe/tasks/examples/android/objectdetector/src/main/BUILD +++ b/mediapipe/tasks/examples/android/objectdetector/src/main/BUILD @@ -31,6 +31,7 @@ android_binary( multidex = "native", resource_files = ["//mediapipe/tasks/examples/android:resource_files"], deps = [ + "//mediapipe/java/com/google/mediapipe/framework:android_framework", "//mediapipe/java/com/google/mediapipe/framework/image", "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:detection", "//mediapipe/tasks/java/com/google/mediapipe/tasks/core", diff --git a/mediapipe/tasks/examples/android/objectdetector/src/main/java/com/google/mediapipe/tasks/examples/objectdetector/MainActivity.java b/mediapipe/tasks/examples/android/objectdetector/src/main/java/com/google/mediapipe/tasks/examples/objectdetector/MainActivity.java index 7f7ec138..18c010a0 100644 --- a/mediapipe/tasks/examples/android/objectdetector/src/main/java/com/google/mediapipe/tasks/examples/objectdetector/MainActivity.java +++ b/mediapipe/tasks/examples/android/objectdetector/src/main/java/com/google/mediapipe/tasks/examples/objectdetector/MainActivity.java @@ -16,7 +16,6 @@ package com.google.mediapipe.tasks.examples.objectdetector; import android.content.Intent; import android.graphics.Bitmap; -import android.graphics.Matrix; import android.media.MediaMetadataRetriever; import android.os.Bundle; import android.provider.MediaStore; @@ -29,9 +28,11 @@ import androidx.activity.result.ActivityResultLauncher; import androidx.activity.result.contract.ActivityResultContracts; import androidx.exifinterface.media.ExifInterface; // ContentResolver dependency +import com.google.mediapipe.framework.MediaPipeException; import com.google.mediapipe.framework.image.BitmapImageBuilder; -import com.google.mediapipe.framework.image.Image; +import com.google.mediapipe.framework.image.MPImage; import com.google.mediapipe.tasks.core.BaseOptions; +import com.google.mediapipe.tasks.vision.core.ImageProcessingOptions; import com.google.mediapipe.tasks.vision.core.RunningMode; import com.google.mediapipe.tasks.vision.objectdetector.ObjectDetectionResult; import com.google.mediapipe.tasks.vision.objectdetector.ObjectDetector; @@ -82,6 +83,7 @@ public class MainActivity extends AppCompatActivity { if (resultIntent != null) { if (result.getResultCode() == RESULT_OK) { Bitmap bitmap = null; + int rotation = 0; try { bitmap = downscaleBitmap( @@ -93,13 +95,16 @@ public class MainActivity extends AppCompatActivity { try { InputStream imageData = this.getContentResolver().openInputStream(resultIntent.getData()); - bitmap = rotateBitmap(bitmap, imageData); - } catch (IOException e) { + rotation = getImageRotation(imageData); + } catch (IOException | MediaPipeException e) { Log.e(TAG, "Bitmap rotation error:" + e); } if (bitmap != null) { - Image image = new BitmapImageBuilder(bitmap).build(); - ObjectDetectionResult detectionResult = objectDetector.detect(image); + MPImage image = new BitmapImageBuilder(bitmap).build(); + ObjectDetectionResult detectionResult = + objectDetector.detect( + image, + ImageProcessingOptions.builder().setRotationDegrees(rotation).build()); imageView.setData(image, detectionResult); runOnUiThread(() -> imageView.update()); } @@ -144,7 +149,8 @@ public class MainActivity extends AppCompatActivity { MediaMetadataRetriever.METADATA_KEY_VIDEO_FRAME_COUNT)); long frameIntervalMs = duration / numFrames; for (int i = 0; i < numFrames; ++i) { - Image image = new BitmapImageBuilder(metaRetriever.getFrameAtIndex(i)).build(); + MPImage image = + new BitmapImageBuilder(metaRetriever.getFrameAtIndex(i)).build(); ObjectDetectionResult detectionResult = objectDetector.detectForVideo(image, frameIntervalMs * i); // Currently only annotates the detection result on the first video frame and @@ -209,28 +215,25 @@ public class MainActivity extends AppCompatActivity { return Bitmap.createScaledBitmap(originalBitmap, width, height, false); } - private Bitmap rotateBitmap(Bitmap inputBitmap, InputStream imageData) throws IOException { + private int getImageRotation(InputStream imageData) throws IOException, MediaPipeException { int orientation = new ExifInterface(imageData) .getAttributeInt(ExifInterface.TAG_ORIENTATION, ExifInterface.ORIENTATION_NORMAL); - if (orientation == ExifInterface.ORIENTATION_NORMAL) { - return inputBitmap; - } - Matrix matrix = new Matrix(); switch (orientation) { + case ExifInterface.ORIENTATION_NORMAL: + return 0; case ExifInterface.ORIENTATION_ROTATE_90: - matrix.postRotate(90); - break; + return 90; case ExifInterface.ORIENTATION_ROTATE_180: - matrix.postRotate(180); - break; + return 180; case ExifInterface.ORIENTATION_ROTATE_270: - matrix.postRotate(270); - break; + return 270; default: - matrix.postRotate(0); + // TODO: use getRotationDegrees() and isFlipped() instead of switch once flip + // is supported. + throw new MediaPipeException( + MediaPipeException.StatusCode.UNIMPLEMENTED.ordinal(), + "Flipped images are not supported yet."); } - return Bitmap.createBitmap( - inputBitmap, 0, 0, inputBitmap.getWidth(), inputBitmap.getHeight(), matrix, true); } } diff --git a/mediapipe/tasks/examples/android/objectdetector/src/main/java/com/google/mediapipe/tasks/examples/objectdetector/ObjectDetectionResultImageView.java b/mediapipe/tasks/examples/android/objectdetector/src/main/java/com/google/mediapipe/tasks/examples/objectdetector/ObjectDetectionResultImageView.java index 94a4a90d..283e4885 100644 --- a/mediapipe/tasks/examples/android/objectdetector/src/main/java/com/google/mediapipe/tasks/examples/objectdetector/ObjectDetectionResultImageView.java +++ b/mediapipe/tasks/examples/android/objectdetector/src/main/java/com/google/mediapipe/tasks/examples/objectdetector/ObjectDetectionResultImageView.java @@ -22,7 +22,7 @@ import android.graphics.Matrix; import android.graphics.Paint; import androidx.appcompat.widget.AppCompatImageView; import com.google.mediapipe.framework.image.BitmapExtractor; -import com.google.mediapipe.framework.image.Image; +import com.google.mediapipe.framework.image.MPImage; import com.google.mediapipe.tasks.components.containers.Detection; import com.google.mediapipe.tasks.vision.objectdetector.ObjectDetectionResult; @@ -40,12 +40,12 @@ public class ObjectDetectionResultImageView extends AppCompatImageView { } /** - * Sets an {@link Image} and an {@link ObjectDetectionResult} to render. + * Sets a {@link MPImage} and an {@link ObjectDetectionResult} to render. * - * @param image an {@link Image} object for annotation. + * @param image a {@link MPImage} object for annotation. * @param result an {@link ObjectDetectionResult} object that contains the detection result. */ - public void setData(Image image, ObjectDetectionResult result) { + public void setData(MPImage image, ObjectDetectionResult result) { if (image == null || result == null) { return; } diff --git a/mediapipe/tasks/java/BUILD b/mediapipe/tasks/java/BUILD index 02451073..7e628326 100644 --- a/mediapipe/tasks/java/BUILD +++ b/mediapipe/tasks/java/BUILD @@ -1 +1,15 @@ -# dummy file for tap test to find the pattern +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +licenses(["notice"]) diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/BUILD b/mediapipe/tasks/java/com/google/mediapipe/tasks/BUILD new file mode 100644 index 00000000..7e628326 --- /dev/null +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/BUILD @@ -0,0 +1,15 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +licenses(["notice"]) diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers/Landmark.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers/Landmark.java index 3f96d777..e4586619 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers/Landmark.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers/Landmark.java @@ -30,13 +30,13 @@ public abstract class Landmark { return new AutoValue_Landmark(x, y, z, normalized); } - // The x coordniates of the landmark. + // The x coordinates of the landmark. public abstract float x(); - // The y coordniates of the landmark. + // The y coordinates of the landmark. public abstract float y(); - // The z coordniates of the landmark. + // The z coordinates of the landmark. public abstract float z(); // Whether this landmark is normalized with respect to the image size. diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/BUILD b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/BUILD index b4ebfe8c..cb9d6742 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/BUILD +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/BUILD @@ -36,3 +36,15 @@ android_library( "@maven//:com_google_guava_guava", ], ) + +load("//mediapipe/tasks/java/com/google/mediapipe/tasks:mediapipe_tasks_aar.bzl", "mediapipe_tasks_core_aar") + +mediapipe_tasks_core_aar( + name = "tasks_core", + srcs = glob(["*.java"]) + [ + "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:java_src", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/processors:java_src", + "//mediapipe/java/com/google/mediapipe/framework/image:java_src", + ], + manifest = "AndroidManifest.xml", +) diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/mediapipe_tasks_aar.bzl b/mediapipe/tasks/java/com/google/mediapipe/tasks/mediapipe_tasks_aar.bzl new file mode 100644 index 00000000..0260e3fa --- /dev/null +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/mediapipe_tasks_aar.bzl @@ -0,0 +1,256 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Building MediaPipe Tasks AARs.""" + +load("//mediapipe/java/com/google/mediapipe:mediapipe_aar.bzl", "mediapipe_build_aar_with_jni", "mediapipe_java_proto_src_extractor", "mediapipe_java_proto_srcs") +load("@build_bazel_rules_android//android:rules.bzl", "android_library") + +_CORE_TASKS_JAVA_PROTO_LITE_TARGETS = [ + "//mediapipe/tasks/cc/components/containers/proto:category_java_proto_lite", + "//mediapipe/tasks/cc/components/containers/proto:classifications_java_proto_lite", + "//mediapipe/tasks/cc/components/containers/proto:embeddings_java_proto_lite", + "//mediapipe/tasks/cc/components/containers/proto:landmarks_detection_result_java_proto_lite", + "//mediapipe/tasks/cc/components/processors/proto:classifier_options_java_proto_lite", + "//mediapipe/tasks/cc/core/proto:acceleration_java_proto_lite", + "//mediapipe/tasks/cc/core/proto:base_options_java_proto_lite", + "//mediapipe/tasks/cc/core/proto:external_file_java_proto_lite", +] + +_VISION_TASKS_JAVA_PROTO_LITE_TARGETS = [ + "//mediapipe/tasks/cc/vision/object_detector/proto:object_detector_options_java_proto_lite", + "//mediapipe/tasks/cc/vision/image_classifier/proto:image_classifier_graph_options_java_proto_lite", + "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_classifier_graph_options_java_proto_lite", + "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_embedder_graph_options_java_proto_lite", + "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_recognizer_graph_options_java_proto_lite", + "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:hand_gesture_recognizer_graph_options_java_proto_lite", + "//mediapipe/tasks/cc/vision/hand_detector/proto:hand_detector_graph_options_java_proto_lite", + "//mediapipe/tasks/cc/vision/hand_landmarker/proto:hand_landmarker_graph_options_java_proto_lite", + "//mediapipe/tasks/cc/vision/hand_landmarker/proto:hand_landmarks_detector_graph_options_java_proto_lite", +] + +_TEXT_TASKS_JAVA_PROTO_LITE_TARGETS = [ + "//mediapipe/tasks/cc/text/text_classifier/proto:text_classifier_graph_options_java_proto_lite", +] + +def mediapipe_tasks_core_aar(name, srcs, manifest): + """Builds medaipipe tasks core AAR. + + Args: + name: The bazel target name. + srcs: MediaPipe Tasks' core layer source files. + manifest: The Android manifest. + """ + + mediapipe_tasks_java_proto_srcs = [] + for target in _CORE_TASKS_JAVA_PROTO_LITE_TARGETS: + mediapipe_tasks_java_proto_srcs.append( + _mediapipe_tasks_java_proto_src_extractor(target = target), + ) + + for target in _VISION_TASKS_JAVA_PROTO_LITE_TARGETS: + mediapipe_tasks_java_proto_srcs.append( + _mediapipe_tasks_java_proto_src_extractor(target = target), + ) + + for target in _TEXT_TASKS_JAVA_PROTO_LITE_TARGETS: + mediapipe_tasks_java_proto_srcs.append( + _mediapipe_tasks_java_proto_src_extractor(target = target), + ) + + mediapipe_tasks_java_proto_srcs.append(mediapipe_java_proto_src_extractor( + target = "//mediapipe/calculators/core:flow_limiter_calculator_java_proto_lite", + src_out = "com/google/mediapipe/calculator/proto/FlowLimiterCalculatorProto.java", + )) + + mediapipe_tasks_java_proto_srcs.append(mediapipe_java_proto_src_extractor( + target = "//mediapipe/calculators/tensor:inference_calculator_java_proto_lite", + src_out = "com/google/mediapipe/calculator/proto/InferenceCalculatorProto.java", + )) + + android_library( + name = name, + srcs = srcs + [ + "//mediapipe/java/com/google/mediapipe/framework:java_src", + ] + mediapipe_java_proto_srcs() + + mediapipe_tasks_java_proto_srcs, + javacopts = [ + "-Xep:AndroidJdkLibsChecker:OFF", + ], + manifest = manifest, + deps = [ + "//mediapipe/calculators/core:flow_limiter_calculator_java_proto_lite", + "//mediapipe/calculators/tensor:inference_calculator_java_proto_lite", + "//mediapipe/framework:calculator_java_proto_lite", + "//mediapipe/framework:calculator_profile_java_proto_lite", + "//mediapipe/framework:calculator_options_java_proto_lite", + "//mediapipe/framework:mediapipe_options_java_proto_lite", + "//mediapipe/framework:packet_factory_java_proto_lite", + "//mediapipe/framework:packet_generator_java_proto_lite", + "//mediapipe/framework:status_handler_java_proto_lite", + "//mediapipe/framework:stream_handler_java_proto_lite", + "//mediapipe/framework/formats:classification_java_proto_lite", + "//mediapipe/framework/formats:detection_java_proto_lite", + "//mediapipe/framework/formats:landmark_java_proto_lite", + "//mediapipe/framework/formats:location_data_java_proto_lite", + "//mediapipe/framework/formats:rect_java_proto_lite", + "//mediapipe/java/com/google/mediapipe/framework:android_framework", + "//mediapipe/java/com/google/mediapipe/framework/image", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/core/jni:model_resources_cache_jni", + "//third_party:androidx_annotation", + "//third_party:autovalue", + "@com_google_protobuf//:protobuf_javalite", + "@maven//:com_google_guava_guava", + "@maven//:com_google_flogger_flogger", + "@maven//:com_google_flogger_flogger_system_backend", + "@maven//:com_google_code_findbugs_jsr305", + ] + + _CORE_TASKS_JAVA_PROTO_LITE_TARGETS + + _VISION_TASKS_JAVA_PROTO_LITE_TARGETS + + _TEXT_TASKS_JAVA_PROTO_LITE_TARGETS, + ) + +def mediapipe_tasks_vision_aar(name, srcs, native_library): + """Builds medaipipe tasks vision AAR. + + Args: + name: The bazel target name. + srcs: MediaPipe Vision Tasks' source files. + native_library: The native library that contains vision tasks' graph and calculators. + """ + + native.genrule( + name = name + "tasks_manifest_generator", + outs = ["AndroidManifest.xml"], + cmd = """ +cat > $(OUTS) < + + + +EOF +""", + ) + + _mediapipe_tasks_aar( + name = name, + srcs = srcs, + manifest = "AndroidManifest.xml", + java_proto_lite_targets = _CORE_TASKS_JAVA_PROTO_LITE_TARGETS + _VISION_TASKS_JAVA_PROTO_LITE_TARGETS, + native_library = native_library, + ) + +def mediapipe_tasks_text_aar(name, srcs, native_library): + """Builds medaipipe tasks text AAR. + + Args: + name: The bazel target name. + srcs: MediaPipe Text Tasks' source files. + native_library: The native library that contains text tasks' graph and calculators. + """ + + native.genrule( + name = name + "tasks_manifest_generator", + outs = ["AndroidManifest.xml"], + cmd = """ +cat > $(OUTS) < + + + +EOF +""", + ) + + _mediapipe_tasks_aar( + name = name, + srcs = srcs, + manifest = "AndroidManifest.xml", + java_proto_lite_targets = _CORE_TASKS_JAVA_PROTO_LITE_TARGETS + _TEXT_TASKS_JAVA_PROTO_LITE_TARGETS, + native_library = native_library, + ) + +def _mediapipe_tasks_aar(name, srcs, manifest, java_proto_lite_targets, native_library): + """Builds medaipipe tasks AAR.""" + + # When "--define EXCLUDE_OPENCV_SO_LIB=1" is set in the build command, + # the OpenCV so libraries will be excluded from the AAR package to + # save the package size. + native.config_setting( + name = "exclude_opencv_so_lib", + define_values = { + "EXCLUDE_OPENCV_SO_LIB": "1", + }, + visibility = ["//visibility:public"], + ) + + native.cc_library( + name = name + "_jni_opencv_cc_lib", + srcs = select({ + "//mediapipe:android_arm64": ["@android_opencv//:libopencv_java3_so_arm64-v8a"], + "//mediapipe:android_armeabi": ["@android_opencv//:libopencv_java3_so_armeabi-v7a"], + "//mediapipe:android_arm": ["@android_opencv//:libopencv_java3_so_armeabi-v7a"], + "//mediapipe:android_x86": ["@android_opencv//:libopencv_java3_so_x86"], + "//mediapipe:android_x86_64": ["@android_opencv//:libopencv_java3_so_x86_64"], + "//conditions:default": [], + }), + alwayslink = 1, + ) + + android_library( + name = name + "_android_lib", + srcs = srcs, + manifest = manifest, + proguard_specs = ["//mediapipe/java/com/google/mediapipe/framework:proguard.pgcfg"], + deps = java_proto_lite_targets + [native_library] + [ + "//mediapipe/java/com/google/mediapipe/framework:android_framework", + "//mediapipe/java/com/google/mediapipe/framework/image", + "//mediapipe/framework:calculator_options_java_proto_lite", + "//mediapipe/framework:calculator_java_proto_lite", + "//mediapipe/framework/formats:classification_java_proto_lite", + "//mediapipe/framework/formats:detection_java_proto_lite", + "//mediapipe/framework/formats:landmark_java_proto_lite", + "//mediapipe/framework/formats:location_data_java_proto_lite", + "//mediapipe/framework/formats:rect_java_proto_lite", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:detection", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:category", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:classification_entry", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:classifications", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:landmark", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/processors:classifieroptions", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/core", + "//third_party:autovalue", + "@maven//:com_google_guava_guava", + ] + select({ + "//conditions:default": [":" + name + "_jni_opencv_cc_lib"], + "//mediapipe/framework/port:disable_opencv": [], + "exclude_opencv_so_lib": [], + }), + ) + + mediapipe_build_aar_with_jni(name, name + "_android_lib") + +def _mediapipe_tasks_java_proto_src_extractor(target): + proto_path = "com/google/" + target.split(":")[0].replace("cc/", "").replace("//", "").replace("_", "") + "/" + proto_name = target.split(":")[-1].replace("_java_proto_lite", "").replace("_", " ").title().replace(" ", "") + "Proto.java" + return mediapipe_java_proto_src_extractor( + target = target, + src_out = proto_path + proto_name, + ) diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/text/BUILD b/mediapipe/tasks/java/com/google/mediapipe/tasks/text/BUILD index 1719707d..fa2a547c 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/text/BUILD +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/text/BUILD @@ -61,3 +61,11 @@ android_library( "@maven//:com_google_guava_guava", ], ) + +load("//mediapipe/tasks/java/com/google/mediapipe/tasks:mediapipe_tasks_aar.bzl", "mediapipe_tasks_text_aar") + +mediapipe_tasks_text_aar( + name = "tasks_text", + srcs = glob(["**/*.java"]), + native_library = ":libmediapipe_tasks_text_jni_lib", +) diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/text/textclassifier/TextClassificationResult.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/text/textclassifier/TextClassificationResult.java index dd9b9a1b..c1e2446c 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/text/textclassifier/TextClassificationResult.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/text/textclassifier/TextClassificationResult.java @@ -15,11 +15,11 @@ package com.google.mediapipe.tasks.text.textclassifier; import com.google.auto.value.AutoValue; -import com.google.mediapipe.tasks.components.container.proto.CategoryProto; -import com.google.mediapipe.tasks.components.container.proto.ClassificationsProto; import com.google.mediapipe.tasks.components.containers.Category; import com.google.mediapipe.tasks.components.containers.ClassificationEntry; import com.google.mediapipe.tasks.components.containers.Classifications; +import com.google.mediapipe.tasks.components.containers.proto.CategoryProto; +import com.google.mediapipe.tasks.components.containers.proto.ClassificationsProto; import com.google.mediapipe.tasks.core.TaskResult; import java.util.ArrayList; import java.util.Collections; diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/text/textclassifier/TextClassifier.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/text/textclassifier/TextClassifier.java index 0ebffe7c..07a4fa48 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/text/textclassifier/TextClassifier.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/text/textclassifier/TextClassifier.java @@ -22,7 +22,7 @@ import com.google.mediapipe.framework.MediaPipeException; import com.google.mediapipe.framework.Packet; import com.google.mediapipe.framework.PacketGetter; import com.google.mediapipe.framework.ProtoUtil; -import com.google.mediapipe.tasks.components.container.proto.ClassificationsProto; +import com.google.mediapipe.tasks.components.containers.proto.ClassificationsProto; import com.google.mediapipe.tasks.components.processors.ClassifierOptions; import com.google.mediapipe.tasks.core.BaseOptions; import com.google.mediapipe.tasks.core.OutputHandler; @@ -31,7 +31,6 @@ import com.google.mediapipe.tasks.core.TaskOptions; import com.google.mediapipe.tasks.core.TaskRunner; import com.google.mediapipe.tasks.core.proto.BaseOptionsProto; import com.google.mediapipe.tasks.text.textclassifier.proto.TextClassifierGraphOptionsProto; -import com.google.protobuf.InvalidProtocolBufferException; import java.io.File; import java.io.IOException; import java.util.Arrays; @@ -154,7 +153,7 @@ public final class TextClassifier implements AutoCloseable { packets.get(CLASSIFICATION_RESULT_OUT_STREAM_INDEX), ClassificationsProto.ClassificationResult.getDefaultInstance()), packets.get(CLASSIFICATION_RESULT_OUT_STREAM_INDEX).getTimestamp()); - } catch (InvalidProtocolBufferException e) { + } catch (IOException e) { throw new MediaPipeException( MediaPipeException.StatusCode.INTERNAL.ordinal(), e.getMessage()); } diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/BUILD b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/BUILD index dcf3b354..d15040ae 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/BUILD +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/BUILD @@ -28,6 +28,7 @@ android_library( "//mediapipe/java/com/google/mediapipe/framework:android_framework_no_mff", "//mediapipe/java/com/google/mediapipe/framework/image", "//mediapipe/tasks/java/com/google/mediapipe/tasks/core", + "//third_party:autovalue", "@maven//:com_google_guava_guava", ], ) @@ -128,6 +129,7 @@ android_library( "//mediapipe/java/com/google/mediapipe/framework/image", "//mediapipe/tasks/cc/components/processors/proto:classifier_options_java_proto_lite", "//mediapipe/tasks/cc/core/proto:base_options_java_proto_lite", + "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_classifier_graph_options_java_proto_lite", "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_recognizer_graph_options_java_proto_lite", "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:hand_gesture_recognizer_graph_options_java_proto_lite", "//mediapipe/tasks/cc/vision/hand_detector/proto:hand_detector_graph_options_java_proto_lite", @@ -135,8 +137,17 @@ android_library( "//mediapipe/tasks/cc/vision/hand_landmarker/proto:hand_landmarks_detector_graph_options_java_proto_lite", "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:category", "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:landmark", + "//mediapipe/tasks/java/com/google/mediapipe/tasks/components/processors:classifieroptions", "//mediapipe/tasks/java/com/google/mediapipe/tasks/core", "//third_party:autovalue", "@maven//:com_google_guava_guava", ], ) + +load("//mediapipe/tasks/java/com/google/mediapipe/tasks:mediapipe_tasks_aar.bzl", "mediapipe_tasks_vision_aar") + +mediapipe_tasks_vision_aar( + name = "tasks_vision", + srcs = glob(["**/*.java"]), + native_library = ":libmediapipe_tasks_vision_jni_lib", +) diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/core/BaseVisionTaskApi.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/core/BaseVisionTaskApi.java index 7ab8e75a..0774b69a 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/core/BaseVisionTaskApi.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/core/BaseVisionTaskApi.java @@ -19,12 +19,11 @@ import com.google.mediapipe.formats.proto.RectProto.NormalizedRect; import com.google.mediapipe.framework.MediaPipeException; import com.google.mediapipe.framework.Packet; import com.google.mediapipe.framework.ProtoUtil; -import com.google.mediapipe.framework.image.Image; +import com.google.mediapipe.framework.image.MPImage; import com.google.mediapipe.tasks.core.TaskResult; import com.google.mediapipe.tasks.core.TaskRunner; import java.util.HashMap; import java.util.Map; -import java.util.Optional; /** The base class of MediaPipe vision tasks. */ public class BaseVisionTaskApi implements AutoCloseable { @@ -32,7 +31,7 @@ public class BaseVisionTaskApi implements AutoCloseable { private final TaskRunner runner; private final RunningMode runningMode; private final String imageStreamName; - private final Optional normRectStreamName; + private final String normRectStreamName; static { System.loadLibrary("mediapipe_tasks_vision_jni"); @@ -40,27 +39,13 @@ public class BaseVisionTaskApi implements AutoCloseable { } /** - * Constructor to initialize a {@link BaseVisionTaskApi} only taking images as input. + * Constructor to initialize a {@link BaseVisionTaskApi}. * * @param runner a {@link TaskRunner}. * @param runningMode a mediapipe vision task {@link RunningMode}. * @param imageStreamName the name of the input image stream. - */ - public BaseVisionTaskApi(TaskRunner runner, RunningMode runningMode, String imageStreamName) { - this.runner = runner; - this.runningMode = runningMode; - this.imageStreamName = imageStreamName; - this.normRectStreamName = Optional.empty(); - } - - /** - * Constructor to initialize a {@link BaseVisionTaskApi} taking images and normalized rects as - * input. - * - * @param runner a {@link TaskRunner}. - * @param runningMode a mediapipe vision task {@link RunningMode}. - * @param imageStreamName the name of the input image stream. - * @param normRectStreamName the name of the input normalized rect image stream. + * @param normRectStreamName the name of the input normalized rect image stream used to provide + * (mandatory) rotation and (optional) region-of-interest. */ public BaseVisionTaskApi( TaskRunner runner, @@ -70,61 +55,31 @@ public class BaseVisionTaskApi implements AutoCloseable { this.runner = runner; this.runningMode = runningMode; this.imageStreamName = imageStreamName; - this.normRectStreamName = Optional.of(normRectStreamName); + this.normRectStreamName = normRectStreamName; } /** * A synchronous method to process single image inputs. The call blocks the current thread until a * failure status or a successful result is returned. * - * @param image a MediaPipe {@link Image} object for processing. - * @throws MediaPipeException if the task is not in the image mode or requires a normalized rect - * input. + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. + * @throws MediaPipeException if the task is not in the image mode. */ - protected TaskResult processImageData(Image image) { + protected TaskResult processImageData( + MPImage image, ImageProcessingOptions imageProcessingOptions) { if (runningMode != RunningMode.IMAGE) { throw new MediaPipeException( MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), "Task is not initialized with the image mode. Current running mode:" + runningMode.name()); } - if (normRectStreamName.isPresent()) { - throw new MediaPipeException( - MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), - "Task expects a normalized rect as input."); - } - Map inputPackets = new HashMap<>(); - inputPackets.put(imageStreamName, runner.getPacketCreator().createImage(image)); - return runner.process(inputPackets); - } - - /** - * A synchronous method to process single image inputs. The call blocks the current thread until a - * failure status or a successful result is returned. - * - * @param image a MediaPipe {@link Image} object for processing. - * @param roi a {@link RectF} defining the region-of-interest to process in the image. Coordinates - * are expected to be specified as normalized values in [0,1]. - * @throws MediaPipeException if the task is not in the image mode or doesn't require a normalized - * rect. - */ - protected TaskResult processImageData(Image image, RectF roi) { - if (runningMode != RunningMode.IMAGE) { - throw new MediaPipeException( - MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), - "Task is not initialized with the image mode. Current running mode:" - + runningMode.name()); - } - if (!normRectStreamName.isPresent()) { - throw new MediaPipeException( - MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), - "Task doesn't expect a normalized rect as input."); - } Map inputPackets = new HashMap<>(); inputPackets.put(imageStreamName, runner.getPacketCreator().createImage(image)); inputPackets.put( - normRectStreamName.get(), - runner.getPacketCreator().createProto(convertToNormalizedRect(roi))); + normRectStreamName, + runner.getPacketCreator().createProto(convertToNormalizedRect(imageProcessingOptions))); return runner.process(inputPackets); } @@ -132,56 +87,25 @@ public class BaseVisionTaskApi implements AutoCloseable { * A synchronous method to process continuous video frames. The call blocks the current thread * until a failure status or a successful result is returned. * - * @param image a MediaPipe {@link Image} object for processing. + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. * @param timestampMs the corresponding timestamp of the input image in milliseconds. - * @throws MediaPipeException if the task is not in the video mode or requires a normalized rect - * input. + * @throws MediaPipeException if the task is not in the video mode. */ - protected TaskResult processVideoData(Image image, long timestampMs) { + protected TaskResult processVideoData( + MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) { if (runningMode != RunningMode.VIDEO) { throw new MediaPipeException( MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), "Task is not initialized with the video mode. Current running mode:" + runningMode.name()); } - if (normRectStreamName.isPresent()) { - throw new MediaPipeException( - MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), - "Task expects a normalized rect as input."); - } - Map inputPackets = new HashMap<>(); - inputPackets.put(imageStreamName, runner.getPacketCreator().createImage(image)); - return runner.process(inputPackets, timestampMs * MICROSECONDS_PER_MILLISECOND); - } - - /** - * A synchronous method to process continuous video frames. The call blocks the current thread - * until a failure status or a successful result is returned. - * - * @param image a MediaPipe {@link Image} object for processing. - * @param roi a {@link RectF} defining the region-of-interest to process in the image. Coordinates - * are expected to be specified as normalized values in [0,1]. - * @param timestampMs the corresponding timestamp of the input image in milliseconds. - * @throws MediaPipeException if the task is not in the video mode or doesn't require a normalized - * rect. - */ - protected TaskResult processVideoData(Image image, RectF roi, long timestampMs) { - if (runningMode != RunningMode.VIDEO) { - throw new MediaPipeException( - MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), - "Task is not initialized with the video mode. Current running mode:" - + runningMode.name()); - } - if (!normRectStreamName.isPresent()) { - throw new MediaPipeException( - MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), - "Task doesn't expect a normalized rect as input."); - } Map inputPackets = new HashMap<>(); inputPackets.put(imageStreamName, runner.getPacketCreator().createImage(image)); inputPackets.put( - normRectStreamName.get(), - runner.getPacketCreator().createProto(convertToNormalizedRect(roi))); + normRectStreamName, + runner.getPacketCreator().createProto(convertToNormalizedRect(imageProcessingOptions))); return runner.process(inputPackets, timestampMs * MICROSECONDS_PER_MILLISECOND); } @@ -189,56 +113,25 @@ public class BaseVisionTaskApi implements AutoCloseable { * An asynchronous method to send live stream data to the {@link TaskRunner}. The results will be * available in the user-defined result listener. * - * @param image a MediaPipe {@link Image} object for processing. + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. * @param timestampMs the corresponding timestamp of the input image in milliseconds. - * @throws MediaPipeException if the task is not in the video mode or requires a normalized rect - * input. + * @throws MediaPipeException if the task is not in the stream mode. */ - protected void sendLiveStreamData(Image image, long timestampMs) { + protected void sendLiveStreamData( + MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) { if (runningMode != RunningMode.LIVE_STREAM) { throw new MediaPipeException( MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), "Task is not initialized with the live stream mode. Current running mode:" + runningMode.name()); } - if (normRectStreamName.isPresent()) { - throw new MediaPipeException( - MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), - "Task expects a normalized rect as input."); - } - Map inputPackets = new HashMap<>(); - inputPackets.put(imageStreamName, runner.getPacketCreator().createImage(image)); - runner.send(inputPackets, timestampMs * MICROSECONDS_PER_MILLISECOND); - } - - /** - * An asynchronous method to send live stream data to the {@link TaskRunner}. The results will be - * available in the user-defined result listener. - * - * @param image a MediaPipe {@link Image} object for processing. - * @param roi a {@link RectF} defining the region-of-interest to process in the image. Coordinates - * are expected to be specified as normalized values in [0,1]. - * @param timestampMs the corresponding timestamp of the input image in milliseconds. - * @throws MediaPipeException if the task is not in the video mode or doesn't require a normalized - * rect. - */ - protected void sendLiveStreamData(Image image, RectF roi, long timestampMs) { - if (runningMode != RunningMode.LIVE_STREAM) { - throw new MediaPipeException( - MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), - "Task is not initialized with the live stream mode. Current running mode:" - + runningMode.name()); - } - if (!normRectStreamName.isPresent()) { - throw new MediaPipeException( - MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(), - "Task doesn't expect a normalized rect as input."); - } Map inputPackets = new HashMap<>(); inputPackets.put(imageStreamName, runner.getPacketCreator().createImage(image)); inputPackets.put( - normRectStreamName.get(), - runner.getPacketCreator().createProto(convertToNormalizedRect(roi))); + normRectStreamName, + runner.getPacketCreator().createProto(convertToNormalizedRect(imageProcessingOptions))); runner.send(inputPackets, timestampMs * MICROSECONDS_PER_MILLISECOND); } @@ -248,13 +141,23 @@ public class BaseVisionTaskApi implements AutoCloseable { runner.close(); } - /** Converts a {@link RectF} object into a {@link NormalizedRect} protobuf message. */ - private static NormalizedRect convertToNormalizedRect(RectF rect) { + /** + * Converts an {@link ImageProcessingOptions} instance into a {@link NormalizedRect} protobuf + * message. + */ + private static NormalizedRect convertToNormalizedRect( + ImageProcessingOptions imageProcessingOptions) { + RectF regionOfInterest = + imageProcessingOptions.regionOfInterest().isPresent() + ? imageProcessingOptions.regionOfInterest().get() + : new RectF(0, 0, 1, 1); return NormalizedRect.newBuilder() - .setXCenter(rect.centerX()) - .setYCenter(rect.centerY()) - .setWidth(rect.width()) - .setHeight(rect.height()) + .setXCenter(regionOfInterest.centerX()) + .setYCenter(regionOfInterest.centerY()) + .setWidth(regionOfInterest.width()) + .setHeight(regionOfInterest.height()) + // Convert to radians anti-clockwise. + .setRotation(-(float) Math.PI * imageProcessingOptions.rotationDegrees() / 180.0f) .build(); } } diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/core/ImageProcessingOptions.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/core/ImageProcessingOptions.java new file mode 100644 index 00000000..a34a9787 --- /dev/null +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/core/ImageProcessingOptions.java @@ -0,0 +1,92 @@ +// Copyright 2022 The MediaPipe Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.mediapipe.tasks.vision.core; + +import android.graphics.RectF; +import com.google.auto.value.AutoValue; +import java.util.Optional; + +// TODO: add support for image flipping. +/** Options for image processing. */ +@AutoValue +public abstract class ImageProcessingOptions { + + /** + * Builder for {@link ImageProcessingOptions}. + * + *

    If both region-of-interest and rotation are specified, the crop around the + * region-of-interest is extracted first, then the specified rotation is applied to the crop. + */ + @AutoValue.Builder + public abstract static class Builder { + /** + * Sets the optional region-of-interest to crop from the image. If not specified, the full image + * is used. + * + *

    Coordinates must be in [0,1], {@code left} must be < {@code right} and {@code top} must be + * < {@code bottom}, otherwise an IllegalArgumentException will be thrown when {@link #build()} + * is called. + */ + public abstract Builder setRegionOfInterest(RectF value); + + /** + * Sets the rotation to apply to the image (or cropped region-of-interest), in degrees + * clockwise. Defaults to 0. + * + *

    The rotation must be a multiple (positive or negative) of 90°, otherwise an + * IllegalArgumentException will be thrown when {@link #build()} is called. + */ + public abstract Builder setRotationDegrees(int value); + + abstract ImageProcessingOptions autoBuild(); + + /** + * Validates and builds the {@link ImageProcessingOptions} instance. + * + * @throws IllegalArgumentException if some of the provided values do not meet their + * requirements. + */ + public final ImageProcessingOptions build() { + ImageProcessingOptions options = autoBuild(); + if (options.regionOfInterest().isPresent()) { + RectF roi = options.regionOfInterest().get(); + if (roi.left >= roi.right || roi.top >= roi.bottom) { + throw new IllegalArgumentException( + String.format( + "Expected left < right and top < bottom, found: %s.", roi.toShortString())); + } + if (roi.left < 0 || roi.right > 1 || roi.top < 0 || roi.bottom > 1) { + throw new IllegalArgumentException( + String.format("Expected RectF values in [0,1], found: %s.", roi.toShortString())); + } + } + if (options.rotationDegrees() % 90 != 0) { + throw new IllegalArgumentException( + String.format( + "Expected rotation to be a multiple of 90°, found: %d.", + options.rotationDegrees())); + } + return options; + } + } + + public abstract Optional regionOfInterest(); + + public abstract int rotationDegrees(); + + public static Builder builder() { + return new AutoValue_ImageProcessingOptions.Builder().setRotationDegrees(0); + } +} diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognitionResult.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognitionResult.java index fd764cb1..e9e1ebe8 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognitionResult.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognitionResult.java @@ -31,6 +31,8 @@ import java.util.List; @AutoValue public abstract class GestureRecognitionResult implements TaskResult { + private static final int kGestureDefaultIndex = -1; + /** * Creates a {@link GestureRecognitionResult} instance from the lists of landmarks, handedness, * and gestures protobuf messages. @@ -97,7 +99,9 @@ public abstract class GestureRecognitionResult implements TaskResult { gestures.add( Category.create( classification.getScore(), - classification.getIndex(), + // Gesture index is not used, because the final gesture result comes from multiple + // classifiers. + kGestureDefaultIndex, classification.getLabel(), classification.getDisplayName())); } @@ -123,6 +127,10 @@ public abstract class GestureRecognitionResult implements TaskResult { /** Handedness of detected hands. */ public abstract List> handednesses(); - /** Recognized hand gestures of detected hands */ + /** + * Recognized hand gestures of detected hands. Note that the index of the gesture is always -1, + * because the raw indices from multiple gesture classifiers cannot consolidate to a meaningful + * index. + */ public abstract List> gestures(); } diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognizer.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognizer.java index 128f6eab..d6faf598 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognizer.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognizer.java @@ -25,8 +25,8 @@ import com.google.mediapipe.framework.AndroidPacketGetter; import com.google.mediapipe.framework.Packet; import com.google.mediapipe.framework.PacketGetter; import com.google.mediapipe.framework.image.BitmapImageBuilder; -import com.google.mediapipe.framework.image.Image; -import com.google.mediapipe.tasks.components.processors.proto.ClassifierOptionsProto; +import com.google.mediapipe.framework.image.MPImage; +import com.google.mediapipe.tasks.components.processors.ClassifierOptions; import com.google.mediapipe.tasks.core.BaseOptions; import com.google.mediapipe.tasks.core.ErrorListener; import com.google.mediapipe.tasks.core.OutputHandler; @@ -36,7 +36,9 @@ import com.google.mediapipe.tasks.core.TaskOptions; import com.google.mediapipe.tasks.core.TaskRunner; import com.google.mediapipe.tasks.core.proto.BaseOptionsProto; import com.google.mediapipe.tasks.vision.core.BaseVisionTaskApi; +import com.google.mediapipe.tasks.vision.core.ImageProcessingOptions; import com.google.mediapipe.tasks.vision.core.RunningMode; +import com.google.mediapipe.tasks.vision.gesturerecognizer.proto.GestureClassifierGraphOptionsProto; import com.google.mediapipe.tasks.vision.gesturerecognizer.proto.GestureRecognizerGraphOptionsProto; import com.google.mediapipe.tasks.vision.gesturerecognizer.proto.HandGestureRecognizerGraphOptionsProto; import com.google.mediapipe.tasks.vision.handdetector.proto.HandDetectorGraphOptionsProto; @@ -58,7 +60,7 @@ import java.util.Optional; * Model Maker. See . * *

      - *
    • Input image {@link Image} + *
    • Input image {@link MPImage} *
        *
      • The image that gesture recognition runs on. *
      @@ -71,8 +73,10 @@ import java.util.Optional; public final class GestureRecognizer extends BaseVisionTaskApi { private static final String TAG = GestureRecognizer.class.getSimpleName(); private static final String IMAGE_IN_STREAM_NAME = "image_in"; + private static final String NORM_RECT_IN_STREAM_NAME = "norm_rect_in"; private static final List INPUT_STREAMS = - Collections.unmodifiableList(Arrays.asList("IMAGE:" + IMAGE_IN_STREAM_NAME)); + Collections.unmodifiableList( + Arrays.asList("IMAGE:" + IMAGE_IN_STREAM_NAME, "NORM_RECT:" + NORM_RECT_IN_STREAM_NAME)); private static final List OUTPUT_STREAMS = Collections.unmodifiableList( Arrays.asList( @@ -148,9 +152,9 @@ public final class GestureRecognizer extends BaseVisionTaskApi { public static GestureRecognizer createFromOptions( Context context, GestureRecognizerOptions recognizerOptions) { // TODO: Consolidate OutputHandler and TaskRunner. - OutputHandler handler = new OutputHandler<>(); + OutputHandler handler = new OutputHandler<>(); handler.setOutputPacketConverter( - new OutputHandler.OutputPacketConverter() { + new OutputHandler.OutputPacketConverter() { @Override public GestureRecognitionResult convertToTaskResult(List packets) { // If there is no hands detected in the image, just returns empty lists. @@ -175,7 +179,7 @@ public final class GestureRecognizer extends BaseVisionTaskApi { } @Override - public Image convertToTaskInput(List packets) { + public MPImage convertToTaskInput(List packets) { return new BitmapImageBuilder( AndroidPacketGetter.getBitmapFromRgb(packets.get(IMAGE_OUT_STREAM_INDEX))) .build(); @@ -205,7 +209,26 @@ public final class GestureRecognizer extends BaseVisionTaskApi { * @param runningMode a mediapipe vision task {@link RunningMode}. */ private GestureRecognizer(TaskRunner taskRunner, RunningMode runningMode) { - super(taskRunner, runningMode, IMAGE_IN_STREAM_NAME); + super(taskRunner, runningMode, IMAGE_IN_STREAM_NAME, NORM_RECT_IN_STREAM_NAME); + } + + /** + * Performs gesture recognition on the provided single image with default image processing + * options, i.e. without any rotation applied. Only use this method when the {@link + * GestureRecognizer} is created with {@link RunningMode.IMAGE}. TODO update java doc + * for input image format. + * + *

      {@link GestureRecognizer} supports the following color space types: + * + *

        + *
      • {@link Bitmap.Config.ARGB_8888} + *
      + * + * @param image a MediaPipe {@link MPImage} object for processing. + * @throws MediaPipeException if there is an internal error. + */ + public GestureRecognitionResult recognize(MPImage image) { + return recognize(image, ImageProcessingOptions.builder().build()); } /** @@ -219,11 +242,41 @@ public final class GestureRecognizer extends BaseVisionTaskApi { *
    • {@link Bitmap.Config.ARGB_8888} *
    * - * @param inputImage a MediaPipe {@link Image} object for processing. + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. Note that region-of-interest is not supported + * by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in + * this method throwing an IllegalArgumentException. + * @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a + * region-of-interest. * @throws MediaPipeException if there is an internal error. */ - public GestureRecognitionResult recognize(Image inputImage) { - return (GestureRecognitionResult) processImageData(inputImage); + public GestureRecognitionResult recognize( + MPImage image, ImageProcessingOptions imageProcessingOptions) { + validateImageProcessingOptions(imageProcessingOptions); + return (GestureRecognitionResult) processImageData(image, imageProcessingOptions); + } + + /** + * Performs gesture recognition on the provided video frame with default image processing options, + * i.e. without any rotation applied. Only use this method when the {@link GestureRecognizer} is + * created with {@link RunningMode.VIDEO}. + * + *

    It's required to provide the video frame's timestamp (in milliseconds). The input timestamps + * must be monotonically increasing. + * + *

    {@link GestureRecognizer} supports the following color space types: + * + *

      + *
    • {@link Bitmap.Config.ARGB_8888} + *
    + * + * @param image a MediaPipe {@link MPImage} object for processing. + * @param timestampMs the input timestamp (in milliseconds). + * @throws MediaPipeException if there is an internal error. + */ + public GestureRecognitionResult recognizeForVideo(MPImage image, long timestampMs) { + return recognizeForVideo(image, ImageProcessingOptions.builder().build(), timestampMs); } /** @@ -239,12 +292,43 @@ public final class GestureRecognizer extends BaseVisionTaskApi { *
  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. - * @param inputTimestampMs the input timestamp (in milliseconds). + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. Note that region-of-interest is not supported + * by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in + * this method throwing an IllegalArgumentException. + * @param timestampMs the input timestamp (in milliseconds). + * @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a + * region-of-interest. * @throws MediaPipeException if there is an internal error. */ - public GestureRecognitionResult recognizeForVideo(Image inputImage, long inputTimestampMs) { - return (GestureRecognitionResult) processVideoData(inputImage, inputTimestampMs); + public GestureRecognitionResult recognizeForVideo( + MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) { + validateImageProcessingOptions(imageProcessingOptions); + return (GestureRecognitionResult) processVideoData(image, imageProcessingOptions, timestampMs); + } + + /** + * Sends live image data to perform gesture recognition with default image processing options, + * i.e. without any rotation applied, and the results will be available via the {@link + * ResultListener} provided in the {@link GestureRecognizerOptions}. Only use this method when the + * {@link GestureRecognition} is created with {@link RunningMode.LIVE_STREAM}. + * + *

    It's required to provide a timestamp (in milliseconds) to indicate when the input image is + * sent to the gesture recognizer. The input timestamps must be monotonically increasing. + * + *

    {@link GestureRecognizer} supports the following color space types: + * + *

      + *
    • {@link Bitmap.Config.ARGB_8888} + *
    + * + * @param image a MediaPipe {@link MPImage} object for processing. + * @param timestampMs the input timestamp (in milliseconds). + * @throws MediaPipeException if there is an internal error. + */ + public void recognizeAsync(MPImage image, long timestampMs) { + recognizeAsync(image, ImageProcessingOptions.builder().build(), timestampMs); } /** @@ -261,12 +345,20 @@ public final class GestureRecognizer extends BaseVisionTaskApi { *
  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. - * @param inputTimestampMs the input timestamp (in milliseconds). + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. Note that region-of-interest is not supported + * by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in + * this method throwing an IllegalArgumentException. + * @param timestampMs the input timestamp (in milliseconds). + * @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a + * region-of-interest. * @throws MediaPipeException if there is an internal error. */ - public void recognizeAsync(Image inputImage, long inputTimestampMs) { - sendLiveStreamData(inputImage, inputTimestampMs); + public void recognizeAsync( + MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) { + validateImageProcessingOptions(imageProcessingOptions); + sendLiveStreamData(image, imageProcessingOptions, timestampMs); } /** Options for setting up an {@link GestureRecognizer}. */ @@ -293,40 +385,46 @@ public final class GestureRecognizer extends BaseVisionTaskApi { */ public abstract Builder setRunningMode(RunningMode value); - // TODO: remove these. Temporary solutions before bundle asset is ready. - public abstract Builder setBaseOptionsHandDetector(BaseOptions value); - - public abstract Builder setBaseOptionsHandLandmarker(BaseOptions value); - - public abstract Builder setBaseOptionsGestureRecognizer(BaseOptions value); - /** Sets the maximum number of hands can be detected by the GestureRecognizer. */ public abstract Builder setNumHands(Integer value); - /** Sets minimum confidence score for the hand detection to be considered successfully */ + /** Sets minimum confidence score for the hand detection to be considered successful */ public abstract Builder setMinHandDetectionConfidence(Float value); /** Sets minimum confidence score of hand presence score in the hand landmark detection. */ public abstract Builder setMinHandPresenceConfidence(Float value); - /** Sets the minimum confidence score for the hand tracking to be considered successfully. */ + /** Sets the minimum confidence score for the hand tracking to be considered successful. */ public abstract Builder setMinTrackingConfidence(Float value); /** - * Sets the minimum confidence score for the gestures to be considered successfully. If < 0, - * the gesture confidence threshold=0.5 for the model is used. + * Sets the optional {@link ClassifierOptions} controling the canned gestures classifier, such + * as score threshold, allow list and deny list of gestures. The categories for canned gesture + * classifiers are: ["None", "Closed_Fist", "Open_Palm", "Pointing_Up", "Thumb_Down", + * "Thumb_Up", "Victory", "ILoveYou"] * *

    TODO Note this option is subject to change, after scoring merging * calculator is implemented. */ - public abstract Builder setMinGestureConfidence(Float value); + public abstract Builder setCannedGesturesClassifierOptions( + ClassifierOptions classifierOptions); + + /** + * Sets the optional {@link ClassifierOptions} controling the custom gestures classifier, such + * as score threshold, allow list and deny list of gestures. + * + *

    TODO Note this option is subject to change, after scoring merging + * calculator is implemented. + */ + public abstract Builder setCustomGesturesClassifierOptions( + ClassifierOptions classifierOptions); /** * Sets the result listener to receive the detection results asynchronously when the gesture * recognizer is in the live stream mode. */ public abstract Builder setResultListener( - ResultListener value); + ResultListener value); /** Sets an optional error listener. */ public abstract Builder setErrorListener(ErrorListener value); @@ -359,13 +457,6 @@ public final class GestureRecognizer extends BaseVisionTaskApi { abstract BaseOptions baseOptions(); - // TODO: remove these. Temporary solutions before bundle asset is ready. - abstract BaseOptions baseOptionsHandDetector(); - - abstract BaseOptions baseOptionsHandLandmarker(); - - abstract BaseOptions baseOptionsGestureRecognizer(); - abstract RunningMode runningMode(); abstract Optional numHands(); @@ -376,10 +467,11 @@ public final class GestureRecognizer extends BaseVisionTaskApi { abstract Optional minTrackingConfidence(); - // TODO update gesture confidence options after score merging calculator is ready. - abstract Optional minGestureConfidence(); + abstract Optional cannedGesturesClassifierOptions(); - abstract Optional> resultListener(); + abstract Optional customGesturesClassifierOptions(); + + abstract Optional> resultListener(); abstract Optional errorListener(); @@ -389,8 +481,7 @@ public final class GestureRecognizer extends BaseVisionTaskApi { .setNumHands(1) .setMinHandDetectionConfidence(0.5f) .setMinHandPresenceConfidence(0.5f) - .setMinTrackingConfidence(0.5f) - .setMinGestureConfidence(-1f); + .setMinTrackingConfidence(0.5f); } /** @@ -398,22 +489,18 @@ public final class GestureRecognizer extends BaseVisionTaskApi { */ @Override public CalculatorOptions convertToCalculatorOptionsProto() { - BaseOptionsProto.BaseOptions.Builder baseOptionsBuilder = - BaseOptionsProto.BaseOptions.newBuilder() - .setUseStreamMode(runningMode() != RunningMode.IMAGE) - .mergeFrom(convertBaseOptionsToProto(baseOptions())); GestureRecognizerGraphOptionsProto.GestureRecognizerGraphOptions.Builder taskOptionsBuilder = GestureRecognizerGraphOptionsProto.GestureRecognizerGraphOptions.newBuilder() - .setBaseOptions(baseOptionsBuilder); + .setBaseOptions( + BaseOptionsProto.BaseOptions.newBuilder() + .setUseStreamMode(runningMode() != RunningMode.IMAGE) + .mergeFrom(convertBaseOptionsToProto(baseOptions())) + .build()); // Setup HandDetectorGraphOptions. HandDetectorGraphOptionsProto.HandDetectorGraphOptions.Builder handDetectorGraphOptionsBuilder = - HandDetectorGraphOptionsProto.HandDetectorGraphOptions.newBuilder() - .setBaseOptions( - BaseOptionsProto.BaseOptions.newBuilder() - .setUseStreamMode(runningMode() != RunningMode.IMAGE) - .mergeFrom(convertBaseOptionsToProto(baseOptionsHandDetector()))); + HandDetectorGraphOptionsProto.HandDetectorGraphOptions.newBuilder(); numHands().ifPresent(handDetectorGraphOptionsBuilder::setNumHands); minHandDetectionConfidence() .ifPresent(handDetectorGraphOptionsBuilder::setMinDetectionConfidence); @@ -421,19 +508,12 @@ public final class GestureRecognizer extends BaseVisionTaskApi { // Setup HandLandmarkerGraphOptions. HandLandmarksDetectorGraphOptionsProto.HandLandmarksDetectorGraphOptions.Builder handLandmarksDetectorGraphOptionsBuilder = - HandLandmarksDetectorGraphOptionsProto.HandLandmarksDetectorGraphOptions.newBuilder() - .setBaseOptions( - BaseOptionsProto.BaseOptions.newBuilder() - .setUseStreamMode(runningMode() != RunningMode.IMAGE) - .mergeFrom(convertBaseOptionsToProto(baseOptionsHandLandmarker()))); + HandLandmarksDetectorGraphOptionsProto.HandLandmarksDetectorGraphOptions.newBuilder(); minHandPresenceConfidence() .ifPresent(handLandmarksDetectorGraphOptionsBuilder::setMinDetectionConfidence); HandLandmarkerGraphOptionsProto.HandLandmarkerGraphOptions.Builder handLandmarkerGraphOptionsBuilder = - HandLandmarkerGraphOptionsProto.HandLandmarkerGraphOptions.newBuilder() - .setBaseOptions( - BaseOptionsProto.BaseOptions.newBuilder() - .setUseStreamMode(runningMode() != RunningMode.IMAGE)); + HandLandmarkerGraphOptionsProto.HandLandmarkerGraphOptions.newBuilder(); minTrackingConfidence() .ifPresent(handLandmarkerGraphOptionsBuilder::setMinTrackingConfidence); handLandmarkerGraphOptionsBuilder @@ -443,17 +523,23 @@ public final class GestureRecognizer extends BaseVisionTaskApi { // Setup HandGestureRecognizerGraphOptions. HandGestureRecognizerGraphOptionsProto.HandGestureRecognizerGraphOptions.Builder handGestureRecognizerGraphOptionsBuilder = - HandGestureRecognizerGraphOptionsProto.HandGestureRecognizerGraphOptions.newBuilder() - .setBaseOptions( - BaseOptionsProto.BaseOptions.newBuilder() - .setUseStreamMode(runningMode() != RunningMode.IMAGE) - .mergeFrom(convertBaseOptionsToProto(baseOptionsGestureRecognizer()))); - ClassifierOptionsProto.ClassifierOptions.Builder classifierOptionsBuilder = - ClassifierOptionsProto.ClassifierOptions.newBuilder(); - minGestureConfidence().ifPresent(classifierOptionsBuilder::setScoreThreshold); - handGestureRecognizerGraphOptionsBuilder.setClassifierOptions( - classifierOptionsBuilder.build()); - + HandGestureRecognizerGraphOptionsProto.HandGestureRecognizerGraphOptions.newBuilder(); + cannedGesturesClassifierOptions() + .ifPresent( + classifierOptions -> { + handGestureRecognizerGraphOptionsBuilder.setCannedGestureClassifierGraphOptions( + GestureClassifierGraphOptionsProto.GestureClassifierGraphOptions.newBuilder() + .setClassifierOptions(classifierOptions.convertToProto()) + .build()); + }); + customGesturesClassifierOptions() + .ifPresent( + classifierOptions -> { + handGestureRecognizerGraphOptionsBuilder.setCustomGestureClassifierGraphOptions( + GestureClassifierGraphOptionsProto.GestureClassifierGraphOptions.newBuilder() + .setClassifierOptions(classifierOptions.convertToProto()) + .build()); + }); taskOptionsBuilder .setHandLandmarkerGraphOptions(handLandmarkerGraphOptionsBuilder.build()) .setHandGestureRecognizerGraphOptions(handGestureRecognizerGraphOptionsBuilder.build()); @@ -464,4 +550,15 @@ public final class GestureRecognizer extends BaseVisionTaskApi { .build(); } } + + /** + * Validates that the provided {@link ImageProcessingOptions} doesn't contain a + * region-of-interest. + */ + private static void validateImageProcessingOptions( + ImageProcessingOptions imageProcessingOptions) { + if (imageProcessingOptions.regionOfInterest().isPresent()) { + throw new IllegalArgumentException("GestureRecognizer doesn't support region-of-interest."); + } + } } diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassificationResult.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassificationResult.java index 09f854ca..d82a47b8 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassificationResult.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassificationResult.java @@ -15,11 +15,11 @@ package com.google.mediapipe.tasks.vision.imageclassifier; import com.google.auto.value.AutoValue; -import com.google.mediapipe.tasks.components.container.proto.CategoryProto; -import com.google.mediapipe.tasks.components.container.proto.ClassificationsProto; import com.google.mediapipe.tasks.components.containers.Category; import com.google.mediapipe.tasks.components.containers.ClassificationEntry; import com.google.mediapipe.tasks.components.containers.Classifications; +import com.google.mediapipe.tasks.components.containers.proto.CategoryProto; +import com.google.mediapipe.tasks.components.containers.proto.ClassificationsProto; import com.google.mediapipe.tasks.core.TaskResult; import java.util.ArrayList; import java.util.Collections; diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassifier.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassifier.java index 68cae151..3863b6fe 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassifier.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassifier.java @@ -15,7 +15,6 @@ package com.google.mediapipe.tasks.vision.imageclassifier; import android.content.Context; -import android.graphics.RectF; import android.os.ParcelFileDescriptor; import com.google.auto.value.AutoValue; import com.google.mediapipe.proto.CalculatorOptionsProto.CalculatorOptions; @@ -25,8 +24,8 @@ import com.google.mediapipe.framework.Packet; import com.google.mediapipe.framework.PacketGetter; import com.google.mediapipe.framework.ProtoUtil; import com.google.mediapipe.framework.image.BitmapImageBuilder; -import com.google.mediapipe.framework.image.Image; -import com.google.mediapipe.tasks.components.container.proto.ClassificationsProto; +import com.google.mediapipe.framework.image.MPImage; +import com.google.mediapipe.tasks.components.containers.proto.ClassificationsProto; import com.google.mediapipe.tasks.components.processors.ClassifierOptions; import com.google.mediapipe.tasks.core.BaseOptions; import com.google.mediapipe.tasks.core.ErrorListener; @@ -37,9 +36,9 @@ import com.google.mediapipe.tasks.core.TaskOptions; import com.google.mediapipe.tasks.core.TaskRunner; import com.google.mediapipe.tasks.core.proto.BaseOptionsProto; import com.google.mediapipe.tasks.vision.core.BaseVisionTaskApi; +import com.google.mediapipe.tasks.vision.core.ImageProcessingOptions; import com.google.mediapipe.tasks.vision.core.RunningMode; import com.google.mediapipe.tasks.vision.imageclassifier.proto.ImageClassifierGraphOptionsProto; -import com.google.protobuf.InvalidProtocolBufferException; import java.io.File; import java.io.IOException; import java.nio.ByteBuffer; @@ -165,9 +164,9 @@ public final class ImageClassifier extends BaseVisionTaskApi { * @throws MediaPipeException if there is an error during {@link ImageClassifier} creation. */ public static ImageClassifier createFromOptions(Context context, ImageClassifierOptions options) { - OutputHandler handler = new OutputHandler<>(); + OutputHandler handler = new OutputHandler<>(); handler.setOutputPacketConverter( - new OutputHandler.OutputPacketConverter() { + new OutputHandler.OutputPacketConverter() { @Override public ImageClassificationResult convertToTaskResult(List packets) { try { @@ -176,14 +175,14 @@ public final class ImageClassifier extends BaseVisionTaskApi { packets.get(CLASSIFICATION_RESULT_OUT_STREAM_INDEX), ClassificationsProto.ClassificationResult.getDefaultInstance()), packets.get(CLASSIFICATION_RESULT_OUT_STREAM_INDEX).getTimestamp()); - } catch (InvalidProtocolBufferException e) { + } catch (IOException e) { throw new MediaPipeException( MediaPipeException.StatusCode.INTERNAL.ordinal(), e.getMessage()); } } @Override - public Image convertToTaskInput(List packets) { + public MPImage convertToTaskInput(List packets) { return new BitmapImageBuilder( AndroidPacketGetter.getBitmapFromRgb(packets.get(IMAGE_OUT_STREAM_INDEX))) .build(); @@ -216,6 +215,24 @@ public final class ImageClassifier extends BaseVisionTaskApi { super(taskRunner, runningMode, IMAGE_IN_STREAM_NAME, NORM_RECT_IN_STREAM_NAME); } + /** + * Performs classification on the provided single image with default image processing options, + * i.e. using the whole image as region-of-interest and without any rotation applied. Only use + * this method when the {@link ImageClassifier} is created with {@link RunningMode.IMAGE}. + * + *

    {@link ImageClassifier} supports the following color space types: + * + *

      + *
    • {@link Bitmap.Config.ARGB_8888} + *
    + * + * @param image a MediaPipe {@link MPImage} object for processing. + * @throws MediaPipeException if there is an internal error. + */ + public ImageClassificationResult classify(MPImage image) { + return classify(image, ImageProcessingOptions.builder().build()); + } + /** * Performs classification on the provided single image. Only use this method when the {@link * ImageClassifier} is created with {@link RunningMode.IMAGE}. @@ -226,16 +243,23 @@ public final class ImageClassifier extends BaseVisionTaskApi { *
  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. * @throws MediaPipeException if there is an internal error. */ - public ImageClassificationResult classify(Image inputImage) { - return (ImageClassificationResult) processImageData(inputImage, buildFullImageRectF()); + public ImageClassificationResult classify( + MPImage image, ImageProcessingOptions imageProcessingOptions) { + return (ImageClassificationResult) processImageData(image, imageProcessingOptions); } /** - * Performs classification on the provided single image and region-of-interest. Only use this - * method when the {@link ImageClassifier} is created with {@link RunningMode.IMAGE}. + * Performs classification on the provided video frame with default image processing options, i.e. + * using the whole image as region-of-interest and without any rotation applied. Only use this + * method when the {@link ImageClassifier} is created with {@link RunningMode.VIDEO}. + * + *

    It's required to provide the video frame's timestamp (in milliseconds). The input timestamps + * must be monotonically increasing. * *

    {@link ImageClassifier} supports the following color space types: * @@ -243,13 +267,12 @@ public final class ImageClassifier extends BaseVisionTaskApi { *

  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. - * @param roi a {@link RectF} specifying the region of interest on which to perform - * classification. Coordinates are expected to be specified as normalized values in [0,1]. + * @param image a MediaPipe {@link MPImage} object for processing. + * @param timestampMs the input timestamp (in milliseconds). * @throws MediaPipeException if there is an internal error. */ - public ImageClassificationResult classify(Image inputImage, RectF roi) { - return (ImageClassificationResult) processImageData(inputImage, roi); + public ImageClassificationResult classifyForVideo(MPImage image, long timestampMs) { + return classifyForVideo(image, ImageProcessingOptions.builder().build(), timestampMs); } /** @@ -265,21 +288,26 @@ public final class ImageClassifier extends BaseVisionTaskApi { *
  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. - * @param inputTimestampMs the input timestamp (in milliseconds). + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. + * @param timestampMs the input timestamp (in milliseconds). * @throws MediaPipeException if there is an internal error. */ - public ImageClassificationResult classifyForVideo(Image inputImage, long inputTimestampMs) { - return (ImageClassificationResult) - processVideoData(inputImage, buildFullImageRectF(), inputTimestampMs); + public ImageClassificationResult classifyForVideo( + MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) { + return (ImageClassificationResult) processVideoData(image, imageProcessingOptions, timestampMs); } /** - * Performs classification on the provided video frame with additional region-of-interest. Only - * use this method when the {@link ImageClassifier} is created with {@link RunningMode.VIDEO}. + * Sends live image data to perform classification with default image processing options, i.e. + * using the whole image as region-of-interest and without any rotation applied, and the results + * will be available via the {@link ResultListener} provided in the {@link + * ImageClassifierOptions}. Only use this method when the {@link ImageClassifier} is created with + * {@link RunningMode.LIVE_STREAM}. * - *

    It's required to provide the video frame's timestamp (in milliseconds). The input timestamps - * must be monotonically increasing. + *

    It's required to provide a timestamp (in milliseconds) to indicate when the input image is + * sent to the object detector. The input timestamps must be monotonically increasing. * *

    {@link ImageClassifier} supports the following color space types: * @@ -287,15 +315,12 @@ public final class ImageClassifier extends BaseVisionTaskApi { *

  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. - * @param roi a {@link RectF} specifying the region of interest on which to perform - * classification. Coordinates are expected to be specified as normalized values in [0,1]. - * @param inputTimestampMs the input timestamp (in milliseconds). + * @param image a MediaPipe {@link MPImage} object for processing. + * @param timestampMs the input timestamp (in milliseconds). * @throws MediaPipeException if there is an internal error. */ - public ImageClassificationResult classifyForVideo( - Image inputImage, RectF roi, long inputTimestampMs) { - return (ImageClassificationResult) processVideoData(inputImage, roi, inputTimestampMs); + public void classifyAsync(MPImage image, long timestampMs) { + classifyAsync(image, ImageProcessingOptions.builder().build(), timestampMs); } /** @@ -312,37 +337,15 @@ public final class ImageClassifier extends BaseVisionTaskApi { *
  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. - * @param inputTimestampMs the input timestamp (in milliseconds). + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. + * @param timestampMs the input timestamp (in milliseconds). * @throws MediaPipeException if there is an internal error. */ - public void classifyAsync(Image inputImage, long inputTimestampMs) { - sendLiveStreamData(inputImage, buildFullImageRectF(), inputTimestampMs); - } - - /** - * Sends live image data and additional region-of-interest to perform classification, and the - * results will be available via the {@link ResultListener} provided in the {@link - * ImageClassifierOptions}. Only use this method when the {@link ImageClassifier} is created with - * {@link RunningMode.LIVE_STREAM}. - * - *

    It's required to provide a timestamp (in milliseconds) to indicate when the input image is - * sent to the object detector. The input timestamps must be monotonically increasing. - * - *

    {@link ImageClassifier} supports the following color space types: - * - *

      - *
    • {@link Bitmap.Config.ARGB_8888} - *
    - * - * @param inputImage a MediaPipe {@link Image} object for processing. - * @param roi a {@link RectF} specifying the region of interest on which to perform - * classification. Coordinates are expected to be specified as normalized values in [0,1]. - * @param inputTimestampMs the input timestamp (in milliseconds). - * @throws MediaPipeException if there is an internal error. - */ - public void classifyAsync(Image inputImage, RectF roi, long inputTimestampMs) { - sendLiveStreamData(inputImage, roi, inputTimestampMs); + public void classifyAsync( + MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) { + sendLiveStreamData(image, imageProcessingOptions, timestampMs); } /** Options for setting up and {@link ImageClassifier}. */ @@ -380,7 +383,7 @@ public final class ImageClassifier extends BaseVisionTaskApi { * the image classifier is in the live stream mode. */ public abstract Builder setResultListener( - ResultListener resultListener); + ResultListener resultListener); /** Sets an optional {@link ErrorListener}. */ public abstract Builder setErrorListener(ErrorListener errorListener); @@ -417,7 +420,7 @@ public final class ImageClassifier extends BaseVisionTaskApi { abstract Optional classifierOptions(); - abstract Optional> resultListener(); + abstract Optional> resultListener(); abstract Optional errorListener(); @@ -448,9 +451,4 @@ public final class ImageClassifier extends BaseVisionTaskApi { .build(); } } - - /** Creates a RectF covering the full image. */ - private static RectF buildFullImageRectF() { - return new RectF(0, 0, 1, 1); - } } diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/objectdetector/ObjectDetector.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/objectdetector/ObjectDetector.java index bfce6279..3f944eae 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/objectdetector/ObjectDetector.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/objectdetector/ObjectDetector.java @@ -22,7 +22,7 @@ import com.google.mediapipe.framework.AndroidPacketGetter; import com.google.mediapipe.framework.Packet; import com.google.mediapipe.framework.PacketGetter; import com.google.mediapipe.framework.image.BitmapImageBuilder; -import com.google.mediapipe.framework.image.Image; +import com.google.mediapipe.framework.image.MPImage; import com.google.mediapipe.tasks.core.BaseOptions; import com.google.mediapipe.tasks.core.ErrorListener; import com.google.mediapipe.tasks.core.OutputHandler; @@ -32,6 +32,7 @@ import com.google.mediapipe.tasks.core.TaskOptions; import com.google.mediapipe.tasks.core.TaskRunner; import com.google.mediapipe.tasks.core.proto.BaseOptionsProto; import com.google.mediapipe.tasks.vision.core.BaseVisionTaskApi; +import com.google.mediapipe.tasks.vision.core.ImageProcessingOptions; import com.google.mediapipe.tasks.vision.core.RunningMode; import com.google.mediapipe.tasks.vision.objectdetector.proto.ObjectDetectorOptionsProto; import com.google.mediapipe.formats.proto.DetectionProto.Detection; @@ -96,8 +97,10 @@ import java.util.Optional; public final class ObjectDetector extends BaseVisionTaskApi { private static final String TAG = ObjectDetector.class.getSimpleName(); private static final String IMAGE_IN_STREAM_NAME = "image_in"; + private static final String NORM_RECT_IN_STREAM_NAME = "norm_rect_in"; private static final List INPUT_STREAMS = - Collections.unmodifiableList(Arrays.asList("IMAGE:" + IMAGE_IN_STREAM_NAME)); + Collections.unmodifiableList( + Arrays.asList("IMAGE:" + IMAGE_IN_STREAM_NAME, "NORM_RECT:" + NORM_RECT_IN_STREAM_NAME)); private static final List OUTPUT_STREAMS = Collections.unmodifiableList(Arrays.asList("DETECTIONS:detections_out", "IMAGE:image_out")); private static final int DETECTIONS_OUT_STREAM_INDEX = 0; @@ -162,9 +165,9 @@ public final class ObjectDetector extends BaseVisionTaskApi { public static ObjectDetector createFromOptions( Context context, ObjectDetectorOptions detectorOptions) { // TODO: Consolidate OutputHandler and TaskRunner. - OutputHandler handler = new OutputHandler<>(); + OutputHandler handler = new OutputHandler<>(); handler.setOutputPacketConverter( - new OutputHandler.OutputPacketConverter() { + new OutputHandler.OutputPacketConverter() { @Override public ObjectDetectionResult convertToTaskResult(List packets) { return ObjectDetectionResult.create( @@ -174,7 +177,7 @@ public final class ObjectDetector extends BaseVisionTaskApi { } @Override - public Image convertToTaskInput(List packets) { + public MPImage convertToTaskInput(List packets) { return new BitmapImageBuilder( AndroidPacketGetter.getBitmapFromRgb(packets.get(IMAGE_OUT_STREAM_INDEX))) .build(); @@ -204,7 +207,25 @@ public final class ObjectDetector extends BaseVisionTaskApi { * @param runningMode a mediapipe vision task {@link RunningMode}. */ private ObjectDetector(TaskRunner taskRunner, RunningMode runningMode) { - super(taskRunner, runningMode, IMAGE_IN_STREAM_NAME); + super(taskRunner, runningMode, IMAGE_IN_STREAM_NAME, NORM_RECT_IN_STREAM_NAME); + } + + /** + * Performs object detection on the provided single image with default image processing options, + * i.e. without any rotation applied. Only use this method when the {@link ObjectDetector} is + * created with {@link RunningMode.IMAGE}. + * + *

    {@link ObjectDetector} supports the following color space types: + * + *

      + *
    • {@link Bitmap.Config.ARGB_8888} + *
    + * + * @param image a MediaPipe {@link MPImage} object for processing. + * @throws MediaPipeException if there is an internal error. + */ + public ObjectDetectionResult detect(MPImage image) { + return detect(image, ImageProcessingOptions.builder().build()); } /** @@ -217,11 +238,41 @@ public final class ObjectDetector extends BaseVisionTaskApi { *
  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. Note that region-of-interest is not supported + * by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in + * this method throwing an IllegalArgumentException. + * @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a + * region-of-interest. * @throws MediaPipeException if there is an internal error. */ - public ObjectDetectionResult detect(Image inputImage) { - return (ObjectDetectionResult) processImageData(inputImage); + public ObjectDetectionResult detect( + MPImage image, ImageProcessingOptions imageProcessingOptions) { + validateImageProcessingOptions(imageProcessingOptions); + return (ObjectDetectionResult) processImageData(image, imageProcessingOptions); + } + + /** + * Performs object detection on the provided video frame with default image processing options, + * i.e. without any rotation applied. Only use this method when the {@link ObjectDetector} is + * created with {@link RunningMode.VIDEO}. + * + *

    It's required to provide the video frame's timestamp (in milliseconds). The input timestamps + * must be monotonically increasing. + * + *

    {@link ObjectDetector} supports the following color space types: + * + *

      + *
    • {@link Bitmap.Config.ARGB_8888} + *
    + * + * @param image a MediaPipe {@link MPImage} object for processing. + * @param timestampMs the input timestamp (in milliseconds). + * @throws MediaPipeException if there is an internal error. + */ + public ObjectDetectionResult detectForVideo(MPImage image, long timestampMs) { + return detectForVideo(image, ImageProcessingOptions.builder().build(), timestampMs); } /** @@ -237,12 +288,43 @@ public final class ObjectDetector extends BaseVisionTaskApi { *
  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. - * @param inputTimestampMs the input timestamp (in milliseconds). + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. Note that region-of-interest is not supported + * by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in + * this method throwing an IllegalArgumentException. + * @param timestampMs the input timestamp (in milliseconds). + * @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a + * region-of-interest. * @throws MediaPipeException if there is an internal error. */ - public ObjectDetectionResult detectForVideo(Image inputImage, long inputTimestampMs) { - return (ObjectDetectionResult) processVideoData(inputImage, inputTimestampMs); + public ObjectDetectionResult detectForVideo( + MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) { + validateImageProcessingOptions(imageProcessingOptions); + return (ObjectDetectionResult) processVideoData(image, imageProcessingOptions, timestampMs); + } + + /** + * Sends live image data to perform object detection with default image processing options, i.e. + * without any rotation applied, and the results will be available via the {@link ResultListener} + * provided in the {@link ObjectDetectorOptions}. Only use this method when the {@link + * ObjectDetector} is created with {@link RunningMode.LIVE_STREAM}. + * + *

    It's required to provide a timestamp (in milliseconds) to indicate when the input image is + * sent to the object detector. The input timestamps must be monotonically increasing. + * + *

    {@link ObjectDetector} supports the following color space types: + * + *

      + *
    • {@link Bitmap.Config.ARGB_8888} + *
    + * + * @param image a MediaPipe {@link MPImage} object for processing. + * @param timestampMs the input timestamp (in milliseconds). + * @throws MediaPipeException if there is an internal error. + */ + public void detectAsync(MPImage image, long timestampMs) { + detectAsync(image, ImageProcessingOptions.builder().build(), timestampMs); } /** @@ -259,12 +341,20 @@ public final class ObjectDetector extends BaseVisionTaskApi { *
  • {@link Bitmap.Config.ARGB_8888} * * - * @param inputImage a MediaPipe {@link Image} object for processing. - * @param inputTimestampMs the input timestamp (in milliseconds). + * @param image a MediaPipe {@link MPImage} object for processing. + * @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the + * input image before running inference. Note that region-of-interest is not supported + * by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in + * this method throwing an IllegalArgumentException. + * @param timestampMs the input timestamp (in milliseconds). + * @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a + * region-of-interest. * @throws MediaPipeException if there is an internal error. */ - public void detectAsync(Image inputImage, long inputTimestampMs) { - sendLiveStreamData(inputImage, inputTimestampMs); + public void detectAsync( + MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) { + validateImageProcessingOptions(imageProcessingOptions); + sendLiveStreamData(image, imageProcessingOptions, timestampMs); } /** Options for setting up an {@link ObjectDetector}. */ @@ -333,7 +423,8 @@ public final class ObjectDetector extends BaseVisionTaskApi { * Sets the {@link ResultListener} to receive the detection results asynchronously when the * object detector is in the live stream mode. */ - public abstract Builder setResultListener(ResultListener value); + public abstract Builder setResultListener( + ResultListener value); /** Sets an optional {@link ErrorListener}}. */ public abstract Builder setErrorListener(ErrorListener value); @@ -378,7 +469,7 @@ public final class ObjectDetector extends BaseVisionTaskApi { abstract List categoryDenylist(); - abstract Optional> resultListener(); + abstract Optional> resultListener(); abstract Optional errorListener(); @@ -414,4 +505,15 @@ public final class ObjectDetector extends BaseVisionTaskApi { .build(); } } + + /** + * Validates that the provided {@link ImageProcessingOptions} doesn't contain a + * region-of-interest. + */ + private static void validateImageProcessingOptions( + ImageProcessingOptions imageProcessingOptions) { + if (imageProcessingOptions.regionOfInterest().isPresent()) { + throw new IllegalArgumentException("ObjectDetector doesn't support region-of-interest."); + } + } } diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/core/AndroidManifest.xml b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/core/AndroidManifest.xml new file mode 100644 index 00000000..aa2df6ba --- /dev/null +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/core/AndroidManifest.xml @@ -0,0 +1,24 @@ + + + + + + + + + + + + + + + diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/core/BUILD b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/core/BUILD new file mode 100644 index 00000000..a7f804c6 --- /dev/null +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/core/BUILD @@ -0,0 +1,19 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +package(default_visibility = ["//mediapipe/tasks:internal"]) + +licenses(["notice"]) + +# TODO: Enable this in OSS diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/core/ImageProcessingOptionsTest.java b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/core/ImageProcessingOptionsTest.java new file mode 100644 index 00000000..078b62af --- /dev/null +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/core/ImageProcessingOptionsTest.java @@ -0,0 +1,70 @@ +// Copyright 2022 The MediaPipe Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.mediapipe.tasks.vision.core; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import android.graphics.RectF; +import androidx.test.ext.junit.runners.AndroidJUnit4; +import org.junit.Test; +import org.junit.runner.RunWith; + +/** Test for {@link ImageProcessingOptions}/ */ +@RunWith(AndroidJUnit4.class) +public final class ImageProcessingOptionsTest { + + @Test + public void succeedsWithValidInputs() throws Exception { + ImageProcessingOptions options = + ImageProcessingOptions.builder() + .setRegionOfInterest(new RectF(0.0f, 0.1f, 1.0f, 0.9f)) + .setRotationDegrees(270) + .build(); + } + + @Test + public void failsWithLeftHigherThanRight() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + ImageProcessingOptions.builder() + .setRegionOfInterest(new RectF(0.9f, 0.0f, 0.1f, 1.0f)) + .build()); + assertThat(exception).hasMessageThat().contains("Expected left < right and top < bottom"); + } + + @Test + public void failsWithBottomHigherThanTop() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + ImageProcessingOptions.builder() + .setRegionOfInterest(new RectF(0.0f, 0.9f, 1.0f, 0.1f)) + .build()); + assertThat(exception).hasMessageThat().contains("Expected left < right and top < bottom"); + } + + @Test + public void failsWithInvalidRotation() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> ImageProcessingOptions.builder().setRotationDegrees(1).build()); + assertThat(exception).hasMessageThat().contains("Expected rotation to be a multiple of 90°"); + } +} diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognizerTest.java b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognizerTest.java index efec02b2..f76c4eaa 100644 --- a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognizerTest.java +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/gesturerecognizer/GestureRecognizerTest.java @@ -19,17 +19,20 @@ import static org.junit.Assert.assertThrows; import android.content.res.AssetManager; import android.graphics.BitmapFactory; +import android.graphics.RectF; import androidx.test.core.app.ApplicationProvider; import androidx.test.ext.junit.runners.AndroidJUnit4; import com.google.common.truth.Correspondence; import com.google.mediapipe.formats.proto.ClassificationProto; import com.google.mediapipe.framework.MediaPipeException; import com.google.mediapipe.framework.image.BitmapImageBuilder; -import com.google.mediapipe.framework.image.Image; +import com.google.mediapipe.framework.image.MPImage; import com.google.mediapipe.tasks.components.containers.Category; import com.google.mediapipe.tasks.components.containers.Landmark; import com.google.mediapipe.tasks.components.containers.proto.LandmarksDetectionResultProto.LandmarksDetectionResult; +import com.google.mediapipe.tasks.components.processors.ClassifierOptions; import com.google.mediapipe.tasks.core.BaseOptions; +import com.google.mediapipe.tasks.vision.core.ImageProcessingOptions; import com.google.mediapipe.tasks.vision.core.RunningMode; import com.google.mediapipe.tasks.vision.gesturerecognizer.GestureRecognizer.GestureRecognizerOptions; import java.io.InputStream; @@ -43,20 +46,25 @@ import org.junit.runners.Suite.SuiteClasses; @RunWith(Suite.class) @SuiteClasses({GestureRecognizerTest.General.class, GestureRecognizerTest.RunningModeTest.class}) public class GestureRecognizerTest { - private static final String HAND_DETECTOR_MODEL_FILE = "palm_detection_full.tflite"; - private static final String HAND_LANDMARKER_MODEL_FILE = "hand_landmark_full.tflite"; - private static final String GESTURE_RECOGNIZER_MODEL_FILE = - "cg_classifier_screen3d_landmark_features_nn_2022_08_04_base_simple_model.tflite"; + private static final String GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE = "gesture_recognizer.task"; + private static final String GESTURE_RECOGNIZER_WITH_CUSTOM_CLASSIFIER_BUNDLE_ASSET_FILE = + "gesture_recognizer_with_custom_classifier.task"; private static final String TWO_HANDS_IMAGE = "right_hands.jpg"; private static final String THUMB_UP_IMAGE = "thumb_up.jpg"; + private static final String POINTING_UP_ROTATED_IMAGE = "pointing_up_rotated.jpg"; private static final String NO_HANDS_IMAGE = "cats_and_dogs.jpg"; + private static final String FIST_IMAGE = "fist.jpg"; private static final String THUMB_UP_LANDMARKS = "thumb_up_landmarks.pb"; + private static final String FIST_LANDMARKS = "fist_landmarks.pb"; private static final String TAG = "Gesture Recognizer Test"; private static final String THUMB_UP_LABEL = "Thumb_Up"; - private static final int THUMB_UP_INDEX = 5; + private static final String POINTING_UP_LABEL = "Pointing_Up"; + private static final String FIST_LABEL = "Closed_Fist"; + private static final String ROCK_LABEL = "Rock"; private static final float LANDMARKS_ERROR_TOLERANCE = 0.03f; private static final int IMAGE_WIDTH = 382; private static final int IMAGE_HEIGHT = 406; + private static final int GESTURE_EXPECTED_INDEX = -1; @RunWith(AndroidJUnit4.class) public static final class General extends GestureRecognizerTest { @@ -66,20 +74,16 @@ public class GestureRecognizerTest { GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .build(); GestureRecognizer gestureRecognizer = GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); GestureRecognitionResult actualResult = gestureRecognizer.recognize(getImageFromAsset(THUMB_UP_IMAGE)); GestureRecognitionResult expectedResult = - getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL, THUMB_UP_INDEX); + getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL); assertActualResultApproximatelyEqualsToExpectedResult(actualResult, expectedResult); } @@ -88,13 +92,9 @@ public class GestureRecognizerTest { GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .build(); GestureRecognizer gestureRecognizer = GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); @@ -107,27 +107,22 @@ public class GestureRecognizerTest { } @Test - public void recognize_successWithMinGestureConfidence() throws Exception { + public void recognize_successWithScoreThreshold() throws Exception { GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - // TODO update the confidence to be in range [0,1] after embedding model - // and scoring calculator is integrated. - .setMinGestureConfidence(3.0f) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) + .setCannedGesturesClassifierOptions( + ClassifierOptions.builder().setScoreThreshold(0.5f).build()) .build(); GestureRecognizer gestureRecognizer = GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); GestureRecognitionResult actualResult = gestureRecognizer.recognize(getImageFromAsset(THUMB_UP_IMAGE)); GestureRecognitionResult expectedResult = - getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL, THUMB_UP_INDEX); + getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL); // Only contains one top scoring gesture. assertThat(actualResult.gestures().get(0)).hasSize(1); assertActualGestureEqualExpectedGesture( @@ -139,13 +134,9 @@ public class GestureRecognizerTest { GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .setNumHands(2) .build(); GestureRecognizer gestureRecognizer = @@ -154,6 +145,198 @@ public class GestureRecognizerTest { gestureRecognizer.recognize(getImageFromAsset(TWO_HANDS_IMAGE)); assertThat(actualResult.handednesses()).hasSize(2); } + + @Test + public void recognize_successWithRotation() throws Exception { + GestureRecognizerOptions options = + GestureRecognizerOptions.builder() + .setBaseOptions( + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) + .setNumHands(1) + .build(); + GestureRecognizer gestureRecognizer = + GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); + ImageProcessingOptions imageProcessingOptions = + ImageProcessingOptions.builder().setRotationDegrees(-90).build(); + GestureRecognitionResult actualResult = + gestureRecognizer.recognize( + getImageFromAsset(POINTING_UP_ROTATED_IMAGE), imageProcessingOptions); + assertThat(actualResult.gestures()).hasSize(1); + assertThat(actualResult.gestures().get(0).get(0).categoryName()).isEqualTo(POINTING_UP_LABEL); + } + + @Test + public void recognize_successWithCannedGestureFist() throws Exception { + GestureRecognizerOptions options = + GestureRecognizerOptions.builder() + .setBaseOptions( + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) + .setNumHands(1) + .build(); + GestureRecognizer gestureRecognizer = + GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); + GestureRecognitionResult actualResult = + gestureRecognizer.recognize(getImageFromAsset(FIST_IMAGE)); + GestureRecognitionResult expectedResult = + getExpectedGestureRecognitionResult(FIST_LANDMARKS, FIST_LABEL); + assertActualResultApproximatelyEqualsToExpectedResult(actualResult, expectedResult); + } + + @Test + public void recognize_successWithCustomGestureRock() throws Exception { + GestureRecognizerOptions options = + GestureRecognizerOptions.builder() + .setBaseOptions( + BaseOptions.builder() + .setModelAssetPath( + GESTURE_RECOGNIZER_WITH_CUSTOM_CLASSIFIER_BUNDLE_ASSET_FILE) + .build()) + .setNumHands(1) + .build(); + GestureRecognizer gestureRecognizer = + GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); + GestureRecognitionResult actualResult = + gestureRecognizer.recognize(getImageFromAsset(FIST_IMAGE)); + GestureRecognitionResult expectedResult = + getExpectedGestureRecognitionResult(FIST_LANDMARKS, ROCK_LABEL); + assertActualResultApproximatelyEqualsToExpectedResult(actualResult, expectedResult); + } + + @Test + public void recognize_successWithAllowGestureFist() throws Exception { + GestureRecognizerOptions options = + GestureRecognizerOptions.builder() + .setBaseOptions( + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) + .setNumHands(1) + .setCannedGesturesClassifierOptions( + ClassifierOptions.builder() + .setScoreThreshold(0.5f) + .setCategoryAllowlist(Arrays.asList("Closed_Fist")) + .build()) + .build(); + GestureRecognizer gestureRecognizer = + GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); + GestureRecognitionResult actualResult = + gestureRecognizer.recognize(getImageFromAsset(FIST_IMAGE)); + GestureRecognitionResult expectedResult = + getExpectedGestureRecognitionResult(FIST_LANDMARKS, FIST_LABEL); + assertActualResultApproximatelyEqualsToExpectedResult(actualResult, expectedResult); + } + + @Test + public void recognize_successWithDenyGestureFist() throws Exception { + GestureRecognizerOptions options = + GestureRecognizerOptions.builder() + .setBaseOptions( + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) + .setNumHands(1) + .setCannedGesturesClassifierOptions( + ClassifierOptions.builder() + .setScoreThreshold(0.5f) + .setCategoryDenylist(Arrays.asList("Closed_Fist")) + .build()) + .build(); + GestureRecognizer gestureRecognizer = + GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); + GestureRecognitionResult actualResult = + gestureRecognizer.recognize(getImageFromAsset(FIST_IMAGE)); + assertThat(actualResult.landmarks()).isEmpty(); + assertThat(actualResult.worldLandmarks()).isEmpty(); + assertThat(actualResult.handednesses()).isEmpty(); + assertThat(actualResult.gestures()).isEmpty(); + } + + @Test + public void recognize_successWithAllowAllGestureExceptFist() throws Exception { + GestureRecognizerOptions options = + GestureRecognizerOptions.builder() + .setBaseOptions( + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) + .setNumHands(1) + .setCannedGesturesClassifierOptions( + ClassifierOptions.builder() + .setScoreThreshold(0.5f) + .setCategoryAllowlist( + Arrays.asList( + "None", + "Open_Palm", + "Pointing_Up", + "Thumb_Down", + "Thumb_Up", + "Victory", + "ILoveYou")) + .build()) + .build(); + GestureRecognizer gestureRecognizer = + GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); + GestureRecognitionResult actualResult = + gestureRecognizer.recognize(getImageFromAsset(FIST_IMAGE)); + assertThat(actualResult.landmarks()).isEmpty(); + assertThat(actualResult.worldLandmarks()).isEmpty(); + assertThat(actualResult.handednesses()).isEmpty(); + assertThat(actualResult.gestures()).isEmpty(); + } + + @Test + public void recognize_successWithPreferAlowListThanDenyList() throws Exception { + GestureRecognizerOptions options = + GestureRecognizerOptions.builder() + .setBaseOptions( + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) + .setNumHands(1) + .setCannedGesturesClassifierOptions( + ClassifierOptions.builder() + .setScoreThreshold(0.5f) + .setCategoryAllowlist(Arrays.asList("Closed_Fist")) + .setCategoryDenylist(Arrays.asList("Closed_Fist")) + .build()) + .build(); + GestureRecognizer gestureRecognizer = + GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); + GestureRecognitionResult actualResult = + gestureRecognizer.recognize(getImageFromAsset(FIST_IMAGE)); + GestureRecognitionResult expectedResult = + getExpectedGestureRecognitionResult(FIST_LANDMARKS, FIST_LABEL); + assertActualResultApproximatelyEqualsToExpectedResult(actualResult, expectedResult); + } + + @Test + public void recognize_failsWithRegionOfInterest() throws Exception { + GestureRecognizerOptions options = + GestureRecognizerOptions.builder() + .setBaseOptions( + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) + .setNumHands(1) + .build(); + GestureRecognizer gestureRecognizer = + GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); + ImageProcessingOptions imageProcessingOptions = + ImageProcessingOptions.builder().setRegionOfInterest(new RectF(0, 0, 1, 1)).build(); + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + gestureRecognizer.recognize( + getImageFromAsset(THUMB_UP_IMAGE), imageProcessingOptions)); + assertThat(exception) + .hasMessageThat() + .contains("GestureRecognizer doesn't support region-of-interest"); + } } @RunWith(AndroidJUnit4.class) @@ -168,19 +351,7 @@ public class GestureRecognizerTest { GestureRecognizerOptions.builder() .setBaseOptions( BaseOptions.builder() - .setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE) - .build()) - .setBaseOptionsHandDetector( - BaseOptions.builder() - .setModelAssetPath(HAND_DETECTOR_MODEL_FILE) - .build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder() - .setModelAssetPath(HAND_LANDMARKER_MODEL_FILE) - .build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder() - .setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE) + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) .build()) .setRunningMode(mode) .setResultListener((gestureRecognitionResult, inputImage) -> {}) @@ -201,15 +372,7 @@ public class GestureRecognizerTest { GestureRecognizerOptions.builder() .setBaseOptions( BaseOptions.builder() - .setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE) - .build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder() - .setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE) + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) .build()) .setRunningMode(RunningMode.LIVE_STREAM) .build()); @@ -223,13 +386,9 @@ public class GestureRecognizerTest { GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .setRunningMode(RunningMode.IMAGE) .build(); @@ -238,12 +397,16 @@ public class GestureRecognizerTest { MediaPipeException exception = assertThrows( MediaPipeException.class, - () -> gestureRecognizer.recognizeForVideo(getImageFromAsset(THUMB_UP_IMAGE), 0)); + () -> + gestureRecognizer.recognizeForVideo( + getImageFromAsset(THUMB_UP_IMAGE), /*timestampsMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the video mode"); exception = assertThrows( MediaPipeException.class, - () -> gestureRecognizer.recognizeAsync(getImageFromAsset(THUMB_UP_IMAGE), 0)); + () -> + gestureRecognizer.recognizeAsync( + getImageFromAsset(THUMB_UP_IMAGE), /*timestampsMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the live stream mode"); } @@ -252,13 +415,9 @@ public class GestureRecognizerTest { GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .setRunningMode(RunningMode.VIDEO) .build(); @@ -272,7 +431,9 @@ public class GestureRecognizerTest { exception = assertThrows( MediaPipeException.class, - () -> gestureRecognizer.recognizeAsync(getImageFromAsset(THUMB_UP_IMAGE), 0)); + () -> + gestureRecognizer.recognizeAsync( + getImageFromAsset(THUMB_UP_IMAGE), /*timestampsMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the live stream mode"); } @@ -281,13 +442,9 @@ public class GestureRecognizerTest { GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .setRunningMode(RunningMode.LIVE_STREAM) .setResultListener((gestureRecognitionResult, inputImage) -> {}) .build(); @@ -302,7 +459,9 @@ public class GestureRecognizerTest { exception = assertThrows( MediaPipeException.class, - () -> gestureRecognizer.recognizeForVideo(getImageFromAsset(THUMB_UP_IMAGE), 0)); + () -> + gestureRecognizer.recognizeForVideo( + getImageFromAsset(THUMB_UP_IMAGE), /*timestampsMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the video mode"); } @@ -311,13 +470,9 @@ public class GestureRecognizerTest { GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .setRunningMode(RunningMode.IMAGE) .build(); @@ -326,7 +481,7 @@ public class GestureRecognizerTest { GestureRecognitionResult actualResult = gestureRecognizer.recognize(getImageFromAsset(THUMB_UP_IMAGE)); GestureRecognitionResult expectedResult = - getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL, THUMB_UP_INDEX); + getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL); assertActualResultApproximatelyEqualsToExpectedResult(actualResult, expectedResult); } @@ -335,41 +490,34 @@ public class GestureRecognizerTest { GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .setRunningMode(RunningMode.VIDEO) .build(); GestureRecognizer gestureRecognizer = GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options); GestureRecognitionResult expectedResult = - getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL, THUMB_UP_INDEX); + getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL); for (int i = 0; i < 3; i++) { GestureRecognitionResult actualResult = - gestureRecognizer.recognizeForVideo(getImageFromAsset(THUMB_UP_IMAGE), i); + gestureRecognizer.recognizeForVideo( + getImageFromAsset(THUMB_UP_IMAGE), /*timestampsMs=*/ i); assertActualResultApproximatelyEqualsToExpectedResult(actualResult, expectedResult); } } @Test public void recognize_failsWithOutOfOrderInputTimestamps() throws Exception { - Image image = getImageFromAsset(THUMB_UP_IMAGE); + MPImage image = getImageFromAsset(THUMB_UP_IMAGE); GestureRecognitionResult expectedResult = - getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL, THUMB_UP_INDEX); + getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL); GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .setRunningMode(RunningMode.LIVE_STREAM) .setResultListener( (actualResult, inputImage) -> { @@ -380,9 +528,11 @@ public class GestureRecognizerTest { .build(); try (GestureRecognizer gestureRecognizer = GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options)) { - gestureRecognizer.recognizeAsync(image, 1); + gestureRecognizer.recognizeAsync(image, /*timestampsMs=*/ 1); MediaPipeException exception = - assertThrows(MediaPipeException.class, () -> gestureRecognizer.recognizeAsync(image, 0)); + assertThrows( + MediaPipeException.class, + () -> gestureRecognizer.recognizeAsync(image, /*timestampsMs=*/ 0)); assertThat(exception) .hasMessageThat() .contains("having a smaller timestamp than the processed timestamp"); @@ -391,19 +541,15 @@ public class GestureRecognizerTest { @Test public void recognize_successWithLiveSteamMode() throws Exception { - Image image = getImageFromAsset(THUMB_UP_IMAGE); + MPImage image = getImageFromAsset(THUMB_UP_IMAGE); GestureRecognitionResult expectedResult = - getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL, THUMB_UP_INDEX); + getExpectedGestureRecognitionResult(THUMB_UP_LANDMARKS, THUMB_UP_LABEL); GestureRecognizerOptions options = GestureRecognizerOptions.builder() .setBaseOptions( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) - .setBaseOptionsHandDetector( - BaseOptions.builder().setModelAssetPath(HAND_DETECTOR_MODEL_FILE).build()) - .setBaseOptionsHandLandmarker( - BaseOptions.builder().setModelAssetPath(HAND_LANDMARKER_MODEL_FILE).build()) - .setBaseOptionsGestureRecognizer( - BaseOptions.builder().setModelAssetPath(GESTURE_RECOGNIZER_MODEL_FILE).build()) + BaseOptions.builder() + .setModelAssetPath(GESTURE_RECOGNIZER_BUNDLE_ASSET_FILE) + .build()) .setRunningMode(RunningMode.LIVE_STREAM) .setResultListener( (actualResult, inputImage) -> { @@ -415,19 +561,19 @@ public class GestureRecognizerTest { try (GestureRecognizer gestureRecognizer = GestureRecognizer.createFromOptions(ApplicationProvider.getApplicationContext(), options)) { for (int i = 0; i < 3; i++) { - gestureRecognizer.recognizeAsync(image, i); + gestureRecognizer.recognizeAsync(image, /*timestampsMs=*/ i); } } } - private static Image getImageFromAsset(String filePath) throws Exception { + private static MPImage getImageFromAsset(String filePath) throws Exception { AssetManager assetManager = ApplicationProvider.getApplicationContext().getAssets(); InputStream istr = assetManager.open(filePath); return new BitmapImageBuilder(BitmapFactory.decodeStream(istr)).build(); } private static GestureRecognitionResult getExpectedGestureRecognitionResult( - String filePath, String gestureLabel, int gestureIndex) throws Exception { + String filePath, String gestureLabel) throws Exception { AssetManager assetManager = ApplicationProvider.getApplicationContext().getAssets(); InputStream istr = assetManager.open(filePath); LandmarksDetectionResult landmarksDetectionResultProto = @@ -435,9 +581,7 @@ public class GestureRecognizerTest { ClassificationProto.ClassificationList gesturesProto = ClassificationProto.ClassificationList.newBuilder() .addClassification( - ClassificationProto.Classification.newBuilder() - .setLabel(gestureLabel) - .setIndex(gestureIndex)) + ClassificationProto.Classification.newBuilder().setLabel(gestureLabel)) .build(); return GestureRecognitionResult.create( Arrays.asList(landmarksDetectionResultProto.getLandmarks()), @@ -483,11 +627,11 @@ public class GestureRecognizerTest { private static void assertActualGestureEqualExpectedGesture( Category actualGesture, Category expectedGesture) { - assertThat(actualGesture.index()).isEqualTo(actualGesture.index()); - assertThat(expectedGesture.categoryName()).isEqualTo(expectedGesture.categoryName()); + assertThat(actualGesture.categoryName()).isEqualTo(expectedGesture.categoryName()); + assertThat(actualGesture.index()).isEqualTo(GESTURE_EXPECTED_INDEX); } - private static void assertImageSizeIsExpected(Image inputImage) { + private static void assertImageSizeIsExpected(MPImage inputImage) { assertThat(inputImage).isNotNull(); assertThat(inputImage.getWidth()).isEqualTo(IMAGE_WIDTH); assertThat(inputImage.getHeight()).isEqualTo(IMAGE_HEIGHT); diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassifierTest.java b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassifierTest.java index e02e8ebe..99ebd977 100644 --- a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassifierTest.java +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/imageclassifier/ImageClassifierTest.java @@ -24,11 +24,12 @@ import androidx.test.core.app.ApplicationProvider; import androidx.test.ext.junit.runners.AndroidJUnit4; import com.google.mediapipe.framework.MediaPipeException; import com.google.mediapipe.framework.image.BitmapImageBuilder; -import com.google.mediapipe.framework.image.Image; +import com.google.mediapipe.framework.image.MPImage; import com.google.mediapipe.tasks.components.containers.Category; import com.google.mediapipe.tasks.components.processors.ClassifierOptions; import com.google.mediapipe.tasks.core.BaseOptions; import com.google.mediapipe.tasks.core.TestUtils; +import com.google.mediapipe.tasks.vision.core.ImageProcessingOptions; import com.google.mediapipe.tasks.vision.core.RunningMode; import com.google.mediapipe.tasks.vision.imageclassifier.ImageClassifier.ImageClassifierOptions; import java.io.InputStream; @@ -47,7 +48,9 @@ public class ImageClassifierTest { private static final String FLOAT_MODEL_FILE = "mobilenet_v2_1.0_224.tflite"; private static final String QUANTIZED_MODEL_FILE = "mobilenet_v1_0.25_224_quant.tflite"; private static final String BURGER_IMAGE = "burger.jpg"; + private static final String BURGER_ROTATED_IMAGE = "burger_rotated.jpg"; private static final String MULTI_OBJECTS_IMAGE = "multi_objects.jpg"; + private static final String MULTI_OBJECTS_ROTATED_IMAGE = "multi_objects_rotated.jpg"; @RunWith(AndroidJUnit4.class) public static final class General extends ImageClassifierTest { @@ -209,13 +212,60 @@ public class ImageClassifierTest { ImageClassifier.createFromOptions(ApplicationProvider.getApplicationContext(), options); // RectF around the soccer ball. RectF roi = new RectF(0.450f, 0.308f, 0.614f, 0.734f); + ImageProcessingOptions imageProcessingOptions = + ImageProcessingOptions.builder().setRegionOfInterest(roi).build(); ImageClassificationResult results = - imageClassifier.classify(getImageFromAsset(MULTI_OBJECTS_IMAGE), roi); + imageClassifier.classify(getImageFromAsset(MULTI_OBJECTS_IMAGE), imageProcessingOptions); assertHasOneHeadAndOneTimestamp(results, 0); assertCategoriesAre( results, Arrays.asList(Category.create(0.9969325f, 806, "soccer ball", ""))); } + + @Test + public void classify_succeedsWithRotation() throws Exception { + ImageClassifierOptions options = + ImageClassifierOptions.builder() + .setBaseOptions(BaseOptions.builder().setModelAssetPath(FLOAT_MODEL_FILE).build()) + .setClassifierOptions(ClassifierOptions.builder().setMaxResults(3).build()) + .build(); + ImageClassifier imageClassifier = + ImageClassifier.createFromOptions(ApplicationProvider.getApplicationContext(), options); + ImageProcessingOptions imageProcessingOptions = + ImageProcessingOptions.builder().setRotationDegrees(-90).build(); + ImageClassificationResult results = + imageClassifier.classify(getImageFromAsset(BURGER_ROTATED_IMAGE), imageProcessingOptions); + + assertHasOneHeadAndOneTimestamp(results, 0); + assertCategoriesAre( + results, + Arrays.asList( + Category.create(0.6390683f, 934, "cheeseburger", ""), + Category.create(0.0495407f, 963, "meat loaf", ""), + Category.create(0.0469720f, 925, "guacamole", ""))); + } + + @Test + public void classify_succeedsWithRegionOfInterestAndRotation() throws Exception { + ImageClassifierOptions options = + ImageClassifierOptions.builder() + .setBaseOptions(BaseOptions.builder().setModelAssetPath(FLOAT_MODEL_FILE).build()) + .setClassifierOptions(ClassifierOptions.builder().setMaxResults(1).build()) + .build(); + ImageClassifier imageClassifier = + ImageClassifier.createFromOptions(ApplicationProvider.getApplicationContext(), options); + // RectF around the chair. + RectF roi = new RectF(0.0f, 0.1763f, 0.5642f, 0.3049f); + ImageProcessingOptions imageProcessingOptions = + ImageProcessingOptions.builder().setRegionOfInterest(roi).setRotationDegrees(-90).build(); + ImageClassificationResult results = + imageClassifier.classify( + getImageFromAsset(MULTI_OBJECTS_ROTATED_IMAGE), imageProcessingOptions); + + assertHasOneHeadAndOneTimestamp(results, 0); + assertCategoriesAre( + results, Arrays.asList(Category.create(0.686824f, 560, "folding chair", ""))); + } } @RunWith(AndroidJUnit4.class) @@ -269,12 +319,16 @@ public class ImageClassifierTest { MediaPipeException exception = assertThrows( MediaPipeException.class, - () -> imageClassifier.classifyForVideo(getImageFromAsset(BURGER_IMAGE), 0)); + () -> + imageClassifier.classifyForVideo( + getImageFromAsset(BURGER_IMAGE), /*timestampMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the video mode"); exception = assertThrows( MediaPipeException.class, - () -> imageClassifier.classifyAsync(getImageFromAsset(BURGER_IMAGE), 0)); + () -> + imageClassifier.classifyAsync( + getImageFromAsset(BURGER_IMAGE), /*timestampMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the live stream mode"); } @@ -296,7 +350,9 @@ public class ImageClassifierTest { exception = assertThrows( MediaPipeException.class, - () -> imageClassifier.classifyAsync(getImageFromAsset(BURGER_IMAGE), 0)); + () -> + imageClassifier.classifyAsync( + getImageFromAsset(BURGER_IMAGE), /*timestampMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the live stream mode"); } @@ -320,7 +376,9 @@ public class ImageClassifierTest { exception = assertThrows( MediaPipeException.class, - () -> imageClassifier.classifyForVideo(getImageFromAsset(BURGER_IMAGE), 0)); + () -> + imageClassifier.classifyForVideo( + getImageFromAsset(BURGER_IMAGE), /*timestampMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the video mode"); } @@ -342,7 +400,7 @@ public class ImageClassifierTest { @Test public void classify_succeedsWithVideoMode() throws Exception { - Image image = getImageFromAsset(BURGER_IMAGE); + MPImage image = getImageFromAsset(BURGER_IMAGE); ImageClassifierOptions options = ImageClassifierOptions.builder() .setBaseOptions(BaseOptions.builder().setModelAssetPath(FLOAT_MODEL_FILE).build()) @@ -352,7 +410,8 @@ public class ImageClassifierTest { ImageClassifier imageClassifier = ImageClassifier.createFromOptions(ApplicationProvider.getApplicationContext(), options); for (int i = 0; i < 3; i++) { - ImageClassificationResult results = imageClassifier.classifyForVideo(image, i); + ImageClassificationResult results = + imageClassifier.classifyForVideo(image, /*timestampMs=*/ i); assertHasOneHeadAndOneTimestamp(results, i); assertCategoriesAre( results, Arrays.asList(Category.create(0.7952058f, 934, "cheeseburger", ""))); @@ -361,7 +420,7 @@ public class ImageClassifierTest { @Test public void classify_failsWithOutOfOrderInputTimestamps() throws Exception { - Image image = getImageFromAsset(BURGER_IMAGE); + MPImage image = getImageFromAsset(BURGER_IMAGE); ImageClassifierOptions options = ImageClassifierOptions.builder() .setBaseOptions(BaseOptions.builder().setModelAssetPath(FLOAT_MODEL_FILE).build()) @@ -377,9 +436,11 @@ public class ImageClassifierTest { .build(); try (ImageClassifier imageClassifier = ImageClassifier.createFromOptions(ApplicationProvider.getApplicationContext(), options)) { - imageClassifier.classifyAsync(getImageFromAsset(BURGER_IMAGE), 1); + imageClassifier.classifyAsync(getImageFromAsset(BURGER_IMAGE), /*timestampMs=*/ 1); MediaPipeException exception = - assertThrows(MediaPipeException.class, () -> imageClassifier.classifyAsync(image, 0)); + assertThrows( + MediaPipeException.class, + () -> imageClassifier.classifyAsync(image, /*timestampMs=*/ 0)); assertThat(exception) .hasMessageThat() .contains("having a smaller timestamp than the processed timestamp"); @@ -388,7 +449,7 @@ public class ImageClassifierTest { @Test public void classify_succeedsWithLiveStreamMode() throws Exception { - Image image = getImageFromAsset(BURGER_IMAGE); + MPImage image = getImageFromAsset(BURGER_IMAGE); ImageClassifierOptions options = ImageClassifierOptions.builder() .setBaseOptions(BaseOptions.builder().setModelAssetPath(FLOAT_MODEL_FILE).build()) @@ -405,13 +466,13 @@ public class ImageClassifierTest { try (ImageClassifier imageClassifier = ImageClassifier.createFromOptions(ApplicationProvider.getApplicationContext(), options)) { for (int i = 0; i < 3; ++i) { - imageClassifier.classifyAsync(image, i); + imageClassifier.classifyAsync(image, /*timestampMs=*/ i); } } } } - private static Image getImageFromAsset(String filePath) throws Exception { + private static MPImage getImageFromAsset(String filePath) throws Exception { AssetManager assetManager = ApplicationProvider.getApplicationContext().getAssets(); InputStream istr = assetManager.open(filePath); return new BitmapImageBuilder(BitmapFactory.decodeStream(istr)).build(); @@ -437,7 +498,7 @@ public class ImageClassifierTest { } } - private static void assertImageSizeIsExpected(Image inputImage) { + private static void assertImageSizeIsExpected(MPImage inputImage) { assertThat(inputImage).isNotNull(); assertThat(inputImage.getWidth()).isEqualTo(480); assertThat(inputImage.getHeight()).isEqualTo(325); diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/objectdetector/ObjectDetectorTest.java b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/objectdetector/ObjectDetectorTest.java index cdec57d7..2878c380 100644 --- a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/objectdetector/ObjectDetectorTest.java +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/vision/objectdetector/ObjectDetectorTest.java @@ -24,11 +24,12 @@ import androidx.test.core.app.ApplicationProvider; import androidx.test.ext.junit.runners.AndroidJUnit4; import com.google.mediapipe.framework.MediaPipeException; import com.google.mediapipe.framework.image.BitmapImageBuilder; -import com.google.mediapipe.framework.image.Image; +import com.google.mediapipe.framework.image.MPImage; import com.google.mediapipe.tasks.components.containers.Category; import com.google.mediapipe.tasks.components.containers.Detection; import com.google.mediapipe.tasks.core.BaseOptions; import com.google.mediapipe.tasks.core.TestUtils; +import com.google.mediapipe.tasks.vision.core.ImageProcessingOptions; import com.google.mediapipe.tasks.vision.core.RunningMode; import com.google.mediapipe.tasks.vision.objectdetector.ObjectDetector.ObjectDetectorOptions; import java.io.InputStream; @@ -45,10 +46,11 @@ import org.junit.runners.Suite.SuiteClasses; public class ObjectDetectorTest { private static final String MODEL_FILE = "coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.tflite"; private static final String CAT_AND_DOG_IMAGE = "cats_and_dogs.jpg"; + private static final String CAT_AND_DOG_ROTATED_IMAGE = "cats_and_dogs_rotated.jpg"; private static final int IMAGE_WIDTH = 1200; private static final int IMAGE_HEIGHT = 600; private static final float CAT_SCORE = 0.69f; - private static final RectF catBoundingBox = new RectF(611, 164, 986, 596); + private static final RectF CAT_BOUNDING_BOX = new RectF(611, 164, 986, 596); // TODO: Figure out why android_x86 and android_arm tests have slightly different // scores (0.6875 vs 0.69921875). private static final float SCORE_DIFF_TOLERANCE = 0.01f; @@ -67,7 +69,7 @@ public class ObjectDetectorTest { ObjectDetector objectDetector = ObjectDetector.createFromOptions(ApplicationProvider.getApplicationContext(), options); ObjectDetectionResult results = objectDetector.detect(getImageFromAsset(CAT_AND_DOG_IMAGE)); - assertContainsOnlyCat(results, catBoundingBox, CAT_SCORE); + assertContainsOnlyCat(results, CAT_BOUNDING_BOX, CAT_SCORE); } @Test @@ -104,7 +106,7 @@ public class ObjectDetectorTest { ObjectDetector.createFromOptions(ApplicationProvider.getApplicationContext(), options); ObjectDetectionResult results = objectDetector.detect(getImageFromAsset(CAT_AND_DOG_IMAGE)); // The score threshold should block all other other objects, except cat. - assertContainsOnlyCat(results, catBoundingBox, CAT_SCORE); + assertContainsOnlyCat(results, CAT_BOUNDING_BOX, CAT_SCORE); } @Test @@ -175,7 +177,7 @@ public class ObjectDetectorTest { ObjectDetector objectDetector = ObjectDetector.createFromOptions(ApplicationProvider.getApplicationContext(), options); ObjectDetectionResult results = objectDetector.detect(getImageFromAsset(CAT_AND_DOG_IMAGE)); - assertContainsOnlyCat(results, catBoundingBox, CAT_SCORE); + assertContainsOnlyCat(results, CAT_BOUNDING_BOX, CAT_SCORE); } @Test @@ -228,6 +230,46 @@ public class ObjectDetectorTest { .contains("`category_allowlist` and `category_denylist` are mutually exclusive options."); } + @Test + public void detect_succeedsWithRotation() throws Exception { + ObjectDetectorOptions options = + ObjectDetectorOptions.builder() + .setBaseOptions(BaseOptions.builder().setModelAssetPath(MODEL_FILE).build()) + .setMaxResults(1) + .setCategoryAllowlist(Arrays.asList("cat")) + .build(); + ObjectDetector objectDetector = + ObjectDetector.createFromOptions(ApplicationProvider.getApplicationContext(), options); + ImageProcessingOptions imageProcessingOptions = + ImageProcessingOptions.builder().setRotationDegrees(-90).build(); + ObjectDetectionResult results = + objectDetector.detect( + getImageFromAsset(CAT_AND_DOG_ROTATED_IMAGE), imageProcessingOptions); + + assertContainsOnlyCat(results, new RectF(22.0f, 611.0f, 452.0f, 890.0f), 0.7109375f); + } + + @Test + public void detect_failsWithRegionOfInterest() throws Exception { + ObjectDetectorOptions options = + ObjectDetectorOptions.builder() + .setBaseOptions(BaseOptions.builder().setModelAssetPath(MODEL_FILE).build()) + .build(); + ObjectDetector objectDetector = + ObjectDetector.createFromOptions(ApplicationProvider.getApplicationContext(), options); + ImageProcessingOptions imageProcessingOptions = + ImageProcessingOptions.builder().setRegionOfInterest(new RectF(0, 0, 1, 1)).build(); + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + objectDetector.detect( + getImageFromAsset(CAT_AND_DOG_IMAGE), imageProcessingOptions)); + assertThat(exception) + .hasMessageThat() + .contains("ObjectDetector doesn't support region-of-interest"); + } + // TODO: Implement detect_succeedsWithFloatImages, detect_succeedsWithOrientation, // detect_succeedsWithNumThreads, detect_successWithNumThreadsFromBaseOptions, // detect_failsWithInvalidNegativeNumThreads, detect_failsWithInvalidNumThreadsAsZero. @@ -282,12 +324,16 @@ public class ObjectDetectorTest { MediaPipeException exception = assertThrows( MediaPipeException.class, - () -> objectDetector.detectForVideo(getImageFromAsset(CAT_AND_DOG_IMAGE), 0)); + () -> + objectDetector.detectForVideo( + getImageFromAsset(CAT_AND_DOG_IMAGE), /*timestampsMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the video mode"); exception = assertThrows( MediaPipeException.class, - () -> objectDetector.detectAsync(getImageFromAsset(CAT_AND_DOG_IMAGE), 0)); + () -> + objectDetector.detectAsync( + getImageFromAsset(CAT_AND_DOG_IMAGE), /*timestampsMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the live stream mode"); } @@ -309,7 +355,9 @@ public class ObjectDetectorTest { exception = assertThrows( MediaPipeException.class, - () -> objectDetector.detectAsync(getImageFromAsset(CAT_AND_DOG_IMAGE), 0)); + () -> + objectDetector.detectAsync( + getImageFromAsset(CAT_AND_DOG_IMAGE), /*timestampsMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the live stream mode"); } @@ -333,7 +381,9 @@ public class ObjectDetectorTest { exception = assertThrows( MediaPipeException.class, - () -> objectDetector.detectForVideo(getImageFromAsset(CAT_AND_DOG_IMAGE), 0)); + () -> + objectDetector.detectForVideo( + getImageFromAsset(CAT_AND_DOG_IMAGE), /*timestampsMs=*/ 0)); assertThat(exception).hasMessageThat().contains("not initialized with the video mode"); } @@ -348,7 +398,7 @@ public class ObjectDetectorTest { ObjectDetector objectDetector = ObjectDetector.createFromOptions(ApplicationProvider.getApplicationContext(), options); ObjectDetectionResult results = objectDetector.detect(getImageFromAsset(CAT_AND_DOG_IMAGE)); - assertContainsOnlyCat(results, catBoundingBox, CAT_SCORE); + assertContainsOnlyCat(results, CAT_BOUNDING_BOX, CAT_SCORE); } @Test @@ -363,30 +413,33 @@ public class ObjectDetectorTest { ObjectDetector.createFromOptions(ApplicationProvider.getApplicationContext(), options); for (int i = 0; i < 3; i++) { ObjectDetectionResult results = - objectDetector.detectForVideo(getImageFromAsset(CAT_AND_DOG_IMAGE), i); - assertContainsOnlyCat(results, catBoundingBox, CAT_SCORE); + objectDetector.detectForVideo( + getImageFromAsset(CAT_AND_DOG_IMAGE), /*timestampsMs=*/ i); + assertContainsOnlyCat(results, CAT_BOUNDING_BOX, CAT_SCORE); } } @Test public void detect_failsWithOutOfOrderInputTimestamps() throws Exception { - Image image = getImageFromAsset(CAT_AND_DOG_IMAGE); + MPImage image = getImageFromAsset(CAT_AND_DOG_IMAGE); ObjectDetectorOptions options = ObjectDetectorOptions.builder() .setBaseOptions(BaseOptions.builder().setModelAssetPath(MODEL_FILE).build()) .setRunningMode(RunningMode.LIVE_STREAM) .setResultListener( (objectDetectionResult, inputImage) -> { - assertContainsOnlyCat(objectDetectionResult, catBoundingBox, CAT_SCORE); + assertContainsOnlyCat(objectDetectionResult, CAT_BOUNDING_BOX, CAT_SCORE); assertImageSizeIsExpected(inputImage); }) .setMaxResults(1) .build(); try (ObjectDetector objectDetector = ObjectDetector.createFromOptions(ApplicationProvider.getApplicationContext(), options)) { - objectDetector.detectAsync(image, 1); + objectDetector.detectAsync(image, /*timestampsMs=*/ 1); MediaPipeException exception = - assertThrows(MediaPipeException.class, () -> objectDetector.detectAsync(image, 0)); + assertThrows( + MediaPipeException.class, + () -> objectDetector.detectAsync(image, /*timestampsMs=*/ 0)); assertThat(exception) .hasMessageThat() .contains("having a smaller timestamp than the processed timestamp"); @@ -395,14 +448,14 @@ public class ObjectDetectorTest { @Test public void detect_successWithLiveSteamMode() throws Exception { - Image image = getImageFromAsset(CAT_AND_DOG_IMAGE); + MPImage image = getImageFromAsset(CAT_AND_DOG_IMAGE); ObjectDetectorOptions options = ObjectDetectorOptions.builder() .setBaseOptions(BaseOptions.builder().setModelAssetPath(MODEL_FILE).build()) .setRunningMode(RunningMode.LIVE_STREAM) .setResultListener( (objectDetectionResult, inputImage) -> { - assertContainsOnlyCat(objectDetectionResult, catBoundingBox, CAT_SCORE); + assertContainsOnlyCat(objectDetectionResult, CAT_BOUNDING_BOX, CAT_SCORE); assertImageSizeIsExpected(inputImage); }) .setMaxResults(1) @@ -410,13 +463,13 @@ public class ObjectDetectorTest { try (ObjectDetector objectDetector = ObjectDetector.createFromOptions(ApplicationProvider.getApplicationContext(), options)) { for (int i = 0; i < 3; i++) { - objectDetector.detectAsync(image, i); + objectDetector.detectAsync(image, /*timestampsMs=*/ i); } } } } - private static Image getImageFromAsset(String filePath) throws Exception { + private static MPImage getImageFromAsset(String filePath) throws Exception { AssetManager assetManager = ApplicationProvider.getApplicationContext().getAssets(); InputStream istr = assetManager.open(filePath); return new BitmapImageBuilder(BitmapFactory.decodeStream(istr)).build(); @@ -448,7 +501,7 @@ public class ObjectDetectorTest { assertThat(boundingBox1.bottom).isWithin(PIXEL_DIFF_TOLERANCE).of(boundingBox2.bottom); } - private static void assertImageSizeIsExpected(Image inputImage) { + private static void assertImageSizeIsExpected(MPImage inputImage) { assertThat(inputImage).isNotNull(); assertThat(inputImage.getWidth()).isEqualTo(IMAGE_WIDTH); assertThat(inputImage.getHeight()).isEqualTo(IMAGE_HEIGHT); diff --git a/mediapipe/tasks/python/__init__.py b/mediapipe/tasks/python/__init__.py index 65c1214a..77683914 100644 --- a/mediapipe/tasks/python/__init__.py +++ b/mediapipe/tasks/python/__init__.py @@ -11,3 +11,14 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + +"""MediaPipe Tasks API.""" + +from . import components +from . import core +from . import vision + +BaseOptions = core.base_options.BaseOptions + +# Remove unnecessary modules to avoid duplication in API docs. +del core diff --git a/mediapipe/tasks/python/audio/core/BUILD b/mediapipe/tasks/python/audio/core/BUILD new file mode 100644 index 00000000..d9f169c6 --- /dev/null +++ b/mediapipe/tasks/python/audio/core/BUILD @@ -0,0 +1,37 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Placeholder for internal Python strict library and test compatibility macro. + +package(default_visibility = ["//mediapipe/tasks:internal"]) + +licenses(["notice"]) + +py_library( + name = "audio_task_running_mode", + srcs = ["audio_task_running_mode.py"], +) + +py_library( + name = "base_audio_task_api", + srcs = [ + "base_audio_task_api.py", + ], + deps = [ + ":audio_task_running_mode", + "//mediapipe/framework:calculator_py_pb2", + "//mediapipe/python:_framework_bindings", + "//mediapipe/tasks/python/core:optional_dependencies", + ], +) diff --git a/mediapipe/tasks/python/audio/core/__init__.py b/mediapipe/tasks/python/audio/core/__init__.py new file mode 100644 index 00000000..6a840518 --- /dev/null +++ b/mediapipe/tasks/python/audio/core/__init__.py @@ -0,0 +1,16 @@ +"""Copyright 2022 The MediaPipe Authors. + +All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" diff --git a/mediapipe/tasks/python/audio/core/audio_task_running_mode.py b/mediapipe/tasks/python/audio/core/audio_task_running_mode.py new file mode 100644 index 00000000..0fa36d40 --- /dev/null +++ b/mediapipe/tasks/python/audio/core/audio_task_running_mode.py @@ -0,0 +1,29 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The running mode of MediaPipe Audio Tasks.""" + +import enum + + +class AudioTaskRunningMode(enum.Enum): + """MediaPipe audio task running mode. + + Attributes: + AUDIO_CLIPS: The mode for running a mediapipe audio task on independent + audio clips. + AUDIO_STREAM: The mode for running a mediapipe audio task on an audio + stream, such as from microphone. + """ + AUDIO_CLIPS = 'AUDIO_CLIPS' + AUDIO_STREAM = 'AUDIO_STREAM' diff --git a/mediapipe/tasks/python/audio/core/base_audio_task_api.py b/mediapipe/tasks/python/audio/core/base_audio_task_api.py new file mode 100644 index 00000000..b6a2e0e4 --- /dev/null +++ b/mediapipe/tasks/python/audio/core/base_audio_task_api.py @@ -0,0 +1,123 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MediaPipe audio task base api.""" + +from typing import Callable, Mapping, Optional + +from mediapipe.framework import calculator_pb2 +from mediapipe.python._framework_bindings import packet as packet_module +from mediapipe.python._framework_bindings import task_runner as task_runner_module +from mediapipe.tasks.python.audio.core import audio_task_running_mode as running_mode_module +from mediapipe.tasks.python.core.optional_dependencies import doc_controls + +_TaskRunner = task_runner_module.TaskRunner +_Packet = packet_module.Packet +_RunningMode = running_mode_module.AudioTaskRunningMode + + +class BaseAudioTaskApi(object): + """The base class of the user-facing mediapipe audio task api classes.""" + + def __init__( + self, + graph_config: calculator_pb2.CalculatorGraphConfig, + running_mode: _RunningMode, + packet_callback: Optional[Callable[[Mapping[str, packet_module.Packet]], + None]] = None + ) -> None: + """Initializes the `BaseAudioTaskApi` object. + + Args: + graph_config: The mediapipe audio task graph config proto. + running_mode: The running mode of the mediapipe audio task. + packet_callback: The optional packet callback for getting results + asynchronously in the audio stream mode. + + Raises: + ValueError: The packet callback is not properly set based on the task's + running mode. + """ + if running_mode == _RunningMode.AUDIO_STREAM: + if packet_callback is None: + raise ValueError( + 'The audio task is in audio stream mode, a user-defined result ' + 'callback must be provided.') + elif packet_callback: + raise ValueError( + 'The audio task is in audio clips mode, a user-defined result ' + 'callback should not be provided.') + self._runner = _TaskRunner.create(graph_config, packet_callback) + self._running_mode = running_mode + + def _process_audio_clip( + self, inputs: Mapping[str, _Packet]) -> Mapping[str, _Packet]: + """A synchronous method to process independent audio clips. + + The call blocks the current thread until a failure status or a successful + result is returned. + + Args: + inputs: A dict contains (input stream name, data packet) pairs. + + Returns: + A dict contains (output stream name, data packet) pairs. + + Raises: + ValueError: If the task's running mode is not set to audio clips mode. + """ + if self._running_mode != _RunningMode.AUDIO_CLIPS: + raise ValueError( + 'Task is not initialized with the audio clips mode. Current running mode:' + + self._running_mode.name) + return self._runner.process(inputs) + + def _send_audio_stream_data(self, inputs: Mapping[str, _Packet]) -> None: + """An asynchronous method to send audio stream data to the runner. + + The results will be available in the user-defined results callback. + + Args: + inputs: A dict contains (input stream name, data packet) pairs. + + Raises: + ValueError: If the task's running mode is not set to the audio stream + mode. + """ + if self._running_mode != _RunningMode.AUDIO_STREAM: + raise ValueError( + 'Task is not initialized with the audio stream mode. Current running mode:' + + self._running_mode.name) + self._runner.send(inputs) + + def close(self) -> None: + """Shuts down the mediapipe audio task instance. + + Raises: + RuntimeError: If the mediapipe audio task failed to close. + """ + self._runner.close() + + @doc_controls.do_not_generate_docs + def __enter__(self): + """Return `self` upon entering the runtime context.""" + return self + + @doc_controls.do_not_generate_docs + def __exit__(self, unused_exc_type, unused_exc_value, unused_traceback): + """Shuts down the mediapipe audio task instance on exit of the context manager. + + Raises: + RuntimeError: If the mediapipe audio task failed to close. + """ + self.close() diff --git a/mediapipe/tasks/python/components/containers/BUILD b/mediapipe/tasks/python/components/containers/BUILD index cb123562..cb398be7 100644 --- a/mediapipe/tasks/python/components/containers/BUILD +++ b/mediapipe/tasks/python/components/containers/BUILD @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Placeholder for internal Python strict library compatibility macro. +# Placeholder for internal Python strict library and test compatibility macro. package(default_visibility = ["//mediapipe/tasks:internal"]) @@ -36,6 +36,29 @@ py_library( ], ) +py_library( + name = "landmark", + srcs = ["landmark.py"], + deps = [ + "//mediapipe/framework/formats:landmark_py_pb2", + "//mediapipe/tasks/python/core:optional_dependencies", + ], +) + +py_library( + name = "landmark_detection_result", + srcs = ["landmark_detection_result.py"], + deps = [ + ":landmark", + ":rect", + "//mediapipe/framework/formats:classification_py_pb2", + "//mediapipe/framework/formats:landmark_py_pb2", + "//mediapipe/tasks/cc/components/containers/proto:landmarks_detection_result_py_pb2", + "//mediapipe/tasks/python/components/containers:category", + "//mediapipe/tasks/python/core:optional_dependencies", + ], +) + py_library( name = "category", srcs = ["category.py"], @@ -57,6 +80,16 @@ py_library( ], ) +py_library( + name = "classifications", + srcs = ["classifications.py"], + deps = [ + ":category", + "//mediapipe/tasks/cc/components/containers/proto:classifications_py_pb2", + "//mediapipe/tasks/python/core:optional_dependencies", + ], +) + py_library( name = "embeddings", srcs = ["embeddings.py"], diff --git a/mediapipe/tasks/python/components/containers/category.py b/mediapipe/tasks/python/components/containers/category.py index 0b347fc1..cfdb8374 100644 --- a/mediapipe/tasks/python/components/containers/category.py +++ b/mediapipe/tasks/python/components/containers/category.py @@ -14,7 +14,7 @@ """Category data class.""" import dataclasses -from typing import Any +from typing import Any, Optional from mediapipe.tasks.cc.components.containers.proto import category_pb2 from mediapipe.tasks.python.core.optional_dependencies import doc_controls @@ -39,10 +39,10 @@ class Category: category_name: The label of this category object. """ - index: int - score: float - display_name: str - category_name: str + index: Optional[int] = None + score: Optional[float] = None + display_name: Optional[str] = None + category_name: Optional[str] = None @doc_controls.do_not_generate_docs def to_pb2(self) -> _CategoryProto: diff --git a/mediapipe/tasks/python/components/containers/classifications.py b/mediapipe/tasks/python/components/containers/classifications.py new file mode 100644 index 00000000..90ab2261 --- /dev/null +++ b/mediapipe/tasks/python/components/containers/classifications.py @@ -0,0 +1,168 @@ +# Copyright 2022 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Classifications data class.""" + +import dataclasses +from typing import Any, List, Optional + +from mediapipe.tasks.cc.components.containers.proto import classifications_pb2 +from mediapipe.tasks.python.components.containers import category as category_module +from mediapipe.tasks.python.core.optional_dependencies import doc_controls + +_ClassificationEntryProto = classifications_pb2.ClassificationEntry +_ClassificationsProto = classifications_pb2.Classifications +_ClassificationResultProto = classifications_pb2.ClassificationResult + + +@dataclasses.dataclass +class ClassificationEntry: + """List of predicted classes (aka labels) for a given classifier head. + + Attributes: + categories: The array of predicted categories, usually sorted by descending + scores (e.g. from high to low probability). + timestamp_ms: The optional timestamp (in milliseconds) associated to the + classification entry. This is useful for time series use cases, e.g., + audio classification. + """ + + categories: List[category_module.Category] + timestamp_ms: Optional[int] = None + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _ClassificationEntryProto: + """Generates a ClassificationEntry protobuf object.""" + return _ClassificationEntryProto( + categories=[category.to_pb2() for category in self.categories], + timestamp_ms=self.timestamp_ms) + + @classmethod + @doc_controls.do_not_generate_docs + def create_from_pb2( + cls, pb2_obj: _ClassificationEntryProto) -> 'ClassificationEntry': + """Creates a `ClassificationEntry` object from the given protobuf object.""" + return ClassificationEntry( + categories=[ + category_module.Category.create_from_pb2(category) + for category in pb2_obj.categories + ], + timestamp_ms=pb2_obj.timestamp_ms) + + def __eq__(self, other: Any) -> bool: + """Checks if this object is equal to the given object. + + Args: + other: The object to be compared with. + + Returns: + True if the objects are equal. + """ + if not isinstance(other, ClassificationEntry): + return False + + return self.to_pb2().__eq__(other.to_pb2()) + + +@dataclasses.dataclass +class Classifications: + """Represents the classifications for a given classifier head. + + Attributes: + entries: A list of `ClassificationEntry` objects. + head_index: The index of the classifier head these categories refer to. This + is useful for multi-head models. + head_name: The name of the classifier head, which is the corresponding + tensor metadata name. + """ + + entries: List[ClassificationEntry] + head_index: int + head_name: str + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _ClassificationsProto: + """Generates a Classifications protobuf object.""" + return _ClassificationsProto( + entries=[entry.to_pb2() for entry in self.entries], + head_index=self.head_index, + head_name=self.head_name) + + @classmethod + @doc_controls.do_not_generate_docs + def create_from_pb2(cls, pb2_obj: _ClassificationsProto) -> 'Classifications': + """Creates a `Classifications` object from the given protobuf object.""" + return Classifications( + entries=[ + ClassificationEntry.create_from_pb2(entry) + for entry in pb2_obj.entries + ], + head_index=pb2_obj.head_index, + head_name=pb2_obj.head_name) + + def __eq__(self, other: Any) -> bool: + """Checks if this object is equal to the given object. + + Args: + other: The object to be compared with. + + Returns: + True if the objects are equal. + """ + if not isinstance(other, Classifications): + return False + + return self.to_pb2().__eq__(other.to_pb2()) + + +@dataclasses.dataclass +class ClassificationResult: + """Contains one set of results per classifier head. + + Attributes: + classifications: A list of `Classifications` objects. + """ + + classifications: List[Classifications] + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _ClassificationResultProto: + """Generates a ClassificationResult protobuf object.""" + return _ClassificationResultProto(classifications=[ + classification.to_pb2() for classification in self.classifications + ]) + + @classmethod + @doc_controls.do_not_generate_docs + def create_from_pb2( + cls, pb2_obj: _ClassificationResultProto) -> 'ClassificationResult': + """Creates a `ClassificationResult` object from the given protobuf object. + """ + return ClassificationResult(classifications=[ + Classifications.create_from_pb2(classification) + for classification in pb2_obj.classifications + ]) + + def __eq__(self, other: Any) -> bool: + """Checks if this object is equal to the given object. + + Args: + other: The object to be compared with. + + Returns: + True if the objects are equal. + """ + if not isinstance(other, ClassificationResult): + return False + + return self.to_pb2().__eq__(other.to_pb2()) diff --git a/mediapipe/tasks/python/components/containers/landmark.py b/mediapipe/tasks/python/components/containers/landmark.py new file mode 100644 index 00000000..dee2a16a --- /dev/null +++ b/mediapipe/tasks/python/components/containers/landmark.py @@ -0,0 +1,122 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Landmark data class.""" + +import dataclasses +from typing import Optional + +from mediapipe.framework.formats import landmark_pb2 +from mediapipe.tasks.python.core.optional_dependencies import doc_controls + +_LandmarkProto = landmark_pb2.Landmark +_NormalizedLandmarkProto = landmark_pb2.NormalizedLandmark + + +@dataclasses.dataclass +class Landmark: + """A landmark that can have 1 to 3 dimensions. + + Use x for 1D points, (x, y) for 2D points and (x, y, z) for 3D points. + + Attributes: + x: The x coordinate. + y: The y coordinate. + z: The z coordinate. + visibility: Landmark visibility. Should stay unset if not supported. Float + score of whether landmark is visible or occluded by other objects. + Landmark considered as invisible also if it is not present on the screen + (out of scene bounds). Depending on the model, visibility value is either + a sigmoid or an argument of sigmoid. + presence: Landmark presence. Should stay unset if not supported. Float score + of whether landmark is present on the scene (located within scene bounds). + Depending on the model, presence value is either a result of sigmoid or an + argument of sigmoid function to get landmark presence probability. + """ + + x: Optional[float] = None + y: Optional[float] = None + z: Optional[float] = None + visibility: Optional[float] = None + presence: Optional[float] = None + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _LandmarkProto: + """Generates a Landmark protobuf object.""" + return _LandmarkProto( + x=self.x, + y=self.y, + z=self.z, + visibility=self.visibility, + presence=self.presence) + + @classmethod + @doc_controls.do_not_generate_docs + def create_from_pb2(cls, pb2_obj: _LandmarkProto) -> 'Landmark': + """Creates a `Landmark` object from the given protobuf object.""" + return Landmark( + x=pb2_obj.x, + y=pb2_obj.y, + z=pb2_obj.z, + visibility=pb2_obj.visibility, + presence=pb2_obj.presence) + + +@dataclasses.dataclass +class NormalizedLandmark: + """A normalized version of above Landmark proto. + + All coordinates should be within [0, 1]. + + Attributes: + x: The normalized x coordinate. + y: The normalized y coordinate. + z: The normalized z coordinate. + visibility: Landmark visibility. Should stay unset if not supported. Float + score of whether landmark is visible or occluded by other objects. + Landmark considered as invisible also if it is not present on the screen + (out of scene bounds). Depending on the model, visibility value is either + a sigmoid or an argument of sigmoid. + presence: Landmark presence. Should stay unset if not supported. Float score + of whether landmark is present on the scene (located within scene bounds). + Depending on the model, presence value is either a result of sigmoid or an + argument of sigmoid function to get landmark presence probability. + """ + + x: Optional[float] = None + y: Optional[float] = None + z: Optional[float] = None + visibility: Optional[float] = None + presence: Optional[float] = None + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _NormalizedLandmarkProto: + """Generates a NormalizedLandmark protobuf object.""" + return _NormalizedLandmarkProto( + x=self.x, + y=self.y, + z=self.z, + visibility=self.visibility, + presence=self.presence) + + @classmethod + @doc_controls.do_not_generate_docs + def create_from_pb2( + cls, pb2_obj: _NormalizedLandmarkProto) -> 'NormalizedLandmark': + """Creates a `NormalizedLandmark` object from the given protobuf object.""" + return NormalizedLandmark( + x=pb2_obj.x, + y=pb2_obj.y, + z=pb2_obj.z, + visibility=pb2_obj.visibility, + presence=pb2_obj.presence) diff --git a/mediapipe/tasks/python/components/containers/landmark_detection_result.py b/mediapipe/tasks/python/components/containers/landmark_detection_result.py new file mode 100644 index 00000000..e058f84a --- /dev/null +++ b/mediapipe/tasks/python/components/containers/landmark_detection_result.py @@ -0,0 +1,96 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Landmarks Detection Result data class.""" + +import dataclasses +from typing import Optional, List + +from mediapipe.framework.formats import classification_pb2 +from mediapipe.framework.formats import landmark_pb2 +from mediapipe.tasks.cc.components.containers.proto import landmarks_detection_result_pb2 +from mediapipe.tasks.python.components.containers import category as category_module +from mediapipe.tasks.python.components.containers import landmark as landmark_module +from mediapipe.tasks.python.components.containers import rect as rect_module +from mediapipe.tasks.python.core.optional_dependencies import doc_controls + +_LandmarksDetectionResultProto = landmarks_detection_result_pb2.LandmarksDetectionResult +_ClassificationProto = classification_pb2.Classification +_ClassificationListProto = classification_pb2.ClassificationList +_LandmarkListProto = landmark_pb2.LandmarkList +_NormalizedLandmarkListProto = landmark_pb2.NormalizedLandmarkList +_NormalizedRect = rect_module.NormalizedRect +_Category = category_module.Category +_NormalizedLandmark = landmark_module.NormalizedLandmark +_Landmark = landmark_module.Landmark + + +@dataclasses.dataclass +class LandmarksDetectionResult: + """Represents the landmarks detection result. + + Attributes: landmarks : A list of `NormalizedLandmark` objects. categories : A + list of `Category` objects. world_landmarks : A list of `Landmark` objects. + rect : A `NormalizedRect` object. + """ + + landmarks: Optional[List[_NormalizedLandmark]] + categories: Optional[List[_Category]] + world_landmarks: Optional[List[_Landmark]] + rect: _NormalizedRect + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _LandmarksDetectionResultProto: + """Generates a LandmarksDetectionResult protobuf object.""" + + classifications = _ClassificationListProto() + for category in self.categories: + classifications.classification.append( + _ClassificationProto( + index=category.index, + score=category.score, + label=category.category_name, + display_name=category.display_name)) + + return _LandmarksDetectionResultProto( + landmarks=_NormalizedLandmarkListProto(self.landmarks), + classifications=classifications, + world_landmarks=_LandmarkListProto(self.world_landmarks), + rect=self.rect.to_pb2()) + + @classmethod + @doc_controls.do_not_generate_docs + def create_from_pb2( + cls, + pb2_obj: _LandmarksDetectionResultProto) -> 'LandmarksDetectionResult': + """Creates a `LandmarksDetectionResult` object from the given protobuf object. + """ + categories = [] + for classification in pb2_obj.classifications.classification: + categories.append( + category_module.Category( + score=classification.score, + index=classification.index, + category_name=classification.label, + display_name=classification.display_name)) + return LandmarksDetectionResult( + landmarks=[ + _NormalizedLandmark.create_from_pb2(landmark) + for landmark in pb2_obj.landmarks.landmark + ], + categories=categories, + world_landmarks=[ + _Landmark.create_from_pb2(landmark) + for landmark in pb2_obj.world_landmarks.landmark + ], + rect=_NormalizedRect.create_from_pb2(pb2_obj.rect)) diff --git a/mediapipe/tasks/python/components/processors/BUILD b/mediapipe/tasks/python/components/processors/BUILD new file mode 100644 index 00000000..f87a579b --- /dev/null +++ b/mediapipe/tasks/python/components/processors/BUILD @@ -0,0 +1,30 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Placeholder for internal Python strict library compatibility macro. + +# Placeholder for internal Python strict library and test compatibility macro. + +package(default_visibility = ["//mediapipe/tasks:internal"]) + +licenses(["notice"]) + +py_library( + name = "classifier_options", + srcs = ["classifier_options.py"], + deps = [ + "//mediapipe/tasks/cc/components/processors/proto:classifier_options_py_pb2", + "//mediapipe/tasks/python/core:optional_dependencies", + ], +) diff --git a/mediapipe/tasks/python/components/processors/__init__.py b/mediapipe/tasks/python/components/processors/__init__.py new file mode 100644 index 00000000..65c1214a --- /dev/null +++ b/mediapipe/tasks/python/components/processors/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/mediapipe/tasks/python/components/processors/classifier_options.py b/mediapipe/tasks/python/components/processors/classifier_options.py new file mode 100644 index 00000000..2e77f93b --- /dev/null +++ b/mediapipe/tasks/python/components/processors/classifier_options.py @@ -0,0 +1,86 @@ +# Copyright 2022 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Classifier options data class.""" + +import dataclasses +from typing import Any, List, Optional + +from mediapipe.tasks.cc.components.processors.proto import classifier_options_pb2 +from mediapipe.tasks.python.core.optional_dependencies import doc_controls + +_ClassifierOptionsProto = classifier_options_pb2.ClassifierOptions + + +@dataclasses.dataclass +class ClassifierOptions: + """Options for classification processor. + + Attributes: + display_names_locale: The locale to use for display names specified through + the TFLite Model Metadata. + max_results: The maximum number of top-scored classification results to + return. + score_threshold: Overrides the ones provided in the model metadata. Results + below this value are rejected. + category_allowlist: Allowlist of category names. If non-empty, detection + results whose category name is not in this set will be filtered out. + Duplicate or unknown category names are ignored. Mutually exclusive with + `category_denylist`. + category_denylist: Denylist of category names. If non-empty, detection + results whose category name is in this set will be filtered out. Duplicate + or unknown category names are ignored. Mutually exclusive with + `category_allowlist`. + """ + + display_names_locale: Optional[str] = None + max_results: Optional[int] = None + score_threshold: Optional[float] = None + category_allowlist: Optional[List[str]] = None + category_denylist: Optional[List[str]] = None + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _ClassifierOptionsProto: + """Generates a ClassifierOptions protobuf object.""" + return _ClassifierOptionsProto( + score_threshold=self.score_threshold, + category_allowlist=self.category_allowlist, + category_denylist=self.category_denylist, + display_names_locale=self.display_names_locale, + max_results=self.max_results) + + @classmethod + @doc_controls.do_not_generate_docs + def create_from_pb2(cls, + pb2_obj: _ClassifierOptionsProto) -> 'ClassifierOptions': + """Creates a `ClassifierOptions` object from the given protobuf object.""" + return ClassifierOptions( + score_threshold=pb2_obj.score_threshold, + category_allowlist=[str(name) for name in pb2_obj.category_allowlist], + category_denylist=[str(name) for name in pb2_obj.category_denylist], + display_names_locale=pb2_obj.display_names_locale, + max_results=pb2_obj.max_results) + + def __eq__(self, other: Any) -> bool: + """Checks if this object is equal to the given object. + + Args: + other: The object to be compared with. + + Returns: + True if the objects are equal. + """ + if not isinstance(other, ClassifierOptions): + return False + + return self.to_pb2().__eq__(other.to_pb2()) diff --git a/mediapipe/tasks/python/core/BUILD b/mediapipe/tasks/python/core/BUILD index d5cdeecd..76e2f4f4 100644 --- a/mediapipe/tasks/python/core/BUILD +++ b/mediapipe/tasks/python/core/BUILD @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Placeholder for internal Python strict library compatibility macro. +# Placeholder for internal Python strict library and test compatibility macro. package(default_visibility = ["//mediapipe/tasks:internal"]) diff --git a/mediapipe/tasks/python/metadata/metadata_writers/BUILD b/mediapipe/tasks/python/metadata/metadata_writers/BUILD index 3e44218a..d2b55d47 100644 --- a/mediapipe/tasks/python/metadata/metadata_writers/BUILD +++ b/mediapipe/tasks/python/metadata/metadata_writers/BUILD @@ -1,4 +1,4 @@ -# Placeholder for internal Python strict library compatibility macro. +# Placeholder for internal Python strict library and test compatibility macro. package( default_visibility = [ @@ -37,3 +37,9 @@ py_library( srcs = ["writer_utils.py"], deps = ["//mediapipe/tasks/metadata:schema_py"], ) + +py_library( + name = "image_classifier", + srcs = ["image_classifier.py"], + deps = [":metadata_writer"], +) diff --git a/mediapipe/tasks/python/metadata/metadata_writers/image_classifier.py b/mediapipe/tasks/python/metadata/metadata_writers/image_classifier.py new file mode 100644 index 00000000..c516a342 --- /dev/null +++ b/mediapipe/tasks/python/metadata/metadata_writers/image_classifier.py @@ -0,0 +1,71 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Writes metadata and label file to the image classifier models.""" + +from typing import List, Optional + +from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer + +_MODEL_NAME = "ImageClassifier" +_MODEL_DESCRIPTION = ("Identify the most prominent object in the image from a " + "known set of categories.") + + +class MetadataWriter(metadata_writer.MetadataWriterBase): + """MetadataWriter to write the metadata for image classifier.""" + + @classmethod + def create( + cls, + model_buffer: bytearray, + input_norm_mean: List[float], + input_norm_std: List[float], + labels: metadata_writer.Labels, + score_calibration: Optional[metadata_writer.ScoreCalibration] = None + ) -> "MetadataWriter": + """Creates MetadataWriter to write the metadata for image classifier. + + The parameters required in this method are mandatory when using MediaPipe + Tasks. + + Note that only the output TFLite is used for deployment. The output JSON + content is used to interpret the metadata content. + + Args: + model_buffer: A valid flatbuffer loaded from the TFLite model file. + input_norm_mean: the mean value used in the input tensor normalization + [1]. + input_norm_std: the std value used in the input tensor normalizarion [1]. + labels: an instance of Labels helper class used in the output + classification tensor [2]. + score_calibration: A container of the score calibration operation [3] in + the classification tensor. Optional if the model does not use score + calibration. + + [1]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389 + [2]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99 + [3]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L456 + + Returns: + An MetadataWrite object. + """ + writer = metadata_writer.MetadataWriter(model_buffer) + writer.add_genernal_info(_MODEL_NAME, _MODEL_DESCRIPTION) + writer.add_image_input(input_norm_mean, input_norm_std) + writer.add_classification_output(labels, score_calibration) + return cls(writer) diff --git a/mediapipe/tasks/python/metadata/metadata_writers/metadata_info.py b/mediapipe/tasks/python/metadata/metadata_writers/metadata_info.py index 07938d86..76b572e8 100644 --- a/mediapipe/tasks/python/metadata/metadata_writers/metadata_info.py +++ b/mediapipe/tasks/python/metadata/metadata_writers/metadata_info.py @@ -197,6 +197,37 @@ class ScoreCalibrationMd: self._FILE_TYPE) +class ScoreThresholdingMd: + """A container for score thresholding [1] metadata information. + + [1]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L468 + """ + + def __init__(self, global_score_threshold: float) -> None: + """Creates a ScoreThresholdingMd object. + + Args: + global_score_threshold: The recommended global threshold below which + results are considered low-confidence and should be filtered out. + """ + self._global_score_threshold = global_score_threshold + + def create_metadata(self) -> _metadata_fb.ProcessUnitT: + """Creates the score thresholding metadata based on the information. + + Returns: + A Flatbuffers Python object of the score thresholding metadata. + """ + score_thresholding = _metadata_fb.ProcessUnitT() + score_thresholding.optionsType = ( + _metadata_fb.ProcessUnitOptions.ScoreThresholdingOptions) + options = _metadata_fb.ScoreThresholdingOptionsT() + options.globalScoreThreshold = self._global_score_threshold + score_thresholding.options = options + return score_thresholding + + class TensorMd: """A container for common tensor metadata information. @@ -374,23 +405,29 @@ class ClassificationTensorMd(TensorMd): tensor. score_calibration_md: information of the score calibration operation [2] in the classification tensor. + score_thresholding_md: information of the score thresholding [3] in the + classification tensor. [1]: https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99 [2]: https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L456 + [3]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L468 """ # Min and max float values for classification results. _MIN_FLOAT = 0.0 _MAX_FLOAT = 1.0 - def __init__(self, - name: Optional[str] = None, - description: Optional[str] = None, - label_files: Optional[List[LabelFileMd]] = None, - tensor_type: Optional[int] = None, - score_calibration_md: Optional[ScoreCalibrationMd] = None, - tensor_name: Optional[str] = None) -> None: + def __init__( + self, + name: Optional[str] = None, + description: Optional[str] = None, + label_files: Optional[List[LabelFileMd]] = None, + tensor_type: Optional[int] = None, + score_calibration_md: Optional[ScoreCalibrationMd] = None, + tensor_name: Optional[str] = None, + score_thresholding_md: Optional[ScoreThresholdingMd] = None) -> None: """Initializes the instance of ClassificationTensorMd. Args: @@ -404,6 +441,8 @@ class ClassificationTensorMd(TensorMd): tensor_name: name of the corresponding tensor [3] in the TFLite model. It is used to locate the corresponding classification tensor and decide the order of the tensor metadata [4] when populating model metadata. + score_thresholding_md: information of the score thresholding [5] in the + classification tensor. [1]: https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99 [2]: @@ -412,8 +451,11 @@ class ClassificationTensorMd(TensorMd): https://github.com/tensorflow/tensorflow/blob/cb67fef35567298b40ac166b0581cd8ad68e5a3a/tensorflow/lite/schema/schema.fbs#L1129-L1136 [4]: https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L623-L640 + [5]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L468 """ self.score_calibration_md = score_calibration_md + self.score_thresholding_md = score_thresholding_md if tensor_type is _schema_fb.TensorType.UINT8: min_values = [_MIN_UINT8] @@ -443,4 +485,12 @@ class ClassificationTensorMd(TensorMd): tensor_metadata.processUnits = [ self.score_calibration_md.create_metadata() ] + if self.score_thresholding_md: + if tensor_metadata.processUnits: + tensor_metadata.processUnits.append( + self.score_thresholding_md.create_metadata()) + else: + tensor_metadata.processUnits = [ + self.score_thresholding_md.create_metadata() + ] return tensor_metadata diff --git a/mediapipe/tasks/python/metadata/metadata_writers/metadata_writer.py b/mediapipe/tasks/python/metadata/metadata_writers/metadata_writer.py index e69efd01..3a9f9123 100644 --- a/mediapipe/tasks/python/metadata/metadata_writers/metadata_writer.py +++ b/mediapipe/tasks/python/metadata/metadata_writers/metadata_writer.py @@ -15,19 +15,22 @@ """Generic metadata writer.""" import collections +import csv import dataclasses import os import tempfile from typing import List, Optional, Tuple import flatbuffers -from mediapipe.tasks.metadata import metadata_schema_py_generated as _metadata_fb -from mediapipe.tasks.python.metadata import metadata as _metadata +from mediapipe.tasks.metadata import metadata_schema_py_generated as metadata_fb +from mediapipe.tasks.python.metadata import metadata from mediapipe.tasks.python.metadata.metadata_writers import metadata_info from mediapipe.tasks.python.metadata.metadata_writers import writer_utils _INPUT_IMAGE_NAME = 'image' _INPUT_IMAGE_DESCRIPTION = 'Input image to be processed.' +_OUTPUT_CLASSIFICATION_NAME = 'score' +_OUTPUT_CLASSIFICATION_DESCRIPTION = 'Score of the labels respectively.' @dataclasses.dataclass @@ -67,6 +70,18 @@ class LabelItem: locale: Optional[str] = None +@dataclasses.dataclass +class ScoreThresholding: + """Parameters to performs thresholding on output tensor values [1]. + + Attributes: + global_score_threshold: The recommended global threshold below which results + are considered low-confidence and should be filtered out. [1]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L468 + """ + global_score_threshold: float + + class Labels(object): """Simple container holding classification labels of a particular tensor. @@ -140,26 +155,85 @@ class Labels(object): class ScoreCalibration: """Simple container holding score calibration related parameters.""" - # A shortcut to avoid client side code importing _metadata_fb - transformation_types = _metadata_fb.ScoreTransformationType + # A shortcut to avoid client side code importing metadata_fb + transformation_types = metadata_fb.ScoreTransformationType def __init__(self, - transformation_type: _metadata_fb.ScoreTransformationType, - parameters: List[CalibrationParameter], + transformation_type: metadata_fb.ScoreTransformationType, + parameters: List[Optional[CalibrationParameter]], default_score: int = 0): self.transformation_type = transformation_type self.parameters = parameters self.default_score = default_score + @classmethod + def create_from_file(cls, + transformation_type: metadata_fb.ScoreTransformationType, + file_path: str, + default_score: int = 0) -> 'ScoreCalibration': + """Creates ScoreCalibration from the file. + + Args: + transformation_type: type of the function used for transforming the + uncalibrated score before applying score calibration. + file_path: file_path of the score calibration file [1]. Contains + sigmoid-based score calibration parameters, formatted as CSV. Lines + contain for each index of an output tensor the scale, slope, offset and + (optional) min_score parameters to be used for sigmoid fitting (in this + order and in `strtof`-compatible [2] format). Scale should be a + non-negative value. A line may be left empty to default calibrated + scores for this index to default_score. In summary, each line should + thus contain 0, 3 or 4 comma-separated values. + default_score: the default calibrated score to apply if the uncalibrated + score is below min_score or if no parameters were specified for a given + index. + [1]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L133 + [2]: + https://en.cppreference.com/w/c/string/byte/strtof + + Returns: + A ScoreCalibration object. + Raises: + ValueError: if the score_calibration file is malformed. + """ + with open(file_path, 'r') as calibration_file: + csv_reader = csv.reader(calibration_file, delimiter=',') + parameters = [] + for row in csv_reader: + if not row: + parameters.append(None) + continue + + if len(row) != 3 and len(row) != 4: + raise ValueError( + f'Expected empty lines or 3 or 4 parameters per line in score' + f' calibration file, but got {len(row)}.') + + if float(row[0]) < 0: + raise ValueError( + f'Expected scale to be a non-negative value, but got ' + f'{float(row[0])}.') + + parameters.append( + CalibrationParameter( + scale=float(row[0]), + slope=float(row[1]), + offset=float(row[2]), + min_score=None if len(row) == 3 else float(row[3]))) + + return cls(transformation_type, parameters, default_score) + def _fill_default_tensor_names( - tensor_metadata: List[_metadata_fb.TensorMetadataT], + tensor_metadata_list: List[metadata_fb.TensorMetadataT], tensor_names_from_model: List[str]): """Fills the default tensor names.""" # If tensor name in metadata is empty, default to the tensor name saved in # the model. - for metadata, name in zip(tensor_metadata, tensor_names_from_model): - metadata.name = metadata.name or name + for tensor_metadata, name in zip(tensor_metadata_list, + tensor_names_from_model): + tensor_metadata.name = tensor_metadata.name or name def _pair_tensor_metadata( @@ -212,7 +286,7 @@ def _create_metadata_buffer( input_metadata = [m.create_metadata() for m in input_md] else: num_input_tensors = writer_utils.get_subgraph(model_buffer).InputsLength() - input_metadata = [_metadata_fb.TensorMetadataT()] * num_input_tensors + input_metadata = [metadata_fb.TensorMetadataT()] * num_input_tensors _fill_default_tensor_names(input_metadata, writer_utils.get_input_tensor_names(model_buffer)) @@ -224,12 +298,12 @@ def _create_metadata_buffer( output_metadata = [m.create_metadata() for m in output_md] else: num_output_tensors = writer_utils.get_subgraph(model_buffer).OutputsLength() - output_metadata = [_metadata_fb.TensorMetadataT()] * num_output_tensors + output_metadata = [metadata_fb.TensorMetadataT()] * num_output_tensors _fill_default_tensor_names(output_metadata, writer_utils.get_output_tensor_names(model_buffer)) # Create the subgraph metadata. - subgraph_metadata = _metadata_fb.SubGraphMetadataT() + subgraph_metadata = metadata_fb.SubGraphMetadataT() subgraph_metadata.inputTensorMetadata = input_metadata subgraph_metadata.outputTensorMetadata = output_metadata @@ -243,7 +317,7 @@ def _create_metadata_buffer( b = flatbuffers.Builder(0) b.Finish( model_metadata.Pack(b), - _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) + metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) return b.Output() @@ -291,7 +365,7 @@ class MetadataWriter(object): name=model_name, description=model_description) return self - color_space_types = _metadata_fb.ColorSpaceType + color_space_types = metadata_fb.ColorSpaceType def add_feature_input(self, name: Optional[str] = None, @@ -305,7 +379,7 @@ class MetadataWriter(object): self, norm_mean: List[float], norm_std: List[float], - color_space_type: Optional[int] = _metadata_fb.ColorSpaceType.RGB, + color_space_type: Optional[int] = metadata_fb.ColorSpaceType.RGB, name: str = _INPUT_IMAGE_NAME, description: str = _INPUT_IMAGE_DESCRIPTION) -> 'MetadataWriter': """Adds an input image metadata for the image input. @@ -341,13 +415,11 @@ class MetadataWriter(object): self._input_mds.append(input_md) return self - _OUTPUT_CLASSIFICATION_NAME = 'score' - _OUTPUT_CLASSIFICATION_DESCRIPTION = 'Score of the labels respectively' - def add_classification_output( self, labels: Optional[Labels] = None, score_calibration: Optional[ScoreCalibration] = None, + score_thresholding: Optional[ScoreThresholding] = None, name: str = _OUTPUT_CLASSIFICATION_NAME, description: str = _OUTPUT_CLASSIFICATION_DESCRIPTION ) -> 'MetadataWriter': @@ -364,6 +436,7 @@ class MetadataWriter(object): Args: labels: an instance of Labels helper class. score_calibration: an instance of ScoreCalibration helper class. + score_thresholding: an instance of ScoreThresholding. name: Metadata name of the tensor. Note that this is different from tensor name in the flatbuffer. description: human readable description of what the output is. @@ -378,6 +451,10 @@ class MetadataWriter(object): default_score=score_calibration.default_score, file_path=self._export_calibration_file('score_calibration.txt', score_calibration.parameters)) + score_thresholding_md = None + if score_thresholding: + score_thresholding_md = metadata_info.ScoreThresholdingMd( + score_thresholding.global_score_threshold) label_files = None if labels: @@ -394,6 +471,7 @@ class MetadataWriter(object): label_files=label_files, tensor_type=self._output_tensor_type(len(self._output_mds)), score_calibration_md=calibration_md, + score_thresholding_md=score_thresholding_md, ) self._output_mds.append(output_md) return self @@ -416,8 +494,7 @@ class MetadataWriter(object): A tuple of (model_with_metadata_in_bytes, metdata_json_content) """ # Populates metadata and associated files into TFLite model buffer. - populator = _metadata.MetadataPopulator.with_model_buffer( - self._model_buffer) + populator = metadata.MetadataPopulator.with_model_buffer(self._model_buffer) metadata_buffer = _create_metadata_buffer( model_buffer=self._model_buffer, general_md=self._general_md, @@ -429,7 +506,7 @@ class MetadataWriter(object): populator.populate() tflite_content = populator.get_model_buffer() - displayer = _metadata.MetadataDisplayer.with_model_buffer(tflite_content) + displayer = metadata.MetadataDisplayer.with_model_buffer(tflite_content) metadata_json_content = displayer.get_metadata_json() return tflite_content, metadata_json_content @@ -452,9 +529,7 @@ class MetadataWriter(object): """Stores calibration parameters in a csv file.""" filepath = os.path.join(self._temp_folder.name, filename) with open(filepath, 'w') as f: - for idx, item in enumerate(calibrations): - if idx != 0: - f.write('\n') + for item in calibrations: if item: if item.scale is None or item.slope is None or item.offset is None: raise ValueError('scale, slope and offset values can not be set to ' @@ -463,6 +538,29 @@ class MetadataWriter(object): f.write(f'{item.scale},{item.slope},{item.offset},{item.min_score}') else: f.write(f'{item.scale},{item.slope},{item.offset}') + f.write('\n') - self._associated_files.append(filepath) + self._associated_files.append(filepath) return filepath + + +class MetadataWriterBase: + """Base MetadataWriter class which contains the apis exposed to users. + + MetadataWriter for Tasks e.g. image classifier / object detector will inherit + this class for their own usage. + """ + + def __init__(self, writer: MetadataWriter) -> None: + self.writer = writer + + def populate(self) -> Tuple[bytearray, str]: + """Populates metadata into the TFLite file. + + Note that only the output tflite is used for deployment. The output JSON + content is used to interpret the metadata content. + + Returns: + A tuple of (model_with_metadata_in_bytes, metdata_json_content) + """ + return self.writer.populate() diff --git a/mediapipe/tasks/python/metadata/metadata_writers/writer_utils.py b/mediapipe/tasks/python/metadata/metadata_writers/writer_utils.py index 0a054812..eff5f553 100644 --- a/mediapipe/tasks/python/metadata/metadata_writers/writer_utils.py +++ b/mediapipe/tasks/python/metadata/metadata_writers/writer_utils.py @@ -14,7 +14,8 @@ # ============================================================================== """Helper methods for writing metadata into TFLite models.""" -from typing import List +from typing import Dict, List +import zipfile from mediapipe.tasks.metadata import schema_py_generated as _schema_fb @@ -83,3 +84,20 @@ def get_subgraph(model_buffer: bytearray) -> _schema_fb.SubGraph: # multiple subgraphs yet, but models with mini-benchmark may have multiple # subgraphs for acceleration evaluation purpose. return model.Subgraphs(0) + + +def create_model_asset_bundle(input_models: Dict[str, bytes], + output_path: str) -> None: + """Creates the model asset bundle. + + Args: + input_models: A dict of input models with key as the model file name and + value as the model content. + output_path: The output file path to save the model asset bundle. + """ + if not input_models or len(input_models) < 2: + raise ValueError("Needs at least two input models for model asset bundle.") + + with zipfile.ZipFile(output_path, mode="w") as zf: + for file_name, file_buffer in input_models.items(): + zf.writestr(file_name, file_buffer) diff --git a/mediapipe/tasks/python/test/BUILD b/mediapipe/tasks/python/test/BUILD index 8e5b91cf..5ad05798 100644 --- a/mediapipe/tasks/python/test/BUILD +++ b/mediapipe/tasks/python/test/BUILD @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Placeholder for internal Python strict library compatibility macro. +# Placeholder for internal Python strict library and test compatibility macro. package(default_visibility = ["//mediapipe/tasks:internal"]) @@ -27,5 +27,8 @@ py_library( "//mediapipe/model_maker/python/vision/gesture_recognizer:__pkg__", "//mediapipe/tasks:internal", ], - deps = ["//mediapipe/python:_framework_bindings"], + deps = [ + "//mediapipe/python:_framework_bindings", + "@com_google_protobuf//:protobuf_python", + ], ) diff --git a/mediapipe/tasks/python/test/metadata/metadata_test.py b/mediapipe/tasks/python/test/metadata/metadata_test.py index 00dbe526..bed9c283 100644 --- a/mediapipe/tasks/python/test/metadata/metadata_test.py +++ b/mediapipe/tasks/python/test/metadata/metadata_test.py @@ -27,6 +27,8 @@ from mediapipe.tasks.metadata import schema_py_generated as _schema_fb from mediapipe.tasks.python.metadata import metadata as _metadata from mediapipe.tasks.python.test import test_utils +_TEST_DATA_DIR = "mediapipe/tasks/testdata/metadata" + class Tokenizer(enum.Enum): BERT_TOKENIZER = 0 @@ -810,7 +812,8 @@ class MetadataDisplayerTest(MetadataTest): actual_json = _metadata.convert_to_json(actual_buffer) # Verifies the generated json file. - golden_json_file_path = test_utils.get_test_data_path("golden_json.json") + golden_json_file_path = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, "golden_json.json")) with open(golden_json_file_path, "r") as f: expected = f.read() self.assertEqual(actual_json, expected) @@ -821,7 +824,8 @@ class MetadataDisplayerTest(MetadataTest): actual = displayer.get_metadata_json() # Verifies the generated json file. - golden_json_file_path = test_utils.get_test_data_path("golden_json.json") + golden_json_file_path = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, "golden_json.json")) expected = _read_file(golden_json_file_path, "r") self.assertEqual(actual, expected) @@ -848,7 +852,8 @@ class MetadataUtilTest(MetadataTest): metadata_json = _metadata.convert_to_json(metadata_buf) # Verifies the generated json file. - golden_json_file_path = test_utils.get_test_data_path("golden_json.json") + golden_json_file_path = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, "golden_json.json")) expected = _read_file(golden_json_file_path, "r") self.assertEqual(metadata_json, expected) diff --git a/mediapipe/tasks/python/test/metadata/metadata_writers/BUILD b/mediapipe/tasks/python/test/metadata/metadata_writers/BUILD index 8779d2fb..a7bfd297 100644 --- a/mediapipe/tasks/python/test/metadata/metadata_writers/BUILD +++ b/mediapipe/tasks/python/test/metadata/metadata_writers/BUILD @@ -28,9 +28,28 @@ py_test( py_test( name = "metadata_writer_test", srcs = ["metadata_writer_test.py"], - data = ["//mediapipe/tasks/testdata/metadata:model_files"], + data = [ + "//mediapipe/tasks/testdata/metadata:data_files", + "//mediapipe/tasks/testdata/metadata:model_files", + ], deps = [ "//mediapipe/tasks/python/metadata/metadata_writers:metadata_writer", "//mediapipe/tasks/python/test:test_utils", ], ) + +py_test( + name = "image_classifier_test", + srcs = ["image_classifier_test.py"], + data = [ + "//mediapipe/tasks/testdata/metadata:data_files", + "//mediapipe/tasks/testdata/metadata:model_files", + ], + deps = [ + "//mediapipe/tasks/metadata:metadata_schema_py", + "//mediapipe/tasks/python/metadata", + "//mediapipe/tasks/python/metadata/metadata_writers:image_classifier", + "//mediapipe/tasks/python/metadata/metadata_writers:metadata_writer", + "//mediapipe/tasks/python/test:test_utils", + ], +) diff --git a/mediapipe/tasks/python/test/metadata/metadata_writers/image_classifier_test.py b/mediapipe/tasks/python/test/metadata/metadata_writers/image_classifier_test.py new file mode 100644 index 00000000..51f248d1 --- /dev/null +++ b/mediapipe/tasks/python/test/metadata/metadata_writers/image_classifier_test.py @@ -0,0 +1,85 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for metadata_writer.image_classifier.""" + +import os + +from absl.testing import absltest +from absl.testing import parameterized + +from mediapipe.tasks.metadata import metadata_schema_py_generated as metadata_fb +from mediapipe.tasks.python.metadata import metadata +from mediapipe.tasks.python.metadata.metadata_writers import image_classifier +from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer +from mediapipe.tasks.python.test import test_utils + +_TEST_DATA_DIR = "mediapipe/tasks/testdata/metadata" +_FLOAT_MODEL = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, + "mobilenet_v2_1.0_224_without_metadata.tflite")) +_QUANT_MODEL = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, + "mobilenet_v2_1.0_224_quant_without_metadata.tflite")) +_LABEL_FILE = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, "labels.txt")) +_SCORE_CALIBRATION_FILE = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, "score_calibration.txt")) +_SCORE_CALIBRATION_FILENAME = "score_calibration.txt" +_DEFAULT_SCORE_CALIBRATION_VALUE = 0.2 +_NORM_MEAN = 127.5 +_NORM_STD = 127.5 +_FLOAT_JSON = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, "mobilenet_v2_1.0_224.json")) +_QUANT_JSON = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, "mobilenet_v2_1.0_224_quant.json")) + + +class ImageClassifierTest(parameterized.TestCase): + + @parameterized.named_parameters( + { + "testcase_name": "test_float_model", + "model_file": _FLOAT_MODEL, + "golden_json": _FLOAT_JSON + }, { + "testcase_name": "test_quant_model", + "model_file": _QUANT_MODEL, + "golden_json": _QUANT_JSON + }) + def test_write_metadata(self, model_file: str, golden_json: str): + with open(model_file, "rb") as f: + model_buffer = f.read() + writer = image_classifier.MetadataWriter.create( + model_buffer, [_NORM_MEAN], [_NORM_STD], + labels=metadata_writer.Labels().add_from_file(_LABEL_FILE), + score_calibration=metadata_writer.ScoreCalibration.create_from_file( + metadata_fb.ScoreTransformationType.LOG, _SCORE_CALIBRATION_FILE, + _DEFAULT_SCORE_CALIBRATION_VALUE)) + tflite_content, metadata_json = writer.populate() + + with open(golden_json, "r") as f: + expected_json = f.read() + self.assertEqual(metadata_json, expected_json) + + displayer = metadata.MetadataDisplayer.with_model_buffer(tflite_content) + file_buffer = displayer.get_associated_file_buffer( + _SCORE_CALIBRATION_FILENAME) + with open(_SCORE_CALIBRATION_FILE, "rb") as f: + expected_file_buffer = f.read() + self.assertEqual(file_buffer, expected_file_buffer) + + +if __name__ == "__main__": + absltest.main() diff --git a/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_info_test.py b/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_info_test.py index 75602c83..0e1d1c36 100644 --- a/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_info_test.py +++ b/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_info_test.py @@ -14,6 +14,7 @@ # ============================================================================== """Tests for metadata info classes.""" +import os import tempfile from absl.testing import absltest @@ -26,13 +27,15 @@ from mediapipe.tasks.python.metadata import metadata as _metadata from mediapipe.tasks.python.metadata.metadata_writers import metadata_info from mediapipe.tasks.python.test import test_utils -_SCORE_CALIBRATION_FILE = test_utils.get_test_data_path("score_calibration.txt") +_TEST_DATA_DIR = "mediapipe/tasks/testdata/metadata" +_SCORE_CALIBRATION_FILE = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, "score_calibration.txt")) class GeneralMdTest(absltest.TestCase): _EXPECTED_GENERAL_META_JSON = test_utils.get_test_data_path( - "general_meta.json") + os.path.join(_TEST_DATA_DIR, "general_meta.json")) def test_create_metadata_should_succeed(self): general_md = metadata_info.GeneralMd( @@ -59,7 +62,7 @@ class GeneralMdTest(absltest.TestCase): class AssociatedFileMdTest(absltest.TestCase): _EXPECTED_META_JSON = test_utils.get_test_data_path( - "associated_file_meta.json") + os.path.join(_TEST_DATA_DIR, "associated_file_meta.json")) def test_create_metadata_should_succeed(self): file_md = metadata_info.AssociatedFileMd( @@ -92,11 +95,11 @@ class TensorMdTest(parameterized.TestCase): _LABEL_FILE_EN = "labels.txt" _LABEL_FILE_CN = "labels_cn.txt" # Locale label file in Chinese. _EXPECTED_FEATURE_TENSOR_JSON = test_utils.get_test_data_path( - "feature_tensor_meta.json") + os.path.join(_TEST_DATA_DIR, "feature_tensor_meta.json")) _EXPECTED_IMAGE_TENSOR_JSON = test_utils.get_test_data_path( - "image_tensor_meta.json") + os.path.join(_TEST_DATA_DIR, "image_tensor_meta.json")) _EXPECTED_BOUNDING_BOX_TENSOR_JSON = test_utils.get_test_data_path( - "bounding_box_tensor_meta.json") + os.path.join(_TEST_DATA_DIR, "bounding_box_tensor_meta.json")) @parameterized.named_parameters( { @@ -142,11 +145,11 @@ class InputImageTensorMdTest(parameterized.TestCase): _NORM_STD = (127.5, 127.5, 127.5) _COLOR_SPACE_TYPE = _metadata_fb.ColorSpaceType.RGB _EXPECTED_FLOAT_TENSOR_JSON = test_utils.get_test_data_path( - "input_image_tensor_float_meta.json") + os.path.join(_TEST_DATA_DIR, "input_image_tensor_float_meta.json")) _EXPECTED_UINT8_TENSOR_JSON = test_utils.get_test_data_path( - "input_image_tensor_uint8_meta.json") + os.path.join(_TEST_DATA_DIR, "input_image_tensor_uint8_meta.json")) _EXPECTED_UNSUPPORTED_TENSOR_JSON = test_utils.get_test_data_path( - "input_image_tensor_unsupported_meta.json") + os.path.join(_TEST_DATA_DIR, "input_image_tensor_unsupported_meta.json")) @parameterized.named_parameters( { @@ -196,11 +199,12 @@ class ClassificationTensorMdTest(parameterized.TestCase): _LABEL_FILE_CN = "labels_cn.txt" # Locale label file in Chinese. _CALIBRATION_DEFAULT_SCORE = 0.2 _EXPECTED_FLOAT_TENSOR_JSON = test_utils.get_test_data_path( - "classification_tensor_float_meta.json") + os.path.join(_TEST_DATA_DIR, "classification_tensor_float_meta.json")) _EXPECTED_UINT8_TENSOR_JSON = test_utils.get_test_data_path( - "classification_tensor_uint8_meta.json") + os.path.join(_TEST_DATA_DIR, "classification_tensor_uint8_meta.json")) _EXPECTED_UNSUPPORTED_TENSOR_JSON = test_utils.get_test_data_path( - "classification_tensor_unsupported_meta.json") + os.path.join(_TEST_DATA_DIR, + "classification_tensor_unsupported_meta.json")) @parameterized.named_parameters( { @@ -243,9 +247,9 @@ class ClassificationTensorMdTest(parameterized.TestCase): class ScoreCalibrationMdTest(absltest.TestCase): _DEFAULT_VALUE = 0.2 _EXPECTED_TENSOR_JSON = test_utils.get_test_data_path( - "score_calibration_tensor_meta.json") + os.path.join(_TEST_DATA_DIR, "score_calibration_tensor_meta.json")) _EXPECTED_MODEL_META_JSON = test_utils.get_test_data_path( - "score_calibration_file_meta.json") + os.path.join(_TEST_DATA_DIR, "score_calibration_file_meta.json")) def test_create_metadata_should_succeed(self): score_calibration_md = metadata_info.ScoreCalibrationMd( @@ -307,6 +311,25 @@ class ScoreCalibrationMdTest(absltest.TestCase): malformed_calibration_file) +class ScoreThresholdingMdTest(absltest.TestCase): + _DEFAULT_GLOBAL_THRESHOLD = 0.5 + _EXPECTED_TENSOR_JSON = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, "score_thresholding_meta.json")) + + def test_create_metadata_should_succeed(self): + score_thresholding_md = metadata_info.ScoreThresholdingMd( + global_score_threshold=self._DEFAULT_GLOBAL_THRESHOLD) + + score_thresholding_metadata = score_thresholding_md.create_metadata() + + metadata_json = _metadata.convert_to_json( + _create_dummy_model_metadata_with_process_uint( + score_thresholding_metadata)) + with open(self._EXPECTED_TENSOR_JSON, "r") as f: + expected_json = f.read() + self.assertEqual(metadata_json, expected_json) + + def _create_dummy_model_metadata_with_tensor( tensor_metadata: _metadata_fb.TensorMetadataT) -> bytes: # Create a dummy model using the tensor metadata. diff --git a/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_writer_test.py b/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_writer_test.py index c39b4a55..8cde318e 100644 --- a/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_writer_test.py +++ b/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_writer_test.py @@ -13,13 +13,20 @@ # limitations under the License. # ============================================================================== """Tests for metadata writer classes.""" +import os +import tempfile + from absl.testing import absltest from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer from mediapipe.tasks.python.test import test_utils +_TEST_DATA_DIR = 'mediapipe/tasks/testdata/metadata' + _IMAGE_CLASSIFIER_MODEL = test_utils.get_test_data_path( - 'mobilenet_v1_0.25_224_1_default_1.tflite') + os.path.join(_TEST_DATA_DIR, 'mobilenet_v1_0.25_224_1_default_1.tflite')) +_SCORE_CALIBRATION_FILE = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, 'score_calibration.txt')) class LabelsTest(absltest.TestCase): @@ -49,6 +56,53 @@ class LabelsTest(absltest.TestCase): ]) +class ScoreCalibrationTest(absltest.TestCase): + + def test_create_from_file_successful(self): + score_calibration = metadata_writer.ScoreCalibration.create_from_file( + metadata_writer.ScoreCalibration.transformation_types.LOG, + _SCORE_CALIBRATION_FILE) + self.assertLen(score_calibration.parameters, 511) + self.assertIsNone(score_calibration.parameters[0]) + self.assertEqual( + score_calibration.parameters[1], + metadata_writer.CalibrationParameter( + scale=0.9876328110694885, + slope=0.36622241139411926, + offset=0.5352765321731567, + min_score=0.71484375)) + self.assertEqual( + score_calibration.parameters[510], + metadata_writer.CalibrationParameter( + scale=0.9901729226112366, + slope=0.8561913371086121, + offset=0.8783953189849854, + min_score=0.5859375)) + + def test_create_from_file_fail(self): + with tempfile.TemporaryDirectory() as temp_dir: + test_file = os.path.join(temp_dir, 'score_calibration.csv') + with open(test_file, 'w') as f: + f.write('0.98,0.5\n') + + with self.assertRaisesRegex( + ValueError, + 'Expected empty lines or 3 or 4 parameters per line in score ' + 'calibration file, but got 2.'): + metadata_writer.ScoreCalibration.create_from_file( + metadata_writer.ScoreCalibration.transformation_types.LOG, + test_file) + + with open(test_file, 'w') as f: + f.write('-0.98,0.5,0.34\n') + with self.assertRaisesRegex( + ValueError, + 'Expected scale to be a non-negative value, but got -0.98.'): + metadata_writer.ScoreCalibration.create_from_file( + metadata_writer.ScoreCalibration.transformation_types.LOG, + test_file) + + class MetadataWriterForTaskTest(absltest.TestCase): def setUp(self): @@ -197,7 +251,7 @@ class MetadataWriterForTaskTest(absltest.TestCase): "output_tensor_metadata": [ { "name": "score", - "description": "Score of the labels respectively", + "description": "Score of the labels respectively.", "content": { "content_properties_type": "FeatureProperties", "content_properties": { @@ -298,7 +352,7 @@ class MetadataWriterForTaskTest(absltest.TestCase): "output_tensor_metadata": [ { "name": "score", - "description": "Score of the labels respectively", + "description": "Score of the labels respectively.", "content": { "content_properties_type": "FeatureProperties", "content_properties": { @@ -353,6 +407,64 @@ class MetadataWriterForTaskTest(absltest.TestCase): } """) + def test_add_classification_output_with_score_thresholding(self): + writer = metadata_writer.MetadataWriter.create( + self.image_classifier_model_buffer) + writer.add_classification_output( + labels=metadata_writer.Labels().add(['a', 'b', 'c']), + score_thresholding=metadata_writer.ScoreThresholding( + global_score_threshold=0.5)) + _, metadata_json = writer.populate() + print(metadata_json) + self.assertJsonEqual( + metadata_json, """{ + "subgraph_metadata": [ + { + "input_tensor_metadata": [ + { + "name": "input" + } + ], + "output_tensor_metadata": [ + { + "name": "score", + "description": "Score of the labels respectively.", + "content": { + "content_properties_type": "FeatureProperties", + "content_properties": { + } + }, + "process_units": [ + { + "options_type": "ScoreThresholdingOptions", + "options": { + "global_score_threshold": 0.5 + } + } + ], + "stats": { + "max": [ + 1.0 + ], + "min": [ + 0.0 + ] + }, + "associated_files": [ + { + "name": "labels.txt", + "description": "Labels for categories that the model can recognize.", + "type": "TENSOR_AXIS_LABELS" + } + ] + } + ] + } + ], + "min_parser_version": "1.0.0" + } + """) + if __name__ == '__main__': absltest.main() diff --git a/mediapipe/tasks/python/test/test_utils.py b/mediapipe/tasks/python/test/test_utils.py index b428f830..6854791b 100644 --- a/mediapipe/tasks/python/test/test_utils.py +++ b/mediapipe/tasks/python/test/test_utils.py @@ -13,9 +13,15 @@ # limitations under the License. """Test util for MediaPipe Tasks.""" +import difflib import os from absl import flags +import six + +from google.protobuf import descriptor +from google.protobuf import descriptor_pool +from google.protobuf import text_format from mediapipe.python._framework_bindings import image as image_module from mediapipe.python._framework_bindings import image_frame as image_frame_module @@ -36,13 +42,15 @@ def test_srcdir(): raise RuntimeError("Missing TEST_SRCDIR environment.") -def get_test_data_path(file_or_dirname: str) -> str: +def get_test_data_path(file_or_dirname_path: str) -> str: """Returns full test data path.""" for (directory, subdirs, files) in os.walk(test_srcdir()): for f in subdirs + files: - if f.endswith(file_or_dirname): - return os.path.join(directory, f) - raise ValueError("No %s in test directory" % file_or_dirname) + path = os.path.join(directory, f) + if path.endswith(file_or_dirname_path): + return path + raise ValueError("No %s in test directory: %s." % + (file_or_dirname_path, test_srcdir())) def create_calibration_file(file_dir: str, @@ -53,3 +61,126 @@ def create_calibration_file(file_dir: str, with open(calibration_file, mode="w") as file: file.write(content) return calibration_file + + +def assert_proto_equals(self, + a, + b, + check_initialized=True, + normalize_numbers=True, + msg=None): + """assert_proto_equals() is useful for unit tests. + + It produces much more helpful output than assertEqual() for proto2 messages. + Fails with a useful error if a and b aren't equal. Comparison of repeated + fields matches the semantics of unittest.TestCase.assertEqual(), ie order and + extra duplicates fields matter. + + This is a fork of https://github.com/tensorflow/tensorflow/blob/ + master/tensorflow/python/util/protobuf/compare.py#L73. We use slightly + different rounding cutoffs to support Mac usage. + + Args: + self: absltest.testing.parameterized.TestCase + a: proto2 PB instance, or text string representing one. + b: proto2 PB instance -- message.Message or subclass thereof. + check_initialized: boolean, whether to fail if either a or b isn't + initialized. + normalize_numbers: boolean, whether to normalize types and precision of + numbers before comparison. + msg: if specified, is used as the error message on failure. + """ + pool = descriptor_pool.Default() + if isinstance(a, six.string_types): + a = text_format.Parse(a, b.__class__(), descriptor_pool=pool) + + for pb in a, b: + if check_initialized: + errors = pb.FindInitializationErrors() + if errors: + self.fail("Initialization errors: %s\n%s" % (errors, pb)) + if normalize_numbers: + _normalize_number_fields(pb) + + a_str = text_format.MessageToString(a, descriptor_pool=pool) + b_str = text_format.MessageToString(b, descriptor_pool=pool) + + # Some Python versions would perform regular diff instead of multi-line + # diff if string is longer than 2**16. We substitute this behavior + # with a call to unified_diff instead to have easier-to-read diffs. + # For context, see: https://bugs.python.org/issue11763. + if len(a_str) < 2**16 and len(b_str) < 2**16: + self.assertMultiLineEqual(a_str, b_str, msg=msg) + else: + diff = "".join( + difflib.unified_diff(a_str.splitlines(True), b_str.splitlines(True))) + if diff: + self.fail("%s :\n%s" % (msg, diff)) + + +def _normalize_number_fields(pb): + """Normalizes types and precisions of number fields in a protocol buffer. + + Due to subtleties in the python protocol buffer implementation, it is possible + for values to have different types and precision depending on whether they + were set and retrieved directly or deserialized from a protobuf. This function + normalizes integer values to ints and longs based on width, 32-bit floats to + five digits of precision to account for python always storing them as 64-bit, + and ensures doubles are floating point for when they're set to integers. + Modifies pb in place. Recurses into nested objects. https://github.com/tensorf + low/tensorflow/blob/master/tensorflow/python/util/protobuf/compare.py#L118 + + Args: + pb: proto2 message. + + Returns: + the given pb, modified in place. + """ + for desc, values in pb.ListFields(): + is_repeated = True + if desc.label != descriptor.FieldDescriptor.LABEL_REPEATED: + is_repeated = False + values = [values] + + normalized_values = None + + # We force 32-bit values to int and 64-bit values to long to make + # alternate implementations where the distinction is more significant + # (e.g. the C++ implementation) simpler. + if desc.type in (descriptor.FieldDescriptor.TYPE_INT64, + descriptor.FieldDescriptor.TYPE_UINT64, + descriptor.FieldDescriptor.TYPE_SINT64): + normalized_values = [int(x) for x in values] + elif desc.type in (descriptor.FieldDescriptor.TYPE_INT32, + descriptor.FieldDescriptor.TYPE_UINT32, + descriptor.FieldDescriptor.TYPE_SINT32, + descriptor.FieldDescriptor.TYPE_ENUM): + normalized_values = [int(x) for x in values] + elif desc.type == descriptor.FieldDescriptor.TYPE_FLOAT: + normalized_values = [round(x, 5) for x in values] + elif desc.type == descriptor.FieldDescriptor.TYPE_DOUBLE: + normalized_values = [round(float(x), 6) for x in values] + + if normalized_values is not None: + if is_repeated: + pb.ClearField(desc.name) + getattr(pb, desc.name).extend(normalized_values) + else: + setattr(pb, desc.name, normalized_values[0]) + + if (desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE or + desc.type == descriptor.FieldDescriptor.TYPE_GROUP): + if (desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE and + desc.message_type.has_options and + desc.message_type.GetOptions().map_entry): + # This is a map, only recurse if the values have a message type. + if (desc.message_type.fields_by_number[2].type == + descriptor.FieldDescriptor.TYPE_MESSAGE): + for v in six.itervalues(values): + _normalize_number_fields(v) + else: + for v in values: + # recursive step + _normalize_number_fields(v) + + return pb diff --git a/mediapipe/tasks/python/test/vision/BUILD b/mediapipe/tasks/python/test/vision/BUILD index 62595d37..6adab680 100644 --- a/mediapipe/tasks/python/test/vision/BUILD +++ b/mediapipe/tasks/python/test/vision/BUILD @@ -37,6 +37,43 @@ py_test( ], ) +py_test( + name = "image_classifier_test", + srcs = ["image_classifier_test.py"], + data = [ + "//mediapipe/tasks/testdata/vision:test_images", + "//mediapipe/tasks/testdata/vision:test_models", + ], + deps = [ + "//mediapipe/python:_framework_bindings", + "//mediapipe/tasks/python/components/containers:category", + "//mediapipe/tasks/python/components/containers:classifications", + "//mediapipe/tasks/python/components/containers:rect", + "//mediapipe/tasks/python/components/processors:classifier_options", + "//mediapipe/tasks/python/core:base_options", + "//mediapipe/tasks/python/test:test_utils", + "//mediapipe/tasks/python/vision:image_classifier", + "//mediapipe/tasks/python/vision/core:image_processing_options", + "//mediapipe/tasks/python/vision/core:vision_task_running_mode", + ], +) + +py_test( + name = "image_segmenter_test", + srcs = ["image_segmenter_test.py"], + data = [ + "//mediapipe/tasks/testdata/vision:test_images", + "//mediapipe/tasks/testdata/vision:test_models", + ], + deps = [ + "//mediapipe/python:_framework_bindings", + "//mediapipe/tasks/python/core:base_options", + "//mediapipe/tasks/python/test:test_utils", + "//mediapipe/tasks/python/vision:image_segmenter", + "//mediapipe/tasks/python/vision/core:vision_task_running_mode", + ], +) + py_test( name = "image_embedder_test", srcs = ["image_embedder_test.py"], diff --git a/mediapipe/tasks/python/test/vision/image_classifier_test.py b/mediapipe/tasks/python/test/vision/image_classifier_test.py new file mode 100644 index 00000000..11941ce2 --- /dev/null +++ b/mediapipe/tasks/python/test/vision/image_classifier_test.py @@ -0,0 +1,524 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for image classifier.""" + +import enum +import os +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized + +import numpy as np + +from mediapipe.python._framework_bindings import image +from mediapipe.tasks.python.components.containers import category +from mediapipe.tasks.python.components.containers import classifications as classifications_module +from mediapipe.tasks.python.components.containers import rect +from mediapipe.tasks.python.components.processors import classifier_options +from mediapipe.tasks.python.core import base_options as base_options_module +from mediapipe.tasks.python.test import test_utils +from mediapipe.tasks.python.vision import image_classifier +from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module +from mediapipe.tasks.python.vision.core import vision_task_running_mode + +_Rect = rect.Rect +_BaseOptions = base_options_module.BaseOptions +_ClassifierOptions = classifier_options.ClassifierOptions +_Category = category.Category +_ClassificationEntry = classifications_module.ClassificationEntry +_Classifications = classifications_module.Classifications +_ClassificationResult = classifications_module.ClassificationResult +_Image = image.Image +_ImageClassifier = image_classifier.ImageClassifier +_ImageClassifierOptions = image_classifier.ImageClassifierOptions +_RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode +_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions + +_MODEL_FILE = 'mobilenet_v2_1.0_224.tflite' +_IMAGE_FILE = 'burger.jpg' +_ALLOW_LIST = ['cheeseburger', 'guacamole'] +_DENY_LIST = ['cheeseburger'] +_SCORE_THRESHOLD = 0.5 +_MAX_RESULTS = 3 +_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision' + + +def _generate_empty_results(timestamp_ms: int) -> _ClassificationResult: + return _ClassificationResult(classifications=[ + _Classifications( + entries=[ + _ClassificationEntry(categories=[], timestamp_ms=timestamp_ms) + ], + head_index=0, + head_name='probability') + ]) + + +def _generate_burger_results(timestamp_ms: int) -> _ClassificationResult: + return _ClassificationResult(classifications=[ + _Classifications( + entries=[ + _ClassificationEntry( + categories=[ + _Category( + index=934, + score=0.793959, + display_name='', + category_name='cheeseburger'), + _Category( + index=932, + score=0.0273929, + display_name='', + category_name='bagel'), + _Category( + index=925, + score=0.0193408, + display_name='', + category_name='guacamole'), + _Category( + index=963, + score=0.00632786, + display_name='', + category_name='meat loaf') + ], + timestamp_ms=timestamp_ms) + ], + head_index=0, + head_name='probability') + ]) + + +def _generate_soccer_ball_results(timestamp_ms: int) -> _ClassificationResult: + return _ClassificationResult(classifications=[ + _Classifications( + entries=[ + _ClassificationEntry( + categories=[ + _Category( + index=806, + score=0.996527, + display_name='', + category_name='soccer ball') + ], + timestamp_ms=timestamp_ms) + ], + head_index=0, + head_name='probability') + ]) + + +class ModelFileType(enum.Enum): + FILE_CONTENT = 1 + FILE_NAME = 2 + + +class ImageClassifierTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + self.test_image = _Image.create_from_file( + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _IMAGE_FILE))) + self.model_path = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _MODEL_FILE)) + + def test_create_from_file_succeeds_with_valid_model_path(self): + # Creates with default option and valid model file successfully. + with _ImageClassifier.create_from_model_path(self.model_path) as classifier: + self.assertIsInstance(classifier, _ImageClassifier) + + def test_create_from_options_succeeds_with_valid_model_path(self): + # Creates with options containing model file successfully. + base_options = _BaseOptions(model_asset_path=self.model_path) + options = _ImageClassifierOptions(base_options=base_options) + with _ImageClassifier.create_from_options(options) as classifier: + self.assertIsInstance(classifier, _ImageClassifier) + + def test_create_from_options_fails_with_invalid_model_path(self): + # Invalid empty model path. + with self.assertRaisesRegex( + ValueError, + r"ExternalFile must specify at least one of 'file_content', " + r"'file_name', 'file_pointer_meta' or 'file_descriptor_meta'."): + base_options = _BaseOptions(model_asset_path='') + options = _ImageClassifierOptions(base_options=base_options) + _ImageClassifier.create_from_options(options) + + def test_create_from_options_succeeds_with_valid_model_content(self): + # Creates with options containing model content successfully. + with open(self.model_path, 'rb') as f: + base_options = _BaseOptions(model_asset_buffer=f.read()) + options = _ImageClassifierOptions(base_options=base_options) + classifier = _ImageClassifier.create_from_options(options) + self.assertIsInstance(classifier, _ImageClassifier) + + @parameterized.parameters( + (ModelFileType.FILE_NAME, 4, _generate_burger_results(0)), + (ModelFileType.FILE_CONTENT, 4, _generate_burger_results(0))) + def test_classify(self, model_file_type, max_results, + expected_classification_result): + # Creates classifier. + if model_file_type is ModelFileType.FILE_NAME: + base_options = _BaseOptions(model_asset_path=self.model_path) + elif model_file_type is ModelFileType.FILE_CONTENT: + with open(self.model_path, 'rb') as f: + model_content = f.read() + base_options = _BaseOptions(model_asset_buffer=model_content) + else: + # Should never happen + raise ValueError('model_file_type is invalid.') + + custom_classifier_options = _ClassifierOptions(max_results=max_results) + options = _ImageClassifierOptions( + base_options=base_options, classifier_options=custom_classifier_options) + classifier = _ImageClassifier.create_from_options(options) + + # Performs image classification on the input. + image_result = classifier.classify(self.test_image) + # Comparing results. + test_utils.assert_proto_equals(self, image_result.to_pb2(), + expected_classification_result.to_pb2()) + # Closes the classifier explicitly when the classifier is not used in + # a context. + classifier.close() + + @parameterized.parameters( + (ModelFileType.FILE_NAME, 4, _generate_burger_results(0)), + (ModelFileType.FILE_CONTENT, 4, _generate_burger_results(0))) + def test_classify_in_context(self, model_file_type, max_results, + expected_classification_result): + if model_file_type is ModelFileType.FILE_NAME: + base_options = _BaseOptions(model_asset_path=self.model_path) + elif model_file_type is ModelFileType.FILE_CONTENT: + with open(self.model_path, 'rb') as f: + model_content = f.read() + base_options = _BaseOptions(model_asset_buffer=model_content) + else: + # Should never happen + raise ValueError('model_file_type is invalid.') + + custom_classifier_options = _ClassifierOptions(max_results=max_results) + options = _ImageClassifierOptions( + base_options=base_options, classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as classifier: + # Performs image classification on the input. + image_result = classifier.classify(self.test_image) + # Comparing results. + test_utils.assert_proto_equals(self, image_result.to_pb2(), + expected_classification_result.to_pb2()) + + def test_classify_succeeds_with_region_of_interest(self): + base_options = _BaseOptions(model_asset_path=self.model_path) + custom_classifier_options = _ClassifierOptions(max_results=1) + options = _ImageClassifierOptions( + base_options=base_options, classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as classifier: + # Load the test image. + test_image = _Image.create_from_file( + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, 'multi_objects.jpg'))) + # Region-of-interest around the soccer ball. + roi = _Rect(left=0.45, top=0.3075, right=0.614, bottom=0.7345) + image_processing_options = _ImageProcessingOptions(roi) + # Performs image classification on the input. + image_result = classifier.classify(test_image, image_processing_options) + # Comparing results. + test_utils.assert_proto_equals(self, image_result.to_pb2(), + _generate_soccer_ball_results(0).to_pb2()) + + def test_score_threshold_option(self): + custom_classifier_options = _ClassifierOptions( + score_threshold=_SCORE_THRESHOLD) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as classifier: + # Performs image classification on the input. + image_result = classifier.classify(self.test_image) + classifications = image_result.classifications + + for classification in classifications: + for entry in classification.entries: + score = entry.categories[0].score + self.assertGreaterEqual( + score, _SCORE_THRESHOLD, + f'Classification with score lower than threshold found. ' + f'{classification}') + + def test_max_results_option(self): + custom_classifier_options = _ClassifierOptions( + score_threshold=_SCORE_THRESHOLD) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as classifier: + # Performs image classification on the input. + image_result = classifier.classify(self.test_image) + categories = image_result.classifications[0].entries[0].categories + + self.assertLessEqual( + len(categories), _MAX_RESULTS, 'Too many results returned.') + + def test_allow_list_option(self): + custom_classifier_options = _ClassifierOptions( + category_allowlist=_ALLOW_LIST) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as classifier: + # Performs image classification on the input. + image_result = classifier.classify(self.test_image) + classifications = image_result.classifications + + for classification in classifications: + for entry in classification.entries: + label = entry.categories[0].category_name + self.assertIn(label, _ALLOW_LIST, + f'Label {label} found but not in label allow list') + + def test_deny_list_option(self): + custom_classifier_options = _ClassifierOptions(category_denylist=_DENY_LIST) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as classifier: + # Performs image classification on the input. + image_result = classifier.classify(self.test_image) + classifications = image_result.classifications + + for classification in classifications: + for entry in classification.entries: + label = entry.categories[0].category_name + self.assertNotIn(label, _DENY_LIST, + f'Label {label} found but in deny list.') + + def test_combined_allowlist_and_denylist(self): + # Fails with combined allowlist and denylist + with self.assertRaisesRegex( + ValueError, + r'`category_allowlist` and `category_denylist` are mutually ' + r'exclusive options.'): + custom_classifier_options = _ClassifierOptions( + category_allowlist=['foo'], category_denylist=['bar']) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as unused_classifier: + pass + + def test_empty_classification_outputs(self): + custom_classifier_options = _ClassifierOptions(score_threshold=1) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as classifier: + # Performs image classification on the input. + image_result = classifier.classify(self.test_image) + self.assertEmpty(image_result.classifications[0].entries[0].categories) + + def test_missing_result_callback(self): + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM) + with self.assertRaisesRegex(ValueError, + r'result callback must be provided'): + with _ImageClassifier.create_from_options(options) as unused_classifier: + pass + + @parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO)) + def test_illegal_result_callback(self, running_mode): + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=running_mode, + result_callback=mock.MagicMock()) + with self.assertRaisesRegex(ValueError, + r'result callback should not be provided'): + with _ImageClassifier.create_from_options(options) as unused_classifier: + pass + + def test_calling_classify_for_video_in_image_mode(self): + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.IMAGE) + with _ImageClassifier.create_from_options(options) as classifier: + with self.assertRaisesRegex(ValueError, + r'not initialized with the video mode'): + classifier.classify_for_video(self.test_image, 0) + + def test_calling_classify_async_in_image_mode(self): + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.IMAGE) + with _ImageClassifier.create_from_options(options) as classifier: + with self.assertRaisesRegex(ValueError, + r'not initialized with the live stream mode'): + classifier.classify_async(self.test_image, 0) + + def test_calling_classify_in_video_mode(self): + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO) + with _ImageClassifier.create_from_options(options) as classifier: + with self.assertRaisesRegex(ValueError, + r'not initialized with the image mode'): + classifier.classify(self.test_image) + + def test_calling_classify_async_in_video_mode(self): + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO) + with _ImageClassifier.create_from_options(options) as classifier: + with self.assertRaisesRegex(ValueError, + r'not initialized with the live stream mode'): + classifier.classify_async(self.test_image, 0) + + def test_classify_for_video_with_out_of_order_timestamp(self): + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO) + with _ImageClassifier.create_from_options(options) as classifier: + unused_result = classifier.classify_for_video(self.test_image, 1) + with self.assertRaisesRegex( + ValueError, r'Input timestamp must be monotonically increasing'): + classifier.classify_for_video(self.test_image, 0) + + def test_classify_for_video(self): + custom_classifier_options = _ClassifierOptions(max_results=4) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO, + classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as classifier: + for timestamp in range(0, 300, 30): + classification_result = classifier.classify_for_video( + self.test_image, timestamp) + test_utils.assert_proto_equals( + self, classification_result.to_pb2(), + _generate_burger_results(timestamp).to_pb2()) + + def test_classify_for_video_succeeds_with_region_of_interest(self): + custom_classifier_options = _ClassifierOptions(max_results=1) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO, + classifier_options=custom_classifier_options) + with _ImageClassifier.create_from_options(options) as classifier: + # Load the test image. + test_image = _Image.create_from_file( + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, 'multi_objects.jpg'))) + # Region-of-interest around the soccer ball. + roi = _Rect(left=0.45, top=0.3075, right=0.614, bottom=0.7345) + image_processing_options = _ImageProcessingOptions(roi) + for timestamp in range(0, 300, 30): + classification_result = classifier.classify_for_video( + test_image, timestamp, image_processing_options) + test_utils.assert_proto_equals( + self, classification_result.to_pb2(), + _generate_soccer_ball_results(timestamp).to_pb2()) + + def test_calling_classify_in_live_stream_mode(self): + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=mock.MagicMock()) + with _ImageClassifier.create_from_options(options) as classifier: + with self.assertRaisesRegex(ValueError, + r'not initialized with the image mode'): + classifier.classify(self.test_image) + + def test_calling_classify_for_video_in_live_stream_mode(self): + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=mock.MagicMock()) + with _ImageClassifier.create_from_options(options) as classifier: + with self.assertRaisesRegex(ValueError, + r'not initialized with the video mode'): + classifier.classify_for_video(self.test_image, 0) + + def test_classify_async_calls_with_illegal_timestamp(self): + custom_classifier_options = _ClassifierOptions(max_results=4) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + classifier_options=custom_classifier_options, + result_callback=mock.MagicMock()) + with _ImageClassifier.create_from_options(options) as classifier: + classifier.classify_async(self.test_image, 100) + with self.assertRaisesRegex( + ValueError, r'Input timestamp must be monotonically increasing'): + classifier.classify_async(self.test_image, 0) + + @parameterized.parameters((0, _generate_burger_results), + (1, _generate_empty_results)) + def test_classify_async_calls(self, threshold, expected_result_fn): + observed_timestamp_ms = -1 + + def check_result(result: _ClassificationResult, output_image: _Image, + timestamp_ms: int): + test_utils.assert_proto_equals(self, result.to_pb2(), + expected_result_fn(timestamp_ms).to_pb2()) + self.assertTrue( + np.array_equal(output_image.numpy_view(), + self.test_image.numpy_view())) + self.assertLess(observed_timestamp_ms, timestamp_ms) + self.observed_timestamp_ms = timestamp_ms + + custom_classifier_options = _ClassifierOptions( + max_results=4, score_threshold=threshold) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + classifier_options=custom_classifier_options, + result_callback=check_result) + with _ImageClassifier.create_from_options(options) as classifier: + for timestamp in range(0, 300, 30): + classifier.classify_async(self.test_image, timestamp) + + def test_classify_async_succeeds_with_region_of_interest(self): + # Load the test image. + test_image = _Image.create_from_file( + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, 'multi_objects.jpg'))) + # Region-of-interest around the soccer ball. + roi = _Rect(left=0.45, top=0.3075, right=0.614, bottom=0.7345) + image_processing_options = _ImageProcessingOptions(roi) + observed_timestamp_ms = -1 + + def check_result(result: _ClassificationResult, output_image: _Image, + timestamp_ms: int): + test_utils.assert_proto_equals( + self, result.to_pb2(), + _generate_soccer_ball_results(timestamp_ms).to_pb2()) + self.assertEqual(output_image.width, test_image.width) + self.assertEqual(output_image.height, test_image.height) + self.assertLess(observed_timestamp_ms, timestamp_ms) + self.observed_timestamp_ms = timestamp_ms + + custom_classifier_options = _ClassifierOptions(max_results=1) + options = _ImageClassifierOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + classifier_options=custom_classifier_options, + result_callback=check_result) + with _ImageClassifier.create_from_options(options) as classifier: + for timestamp in range(0, 300, 30): + classifier.classify_async(test_image, timestamp, + image_processing_options) + + +if __name__ == '__main__': + absltest.main() diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py new file mode 100644 index 00000000..dde751c1 --- /dev/null +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -0,0 +1,359 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for image segmenter.""" + +import enum +import os +from typing import List +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized +import cv2 +import numpy as np + +from mediapipe.python._framework_bindings import image as image_module +from mediapipe.python._framework_bindings import image_frame +from mediapipe.tasks.python.core import base_options as base_options_module +from mediapipe.tasks.python.test import test_utils +from mediapipe.tasks.python.vision import image_segmenter +from mediapipe.tasks.python.vision.core import vision_task_running_mode + +_BaseOptions = base_options_module.BaseOptions +_Image = image_module.Image +_ImageFormat = image_frame.ImageFormat +_OutputType = image_segmenter.OutputType +_Activation = image_segmenter.Activation +_ImageSegmenter = image_segmenter.ImageSegmenter +_ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions +_RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode + +_MODEL_FILE = 'deeplabv3.tflite' +_IMAGE_FILE = 'segmentation_input_rotation0.jpg' +_SEGMENTATION_FILE = 'segmentation_golden_rotation0.png' +_MASK_MAGNIFICATION_FACTOR = 10 +_MASK_SIMILARITY_THRESHOLD = 0.98 +_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision' + + +def _similar_to_uint8_mask(actual_mask, expected_mask): + actual_mask_pixels = actual_mask.numpy_view().flatten() + expected_mask_pixels = expected_mask.numpy_view().flatten() + + consistent_pixels = 0 + num_pixels = len(expected_mask_pixels) + + for index in range(num_pixels): + consistent_pixels += ( + actual_mask_pixels[index] * + _MASK_MAGNIFICATION_FACTOR == expected_mask_pixels[index]) + + return consistent_pixels / num_pixels >= _MASK_SIMILARITY_THRESHOLD + + +class ModelFileType(enum.Enum): + FILE_CONTENT = 1 + FILE_NAME = 2 + + +class ImageSegmenterTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + # Load the test input image. + self.test_image = _Image.create_from_file( + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _IMAGE_FILE))) + # Loads ground truth segmentation file. + gt_segmentation_data = cv2.imread( + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _SEGMENTATION_FILE)), + cv2.IMREAD_GRAYSCALE) + self.test_seg_image = _Image(_ImageFormat.GRAY8, gt_segmentation_data) + self.model_path = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _MODEL_FILE)) + + def test_create_from_file_succeeds_with_valid_model_path(self): + # Creates with default option and valid model file successfully. + with _ImageSegmenter.create_from_model_path(self.model_path) as segmenter: + self.assertIsInstance(segmenter, _ImageSegmenter) + + def test_create_from_options_succeeds_with_valid_model_path(self): + # Creates with options containing model file successfully. + base_options = _BaseOptions(model_asset_path=self.model_path) + options = _ImageSegmenterOptions(base_options=base_options) + with _ImageSegmenter.create_from_options(options) as segmenter: + self.assertIsInstance(segmenter, _ImageSegmenter) + + def test_create_from_options_fails_with_invalid_model_path(self): + # Invalid empty model path. + with self.assertRaisesRegex( + ValueError, + r"ExternalFile must specify at least one of 'file_content', " + r"'file_name', 'file_pointer_meta' or 'file_descriptor_meta'."): + base_options = _BaseOptions(model_asset_path='') + options = _ImageSegmenterOptions(base_options=base_options) + _ImageSegmenter.create_from_options(options) + + def test_create_from_options_succeeds_with_valid_model_content(self): + # Creates with options containing model content successfully. + with open(self.model_path, 'rb') as f: + base_options = _BaseOptions(model_asset_buffer=f.read()) + options = _ImageSegmenterOptions(base_options=base_options) + segmenter = _ImageSegmenter.create_from_options(options) + self.assertIsInstance(segmenter, _ImageSegmenter) + + @parameterized.parameters((ModelFileType.FILE_NAME,), + (ModelFileType.FILE_CONTENT,)) + def test_segment_succeeds_with_category_mask(self, model_file_type): + # Creates segmenter. + if model_file_type is ModelFileType.FILE_NAME: + base_options = _BaseOptions(model_asset_path=self.model_path) + elif model_file_type is ModelFileType.FILE_CONTENT: + with open(self.model_path, 'rb') as f: + model_content = f.read() + base_options = _BaseOptions(model_asset_buffer=model_content) + else: + # Should never happen + raise ValueError('model_file_type is invalid.') + + options = _ImageSegmenterOptions( + base_options=base_options, output_type=_OutputType.CATEGORY_MASK) + segmenter = _ImageSegmenter.create_from_options(options) + + # Performs image segmentation on the input. + category_masks = segmenter.segment(self.test_image) + self.assertLen(category_masks, 1) + category_mask = category_masks[0] + result_pixels = category_mask.numpy_view().flatten() + + # Check if data type of `category_mask` is correct. + self.assertEqual(result_pixels.dtype, np.uint8) + + self.assertTrue( + _similar_to_uint8_mask(category_masks[0], self.test_seg_image), + f'Number of pixels in the candidate mask differing from that of the ' + f'ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.') + + # Closes the segmenter explicitly when the segmenter is not used in + # a context. + segmenter.close() + + def test_segment_succeeds_with_confidence_mask(self): + # Creates segmenter. + base_options = _BaseOptions(model_asset_path=self.model_path) + + # Run segmentation on the model in CATEGORY_MASK mode. + options = _ImageSegmenterOptions( + base_options=base_options, output_type=_OutputType.CATEGORY_MASK) + segmenter = _ImageSegmenter.create_from_options(options) + category_masks = segmenter.segment(self.test_image) + category_mask = category_masks[0].numpy_view() + + # Run segmentation on the model in CONFIDENCE_MASK mode. + options = _ImageSegmenterOptions( + base_options=base_options, + output_type=_OutputType.CONFIDENCE_MASK, + activation=_Activation.SOFTMAX) + segmenter = _ImageSegmenter.create_from_options(options) + confidence_masks = segmenter.segment(self.test_image) + + # Check if confidence mask shape is correct. + self.assertLen( + confidence_masks, 21, + 'Number of confidence masks must match with number of categories.') + + # Gather the confidence masks in a single array `confidence_mask_array`. + confidence_mask_array = np.array( + [confidence_mask.numpy_view() for confidence_mask in confidence_masks]) + + # Check if data type of `confidence_masks` are correct. + self.assertEqual(confidence_mask_array.dtype, np.float32) + + # Compute the category mask from the created confidence mask. + calculated_category_mask = np.argmax(confidence_mask_array, axis=0) + self.assertListEqual( + calculated_category_mask.tolist(), category_mask.tolist(), + 'Confidence mask does not match with the category mask.') + + # Closes the segmenter explicitly when the segmenter is not used in + # a context. + segmenter.close() + + @parameterized.parameters((ModelFileType.FILE_NAME), + (ModelFileType.FILE_CONTENT)) + def test_segment_in_context(self, model_file_type): + if model_file_type is ModelFileType.FILE_NAME: + base_options = _BaseOptions(model_asset_path=self.model_path) + elif model_file_type is ModelFileType.FILE_CONTENT: + with open(self.model_path, 'rb') as f: + model_contents = f.read() + base_options = _BaseOptions(model_asset_buffer=model_contents) + else: + # Should never happen + raise ValueError('model_file_type is invalid.') + + options = _ImageSegmenterOptions( + base_options=base_options, output_type=_OutputType.CATEGORY_MASK) + with _ImageSegmenter.create_from_options(options) as segmenter: + # Performs image segmentation on the input. + category_masks = segmenter.segment(self.test_image) + self.assertLen(category_masks, 1) + + self.assertTrue( + _similar_to_uint8_mask(category_masks[0], self.test_seg_image), + f'Number of pixels in the candidate mask differing from that of the ' + f'ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.') + + def test_missing_result_callback(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM) + with self.assertRaisesRegex(ValueError, + r'result callback must be provided'): + with _ImageSegmenter.create_from_options(options) as unused_segmenter: + pass + + @parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO)) + def test_illegal_result_callback(self, running_mode): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=running_mode, + result_callback=mock.MagicMock()) + with self.assertRaisesRegex(ValueError, + r'result callback should not be provided'): + with _ImageSegmenter.create_from_options(options) as unused_segmenter: + pass + + def test_calling_segment_for_video_in_image_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.IMAGE) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the video mode'): + segmenter.segment_for_video(self.test_image, 0) + + def test_calling_segment_async_in_image_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.IMAGE) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the live stream mode'): + segmenter.segment_async(self.test_image, 0) + + def test_calling_segment_in_video_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the image mode'): + segmenter.segment(self.test_image) + + def test_calling_segment_async_in_video_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the live stream mode'): + segmenter.segment_async(self.test_image, 0) + + def test_segment_for_video_with_out_of_order_timestamp(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO) + with _ImageSegmenter.create_from_options(options) as segmenter: + unused_result = segmenter.segment_for_video(self.test_image, 1) + with self.assertRaisesRegex( + ValueError, r'Input timestamp must be monotonically increasing'): + segmenter.segment_for_video(self.test_image, 0) + + def test_segment_for_video(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + output_type=_OutputType.CATEGORY_MASK, + running_mode=_RUNNING_MODE.VIDEO) + with _ImageSegmenter.create_from_options(options) as segmenter: + for timestamp in range(0, 300, 30): + category_masks = segmenter.segment_for_video(self.test_image, timestamp) + self.assertLen(category_masks, 1) + self.assertTrue( + _similar_to_uint8_mask(category_masks[0], self.test_seg_image), + f'Number of pixels in the candidate mask differing from that of the ' + f'ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.') + + def test_calling_segment_in_live_stream_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=mock.MagicMock()) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the image mode'): + segmenter.segment(self.test_image) + + def test_calling_segment_for_video_in_live_stream_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=mock.MagicMock()) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the video mode'): + segmenter.segment_for_video(self.test_image, 0) + + def test_segment_async_calls_with_illegal_timestamp(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=mock.MagicMock()) + with _ImageSegmenter.create_from_options(options) as segmenter: + segmenter.segment_async(self.test_image, 100) + with self.assertRaisesRegex( + ValueError, r'Input timestamp must be monotonically increasing'): + segmenter.segment_async(self.test_image, 0) + + def test_segment_async_calls(self): + observed_timestamp_ms = -1 + + def check_result(result: List[image_module.Image], output_image: _Image, + timestamp_ms: int): + # Get the output category mask. + category_mask = result[0] + self.assertEqual(output_image.width, self.test_image.width) + self.assertEqual(output_image.height, self.test_image.height) + self.assertEqual(output_image.width, self.test_seg_image.width) + self.assertEqual(output_image.height, self.test_seg_image.height) + self.assertTrue( + _similar_to_uint8_mask(category_mask, self.test_seg_image), + f'Number of pixels in the candidate mask differing from that of the ' + f'ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.') + self.assertLess(observed_timestamp_ms, timestamp_ms) + self.observed_timestamp_ms = timestamp_ms + + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + output_type=_OutputType.CATEGORY_MASK, + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=check_result) + with _ImageSegmenter.create_from_options(options) as segmenter: + for timestamp in range(0, 300, 30): + segmenter.segment_async(self.test_image, timestamp) + + +if __name__ == '__main__': + absltest.main() diff --git a/mediapipe/tasks/python/test/vision/object_detector_test.py b/mediapipe/tasks/python/test/vision/object_detector_test.py index d5cebd94..53c64427 100644 --- a/mediapipe/tasks/python/test/vision/object_detector_test.py +++ b/mediapipe/tasks/python/test/vision/object_detector_test.py @@ -14,6 +14,7 @@ """Tests for object detector.""" import enum +import os from unittest import mock from absl.testing import absltest @@ -87,6 +88,7 @@ _ALLOW_LIST = ['cat', 'dog'] _DENY_LIST = ['cat'] _SCORE_THRESHOLD = 0.3 _MAX_RESULTS = 3 +_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision' class ModelFileType(enum.Enum): @@ -99,8 +101,10 @@ class ObjectDetectorTest(parameterized.TestCase): def setUp(self): super().setUp() self.test_image = _Image.create_from_file( - test_utils.get_test_data_path(_IMAGE_FILE)) - self.model_path = test_utils.get_test_data_path(_MODEL_FILE) + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _IMAGE_FILE))) + self.model_path = test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _MODEL_FILE)) def test_create_from_file_succeeds_with_valid_model_path(self): # Creates with default option and valid model file successfully. @@ -395,5 +399,6 @@ class ObjectDetectorTest(parameterized.TestCase): detector.detect_async(self.test_image, timestamp) detector.close() + if __name__ == '__main__': absltest.main() diff --git a/mediapipe/tasks/python/vision/BUILD b/mediapipe/tasks/python/vision/BUILD index 08c2709f..96cf6288 100644 --- a/mediapipe/tasks/python/vision/BUILD +++ b/mediapipe/tasks/python/vision/BUILD @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Placeholder for internal Python strict library compatibility macro. +# Placeholder for internal Python strict library and test compatibility macro. package(default_visibility = ["//mediapipe/tasks:internal"]) @@ -37,6 +37,72 @@ py_library( ], ) +py_library( + name = "image_classifier", + srcs = [ + "image_classifier.py", + ], + deps = [ + "//mediapipe/python:_framework_bindings", + "//mediapipe/python:packet_creator", + "//mediapipe/python:packet_getter", + "//mediapipe/tasks/cc/components/containers/proto:classifications_py_pb2", + "//mediapipe/tasks/cc/vision/image_classifier/proto:image_classifier_graph_options_py_pb2", + "//mediapipe/tasks/python/components/containers:classifications", + "//mediapipe/tasks/python/components/containers:rect", + "//mediapipe/tasks/python/components/processors:classifier_options", + "//mediapipe/tasks/python/core:base_options", + "//mediapipe/tasks/python/core:optional_dependencies", + "//mediapipe/tasks/python/core:task_info", + "//mediapipe/tasks/python/vision/core:base_vision_task_api", + "//mediapipe/tasks/python/vision/core:image_processing_options", + "//mediapipe/tasks/python/vision/core:vision_task_running_mode", + ], +) + +py_library( + name = "image_segmenter", + srcs = [ + "image_segmenter.py", + ], + deps = [ + "//mediapipe/python:_framework_bindings", + "//mediapipe/python:packet_creator", + "//mediapipe/python:packet_getter", + "//mediapipe/tasks/cc/components/proto:segmenter_options_py_pb2", + "//mediapipe/tasks/cc/vision/image_segmenter/proto:image_segmenter_options_py_pb2", + "//mediapipe/tasks/python/core:base_options", + "//mediapipe/tasks/python/core:optional_dependencies", + "//mediapipe/tasks/python/core:task_info", + "//mediapipe/tasks/python/vision/core:base_vision_task_api", + "//mediapipe/tasks/python/vision/core:vision_task_running_mode", + ], +) + +py_library( + name = "gesture_recognizer", + srcs = [ + "gesture_recognizer.py", + ], + deps = [ + "//mediapipe/framework/formats:classification_py_pb2", + "//mediapipe/framework/formats:landmark_py_pb2", + "//mediapipe/python:_framework_bindings", + "//mediapipe/python:packet_creator", + "//mediapipe/python:packet_getter", + "//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_recognizer_graph_options_py_pb2", + "//mediapipe/tasks/python/components/containers:category", + "//mediapipe/tasks/python/components/containers:landmark", + "//mediapipe/tasks/python/components/processors:classifier_options", + "//mediapipe/tasks/python/core:base_options", + "//mediapipe/tasks/python/core:optional_dependencies", + "//mediapipe/tasks/python/core:task_info", + "//mediapipe/tasks/python/vision/core:base_vision_task_api", + "//mediapipe/tasks/python/vision/core:image_processing_options", + "//mediapipe/tasks/python/vision/core:vision_task_running_mode", + ], +) + py_library( name = "image_embedder", srcs = [ diff --git a/mediapipe/tasks/python/vision/__init__.py b/mediapipe/tasks/python/vision/__init__.py index 65c1214a..def11317 100644 --- a/mediapipe/tasks/python/vision/__init__.py +++ b/mediapipe/tasks/python/vision/__init__.py @@ -11,3 +11,21 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + +"""MediaPipe Tasks Vision API.""" + +import mediapipe.tasks.python.vision.core +import mediapipe.tasks.python.vision.image_classifier +import mediapipe.tasks.python.vision.object_detector + +ImageClassifier = image_classifier.ImageClassifier +ImageClassifierOptions = image_classifier.ImageClassifierOptions +ObjectDetector = object_detector.ObjectDetector +ObjectDetectorOptions = object_detector.ObjectDetectorOptions +RunningMode = core.vision_task_running_mode.VisionTaskRunningMode + +# Remove unnecessary modules to avoid duplication in API docs. +del core +del image_classifier +del object_detector +del mediapipe diff --git a/mediapipe/tasks/python/vision/core/BUILD b/mediapipe/tasks/python/vision/core/BUILD index c7422969..e2b2b3de 100644 --- a/mediapipe/tasks/python/vision/core/BUILD +++ b/mediapipe/tasks/python/vision/core/BUILD @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Placeholder for internal Python strict library compatibility macro. +# Placeholder for internal Python strict library and test compatibility macro. package(default_visibility = ["//mediapipe/tasks:internal"]) @@ -23,15 +23,25 @@ py_library( srcs = ["vision_task_running_mode.py"], ) +py_library( + name = "image_processing_options", + srcs = ["image_processing_options.py"], + deps = [ + "//mediapipe/tasks/python/components/containers:rect", + ], +) + py_library( name = "base_vision_task_api", srcs = [ "base_vision_task_api.py", ], deps = [ + ":image_processing_options", ":vision_task_running_mode", "//mediapipe/framework:calculator_py_pb2", "//mediapipe/python:_framework_bindings", + "//mediapipe/tasks/python/components/containers:rect", "//mediapipe/tasks/python/core:optional_dependencies", ], ) diff --git a/mediapipe/tasks/python/vision/core/base_vision_task_api.py b/mediapipe/tasks/python/vision/core/base_vision_task_api.py index b2f8a366..01617039 100644 --- a/mediapipe/tasks/python/vision/core/base_vision_task_api.py +++ b/mediapipe/tasks/python/vision/core/base_vision_task_api.py @@ -13,17 +13,22 @@ # limitations under the License. """MediaPipe vision task base api.""" +import math from typing import Callable, Mapping, Optional from mediapipe.framework import calculator_pb2 from mediapipe.python._framework_bindings import packet as packet_module from mediapipe.python._framework_bindings import task_runner as task_runner_module +from mediapipe.tasks.python.components.containers import rect as rect_module from mediapipe.tasks.python.core.optional_dependencies import doc_controls +from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module _TaskRunner = task_runner_module.TaskRunner _Packet = packet_module.Packet +_NormalizedRect = rect_module.NormalizedRect _RunningMode = running_mode_module.VisionTaskRunningMode +_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions class BaseVisionTaskApi(object): @@ -122,6 +127,49 @@ class BaseVisionTaskApi(object): + self._running_mode.name) self._runner.send(inputs) + def convert_to_normalized_rect(self, + options: _ImageProcessingOptions, + roi_allowed: bool = True) -> _NormalizedRect: + """Converts from ImageProcessingOptions to NormalizedRect, performing sanity checks on-the-fly. + + If the input ImageProcessingOptions is not present, returns a default + NormalizedRect covering the whole image with rotation set to 0. If + 'roi_allowed' is false, an error will be returned if the input + ImageProcessingOptions has its 'region_of_interest' field set. + + Args: + options: Options for image processing. + roi_allowed: Indicates if the `region_of_interest` field is allowed to be + set. By default, it's set to True. + + Returns: + A normalized rect proto that repesents the image processing options. + """ + normalized_rect = _NormalizedRect( + rotation=0, x_center=0.5, y_center=0.5, width=1, height=1) + if options is None: + return normalized_rect + + if options.rotation_degrees % 90 != 0: + raise ValueError('Expected rotation to be a multiple of 90°.') + + # Convert to radians counter-clockwise. + normalized_rect.rotation = -options.rotation_degrees * math.pi / 180.0 + + if options.region_of_interest: + if not roi_allowed: + raise ValueError("This task doesn't support region-of-interest.") + roi = options.region_of_interest + if roi.left >= roi.right or roi.top >= roi.bottom: + raise ValueError('Expected Rect with left < right and top < bottom.') + if roi.left < 0 or roi.top < 0 or roi.right > 1 or roi.bottom > 1: + raise ValueError('Expected Rect values to be in [0,1].') + normalized_rect.x_center = (roi.left + roi.right) / 2.0 + normalized_rect.y_center = (roi.top + roi.bottom) / 2.0 + normalized_rect.width = roi.right - roi.left + normalized_rect.height = roi.bottom - roi.top + return normalized_rect + def close(self) -> None: """Shuts down the mediapipe vision task instance. diff --git a/mediapipe/tasks/python/vision/core/image_processing_options.py b/mediapipe/tasks/python/vision/core/image_processing_options.py new file mode 100644 index 00000000..9a0cd2b4 --- /dev/null +++ b/mediapipe/tasks/python/vision/core/image_processing_options.py @@ -0,0 +1,39 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MediaPipe vision options for image processing.""" + +import dataclasses +from typing import Optional + +from mediapipe.tasks.python.components.containers import rect as rect_module + + +@dataclasses.dataclass +class ImageProcessingOptions: + """Options for image processing. + + If both region-of-interest and rotation are specified, the crop around the + region-of-interest is extracted first, then the specified rotation is applied + to the crop. + + Attributes: + region_of_interest: The optional region-of-interest to crop from the image. + If not specified, the full image is used. Coordinates must be in [0,1] + with 'left' < 'right' and 'top' < 'bottom'. + rotation_degrees: The rotation to apply to the image (or cropped + region-of-interest), in degrees clockwise. The rotation must be a multiple + (positive or negative) of 90°. + """ + region_of_interest: Optional[rect_module.Rect] = None + rotation_degrees: int = 0 diff --git a/mediapipe/tasks/python/vision/gesture_recognizer.py b/mediapipe/tasks/python/vision/gesture_recognizer.py new file mode 100644 index 00000000..081ecc42 --- /dev/null +++ b/mediapipe/tasks/python/vision/gesture_recognizer.py @@ -0,0 +1,426 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MediaPipe gesture recognizer task.""" + +import dataclasses +from typing import Callable, Mapping, Optional, List + +from mediapipe.framework.formats import classification_pb2 +from mediapipe.framework.formats import landmark_pb2 +from mediapipe.python import packet_creator +from mediapipe.python import packet_getter +from mediapipe.python._framework_bindings import image as image_module +from mediapipe.python._framework_bindings import packet as packet_module +from mediapipe.tasks.cc.vision.gesture_recognizer.proto import gesture_recognizer_graph_options_pb2 +from mediapipe.tasks.python.components.containers import category as category_module +from mediapipe.tasks.python.components.containers import landmark as landmark_module +from mediapipe.tasks.python.components.processors import classifier_options +from mediapipe.tasks.python.core import base_options as base_options_module +from mediapipe.tasks.python.core import task_info as task_info_module +from mediapipe.tasks.python.core.optional_dependencies import doc_controls +from mediapipe.tasks.python.vision.core import base_vision_task_api +from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module +from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module + +_BaseOptions = base_options_module.BaseOptions +_GestureRecognizerGraphOptionsProto = gesture_recognizer_graph_options_pb2.GestureRecognizerGraphOptions +_ClassifierOptions = classifier_options.ClassifierOptions +_RunningMode = running_mode_module.VisionTaskRunningMode +_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions +_TaskInfo = task_info_module.TaskInfo + +_IMAGE_IN_STREAM_NAME = 'image_in' +_IMAGE_OUT_STREAM_NAME = 'image_out' +_IMAGE_TAG = 'IMAGE' +_NORM_RECT_STREAM_NAME = 'norm_rect_in' +_NORM_RECT_TAG = 'NORM_RECT' +_HAND_GESTURE_STREAM_NAME = 'hand_gestures' +_HAND_GESTURE_TAG = 'HAND_GESTURES' +_HANDEDNESS_STREAM_NAME = 'handedness' +_HANDEDNESS_TAG = 'HANDEDNESS' +_HAND_LANDMARKS_STREAM_NAME = 'landmarks' +_HAND_LANDMARKS_TAG = 'LANDMARKS' +_HAND_WORLD_LANDMARKS_STREAM_NAME = 'world_landmarks' +_HAND_WORLD_LANDMARKS_TAG = 'WORLD_LANDMARKS' +_TASK_GRAPH_NAME = 'mediapipe.tasks.vision.gesture_recognizer.GestureRecognizerGraph' +_MICRO_SECONDS_PER_MILLISECOND = 1000 +_GESTURE_DEFAULT_INDEX = -1 + + +@dataclasses.dataclass +class GestureRecognitionResult: + """The gesture recognition result from GestureRecognizer, where each vector element represents a single hand detected in the image. + + Attributes: + gestures: Recognized hand gestures of detected hands. Note that the index of + the gesture is always -1, because the raw indices from multiple gesture + classifiers cannot consolidate to a meaningful index. + handedness: Classification of handedness. + hand_landmarks: Detected hand landmarks in normalized image coordinates. + hand_world_landmarks: Detected hand landmarks in world coordinates. + """ + + gestures: List[List[category_module.Category]] + handedness: List[List[category_module.Category]] + hand_landmarks: List[List[landmark_module.NormalizedLandmark]] + hand_world_landmarks: List[List[landmark_module.Landmark]] + + +def _build_recognition_result( + output_packets: Mapping[str, + packet_module.Packet]) -> GestureRecognitionResult: + """Consturcts a `GestureRecognitionResult` from output packets.""" + gestures_proto_list = packet_getter.get_proto_list( + output_packets[_HAND_GESTURE_STREAM_NAME]) + handedness_proto_list = packet_getter.get_proto_list( + output_packets[_HANDEDNESS_STREAM_NAME]) + hand_landmarks_proto_list = packet_getter.get_proto_list( + output_packets[_HAND_LANDMARKS_STREAM_NAME]) + hand_world_landmarks_proto_list = packet_getter.get_proto_list( + output_packets[_HAND_WORLD_LANDMARKS_STREAM_NAME]) + + gesture_results = [] + for proto in gestures_proto_list: + gesture_categories = [] + gesture_classifications = classification_pb2.ClassificationList() + gesture_classifications.MergeFrom(proto) + for gesture in gesture_classifications.classification: + gesture_categories.append( + category_module.Category( + index=_GESTURE_DEFAULT_INDEX, + score=gesture.score, + display_name=gesture.display_name, + category_name=gesture.label)) + gesture_results.append(gesture_categories) + + handedness_results = [] + for proto in handedness_proto_list: + handedness_categories = [] + handedness_classifications = classification_pb2.ClassificationList() + handedness_classifications.MergeFrom(proto) + for handedness in handedness_classifications.classification: + handedness_categories.append( + category_module.Category( + index=handedness.index, + score=handedness.score, + display_name=handedness.display_name, + category_name=handedness.label)) + handedness_results.append(handedness_categories) + + hand_landmarks_results = [] + for proto in hand_landmarks_proto_list: + hand_landmarks = landmark_pb2.NormalizedLandmarkList() + hand_landmarks.MergeFrom(proto) + hand_landmarks_results.append([ + landmark_module.NormalizedLandmark.create_from_pb2(hand_landmark) + for hand_landmark in hand_landmarks.landmark + ]) + + hand_world_landmarks_results = [] + for proto in hand_world_landmarks_proto_list: + hand_world_landmarks = landmark_pb2.LandmarkList() + hand_world_landmarks.MergeFrom(proto) + hand_world_landmarks_results.append([ + landmark_module.Landmark.create_from_pb2(hand_world_landmark) + for hand_world_landmark in hand_world_landmarks.landmark + ]) + + return GestureRecognitionResult(gesture_results, handedness_results, + hand_landmarks_results, + hand_world_landmarks_results) + + +@dataclasses.dataclass +class GestureRecognizerOptions: + """Options for the gesture recognizer task. + + Attributes: + base_options: Base options for the hand gesture recognizer task. + running_mode: The running mode of the task. Default to the image mode. + Gesture recognizer task has three running modes: 1) The image mode for + recognizing hand gestures on single image inputs. 2) The video mode for + recognizing hand gestures on the decoded frames of a video. 3) The live + stream mode for recognizing hand gestures on a live stream of input data, + such as from camera. + num_hands: The maximum number of hands can be detected by the recognizer. + min_hand_detection_confidence: The minimum confidence score for the hand + detection to be considered successful. + min_hand_presence_confidence: The minimum confidence score of hand presence + score in the hand landmark detection. + min_tracking_confidence: The minimum confidence score for the hand tracking + to be considered successful. + canned_gesture_classifier_options: Options for configuring the canned + gestures classifier, such as score threshold, allow list and deny list of + gestures. The categories for canned gesture classifiers are: ["None", + "Closed_Fist", "Open_Palm", "Pointing_Up", "Thumb_Down", "Thumb_Up", + "Victory", "ILoveYou"]. Note this option is subject to change. + custom_gesture_classifier_options: Options for configuring the custom + gestures classifier, such as score threshold, allow list and deny list of + gestures. Note this option is subject to change. + result_callback: The user-defined result callback for processing live stream + data. The result callback should only be specified when the running mode + is set to the live stream mode. + """ + base_options: _BaseOptions + running_mode: _RunningMode = _RunningMode.IMAGE + num_hands: Optional[int] = 1 + min_hand_detection_confidence: Optional[float] = 0.5 + min_hand_presence_confidence: Optional[float] = 0.5 + min_tracking_confidence: Optional[float] = 0.5 + canned_gesture_classifier_options: Optional[ + _ClassifierOptions] = _ClassifierOptions() + custom_gesture_classifier_options: Optional[ + _ClassifierOptions] = _ClassifierOptions() + result_callback: Optional[Callable[ + [GestureRecognitionResult, image_module.Image, int], None]] = None + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _GestureRecognizerGraphOptionsProto: + """Generates an GestureRecognizerOptions protobuf object.""" + base_options_proto = self.base_options.to_pb2() + base_options_proto.use_stream_mode = False if self.running_mode == _RunningMode.IMAGE else True + + # Initialize gesture recognizer options from base options. + gesture_recognizer_options_proto = _GestureRecognizerGraphOptionsProto( + base_options=base_options_proto) + # Configure hand detector and hand landmarker options. + hand_landmarker_options_proto = gesture_recognizer_options_proto.hand_landmarker_graph_options + hand_landmarker_options_proto.min_tracking_confidence = self.min_tracking_confidence + hand_landmarker_options_proto.hand_detector_graph_options.num_hands = self.num_hands + hand_landmarker_options_proto.hand_detector_graph_options.min_detection_confidence = self.min_hand_detection_confidence + hand_landmarker_options_proto.hand_landmarks_detector_graph_options.min_detection_confidence = self.min_hand_presence_confidence + + # Configure hand gesture recognizer options. + hand_gesture_recognizer_options_proto = gesture_recognizer_options_proto.hand_gesture_recognizer_graph_options + hand_gesture_recognizer_options_proto.canned_gesture_classifier_graph_options.classifier_options.CopyFrom( + self.canned_gesture_classifier_options.to_pb2()) + hand_gesture_recognizer_options_proto.custom_gesture_classifier_graph_options.classifier_options.CopyFrom( + self.custom_gesture_classifier_options.to_pb2()) + + return gesture_recognizer_options_proto + + +class GestureRecognizer(base_vision_task_api.BaseVisionTaskApi): + """Class that performs gesture recognition on images.""" + + @classmethod + def create_from_model_path(cls, model_path: str) -> 'GestureRecognizer': + """Creates an `GestureRecognizer` object from a TensorFlow Lite model and the default `GestureRecognizerOptions`. + + Note that the created `GestureRecognizer` instance is in image mode, for + recognizing hand gestures on single image inputs. + + Args: + model_path: Path to the model. + + Returns: + `GestureRecognizer` object that's created from the model file and the + default `GestureRecognizerOptions`. + + Raises: + ValueError: If failed to create `GestureRecognizer` object from the + provided file such as invalid file path. + RuntimeError: If other types of error occurred. + """ + base_options = _BaseOptions(model_asset_path=model_path) + options = GestureRecognizerOptions( + base_options=base_options, running_mode=_RunningMode.IMAGE) + return cls.create_from_options(options) + + @classmethod + def create_from_options( + cls, options: GestureRecognizerOptions) -> 'GestureRecognizer': + """Creates the `GestureRecognizer` object from gesture recognizer options. + + Args: + options: Options for the gesture recognizer task. + + Returns: + `GestureRecognizer` object that's created from `options`. + + Raises: + ValueError: If failed to create `GestureRecognizer` object from + `GestureRecognizerOptions` such as missing the model. + RuntimeError: If other types of error occurred. + """ + + def packets_callback(output_packets: Mapping[str, packet_module.Packet]): + if output_packets[_IMAGE_OUT_STREAM_NAME].is_empty(): + return + + image = packet_getter.get_image(output_packets[_IMAGE_OUT_STREAM_NAME]) + + if output_packets[_HAND_GESTURE_STREAM_NAME].is_empty(): + empty_packet = output_packets[_HAND_GESTURE_STREAM_NAME] + options.result_callback( + GestureRecognitionResult([], [], [], []), image, + empty_packet.timestamp.value // _MICRO_SECONDS_PER_MILLISECOND) + return + + gesture_recognition_result = _build_recognition_result(output_packets) + timestamp = output_packets[_HAND_GESTURE_STREAM_NAME].timestamp + options.result_callback(gesture_recognition_result, image, + timestamp.value // _MICRO_SECONDS_PER_MILLISECOND) + + task_info = _TaskInfo( + task_graph=_TASK_GRAPH_NAME, + input_streams=[ + ':'.join([_IMAGE_TAG, _IMAGE_IN_STREAM_NAME]), + ':'.join([_NORM_RECT_TAG, _NORM_RECT_STREAM_NAME]), + ], + output_streams=[ + ':'.join([_HAND_GESTURE_TAG, _HAND_GESTURE_STREAM_NAME]), + ':'.join([_HANDEDNESS_TAG, _HANDEDNESS_STREAM_NAME]), + ':'.join([_HAND_LANDMARKS_TAG, + _HAND_LANDMARKS_STREAM_NAME]), ':'.join([ + _HAND_WORLD_LANDMARKS_TAG, + _HAND_WORLD_LANDMARKS_STREAM_NAME + ]), ':'.join([_IMAGE_TAG, _IMAGE_OUT_STREAM_NAME]) + ], + task_options=options) + return cls( + task_info.generate_graph_config( + enable_flow_limiting=options.running_mode == + _RunningMode.LIVE_STREAM), options.running_mode, + packets_callback if options.result_callback else None) + + def recognize( + self, + image: image_module.Image, + image_processing_options: Optional[_ImageProcessingOptions] = None + ) -> GestureRecognitionResult: + """Performs hand gesture recognition on the given image. + + Only use this method when the GestureRecognizer is created with the image + running mode. + + The image can be of any size with format RGB or RGBA. + TODO: Describes how the input image will be preprocessed after the yuv + support is implemented. + + Args: + image: MediaPipe Image. + image_processing_options: Options for image processing. + + Returns: + The hand gesture recognition results. + + Raises: + ValueError: If any of the input arguments is invalid. + RuntimeError: If gesture recognition failed to run. + """ + normalized_rect = self.convert_to_normalized_rect( + image_processing_options, roi_allowed=False) + output_packets = self._process_image_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image), + _NORM_RECT_STREAM_NAME: + packet_creator.create_proto(normalized_rect.to_pb2()) + }) + + if output_packets[_HAND_GESTURE_STREAM_NAME].is_empty(): + return GestureRecognitionResult([], [], [], []) + + return _build_recognition_result(output_packets) + + def recognize_for_video( + self, + image: image_module.Image, + timestamp_ms: int, + image_processing_options: Optional[_ImageProcessingOptions] = None + ) -> GestureRecognitionResult: + """Performs gesture recognition on the provided video frame. + + Only use this method when the GestureRecognizer is created with the video + running mode. + + Only use this method when the GestureRecognizer is created with the video + running mode. It's required to provide the video frame's timestamp (in + milliseconds) along with the video frame. The input timestamps should be + monotonically increasing for adjacent calls of this method. + + Args: + image: MediaPipe Image. + timestamp_ms: The timestamp of the input video frame in milliseconds. + image_processing_options: Options for image processing. + + Returns: + The hand gesture recognition results. + + Raises: + ValueError: If any of the input arguments is invalid. + RuntimeError: If gesture recognition failed to run. + """ + normalized_rect = self.convert_to_normalized_rect( + image_processing_options, roi_allowed=False) + output_packets = self._process_video_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND), + _NORM_RECT_STREAM_NAME: + packet_creator.create_proto(normalized_rect.to_pb2()).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND) + }) + + if output_packets[_HAND_GESTURE_STREAM_NAME].is_empty(): + return GestureRecognitionResult([], [], [], []) + + return _build_recognition_result(output_packets) + + def recognize_async( + self, + image: image_module.Image, + timestamp_ms: int, + image_processing_options: Optional[_ImageProcessingOptions] = None + ) -> None: + """Sends live image data to perform gesture recognition. + + The results will be available via the "result_callback" provided in the + GestureRecognizerOptions. Only use this method when the GestureRecognizer + is created with the live stream running mode. + + Only use this method when the GestureRecognizer is created with the live + stream running mode. The input timestamps should be monotonically increasing + for adjacent calls of this method. This method will return immediately after + the input image is accepted. The results will be available via the + `result_callback` provided in the `GestureRecognizerOptions`. The + `recognize_async` method is designed to process live stream data such as + camera input. To lower the overall latency, gesture recognizer may drop the + input images if needed. In other words, it's not guaranteed to have output + per input image. + + The `result_callback` provides: + - The hand gesture recognition results. + - The input image that the gesture recognizer runs on. + - The input timestamp in milliseconds. + + Args: + image: MediaPipe Image. + timestamp_ms: The timestamp of the input image in milliseconds. + image_processing_options: Options for image processing. + + Raises: + ValueError: If the current input timestamp is smaller than what the + gesture recognizer has already processed. + """ + normalized_rect = self.convert_to_normalized_rect( + image_processing_options, roi_allowed=False) + self._send_live_stream_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND), + _NORM_RECT_STREAM_NAME: + packet_creator.create_proto(normalized_rect.to_pb2()).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND) + }) diff --git a/mediapipe/tasks/python/vision/image_classifier.py b/mediapipe/tasks/python/vision/image_classifier.py new file mode 100644 index 00000000..02819ddf --- /dev/null +++ b/mediapipe/tasks/python/vision/image_classifier.py @@ -0,0 +1,292 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MediaPipe image classifier task.""" + +import dataclasses +from typing import Callable, Mapping, Optional + +from mediapipe.python import packet_creator +from mediapipe.python import packet_getter +# TODO: Import MPImage directly one we have an alias +from mediapipe.python._framework_bindings import image as image_module +from mediapipe.python._framework_bindings import packet +from mediapipe.tasks.cc.components.containers.proto import classifications_pb2 +from mediapipe.tasks.cc.vision.image_classifier.proto import image_classifier_graph_options_pb2 +from mediapipe.tasks.python.components.containers import classifications +from mediapipe.tasks.python.components.containers import rect +from mediapipe.tasks.python.components.processors import classifier_options +from mediapipe.tasks.python.core import base_options as base_options_module +from mediapipe.tasks.python.core import task_info as task_info_module +from mediapipe.tasks.python.core.optional_dependencies import doc_controls +from mediapipe.tasks.python.vision.core import base_vision_task_api +from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module +from mediapipe.tasks.python.vision.core import vision_task_running_mode + +_NormalizedRect = rect.NormalizedRect +_BaseOptions = base_options_module.BaseOptions +_ImageClassifierGraphOptionsProto = image_classifier_graph_options_pb2.ImageClassifierGraphOptions +_ClassifierOptions = classifier_options.ClassifierOptions +_RunningMode = vision_task_running_mode.VisionTaskRunningMode +_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions +_TaskInfo = task_info_module.TaskInfo + +_CLASSIFICATION_RESULT_OUT_STREAM_NAME = 'classification_result_out' +_CLASSIFICATION_RESULT_TAG = 'CLASSIFICATION_RESULT' +_IMAGE_IN_STREAM_NAME = 'image_in' +_IMAGE_OUT_STREAM_NAME = 'image_out' +_IMAGE_TAG = 'IMAGE' +_NORM_RECT_STREAM_NAME = 'norm_rect_in' +_NORM_RECT_TAG = 'NORM_RECT' +_TASK_GRAPH_NAME = 'mediapipe.tasks.vision.image_classifier.ImageClassifierGraph' +_MICRO_SECONDS_PER_MILLISECOND = 1000 + + +@dataclasses.dataclass +class ImageClassifierOptions: + """Options for the image classifier task. + + Attributes: + base_options: Base options for the image classifier task. + running_mode: The running mode of the task. Default to the image mode. Image + classifier task has three running modes: 1) The image mode for classifying + objects on single image inputs. 2) The video mode for classifying objects + on the decoded frames of a video. 3) The live stream mode for classifying + objects on a live stream of input data, such as from camera. + classifier_options: Options for the image classification task. + result_callback: The user-defined result callback for processing live stream + data. The result callback should only be specified when the running mode + is set to the live stream mode. + """ + base_options: _BaseOptions + running_mode: _RunningMode = _RunningMode.IMAGE + classifier_options: _ClassifierOptions = _ClassifierOptions() + result_callback: Optional[ + Callable[[classifications.ClassificationResult, image_module.Image, int], + None]] = None + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _ImageClassifierGraphOptionsProto: + """Generates an ImageClassifierOptions protobuf object.""" + base_options_proto = self.base_options.to_pb2() + base_options_proto.use_stream_mode = False if self.running_mode == _RunningMode.IMAGE else True + classifier_options_proto = self.classifier_options.to_pb2() + + return _ImageClassifierGraphOptionsProto( + base_options=base_options_proto, + classifier_options=classifier_options_proto) + + +class ImageClassifier(base_vision_task_api.BaseVisionTaskApi): + """Class that performs image classification on images.""" + + @classmethod + def create_from_model_path(cls, model_path: str) -> 'ImageClassifier': + """Creates an `ImageClassifier` object from a TensorFlow Lite model and the default `ImageClassifierOptions`. + + Note that the created `ImageClassifier` instance is in image mode, for + classifying objects on single image inputs. + + Args: + model_path: Path to the model. + + Returns: + `ImageClassifier` object that's created from the model file and the + default `ImageClassifierOptions`. + + Raises: + ValueError: If failed to create `ImageClassifier` object from the provided + file such as invalid file path. + RuntimeError: If other types of error occurred. + """ + base_options = _BaseOptions(model_asset_path=model_path) + options = ImageClassifierOptions( + base_options=base_options, running_mode=_RunningMode.IMAGE) + return cls.create_from_options(options) + + @classmethod + def create_from_options(cls, + options: ImageClassifierOptions) -> 'ImageClassifier': + """Creates the `ImageClassifier` object from image classifier options. + + Args: + options: Options for the image classifier task. + + Returns: + `ImageClassifier` object that's created from `options`. + + Raises: + ValueError: If failed to create `ImageClassifier` object from + `ImageClassifierOptions` such as missing the model. + RuntimeError: If other types of error occurred. + """ + + def packets_callback(output_packets: Mapping[str, packet.Packet]): + if output_packets[_IMAGE_OUT_STREAM_NAME].is_empty(): + return + + classification_result_proto = classifications_pb2.ClassificationResult() + classification_result_proto.CopyFrom( + packet_getter.get_proto( + output_packets[_CLASSIFICATION_RESULT_OUT_STREAM_NAME])) + + classification_result = classifications.ClassificationResult([ + classifications.Classifications.create_from_pb2(classification) + for classification in classification_result_proto.classifications + ]) + image = packet_getter.get_image(output_packets[_IMAGE_OUT_STREAM_NAME]) + timestamp = output_packets[_IMAGE_OUT_STREAM_NAME].timestamp + options.result_callback(classification_result, image, + timestamp.value // _MICRO_SECONDS_PER_MILLISECOND) + + task_info = _TaskInfo( + task_graph=_TASK_GRAPH_NAME, + input_streams=[ + ':'.join([_IMAGE_TAG, _IMAGE_IN_STREAM_NAME]), + ':'.join([_NORM_RECT_TAG, _NORM_RECT_STREAM_NAME]), + ], + output_streams=[ + ':'.join([ + _CLASSIFICATION_RESULT_TAG, + _CLASSIFICATION_RESULT_OUT_STREAM_NAME + ]), ':'.join([_IMAGE_TAG, _IMAGE_OUT_STREAM_NAME]) + ], + task_options=options) + return cls( + task_info.generate_graph_config( + enable_flow_limiting=options.running_mode == + _RunningMode.LIVE_STREAM), options.running_mode, + packets_callback if options.result_callback else None) + + def classify( + self, + image: image_module.Image, + image_processing_options: Optional[_ImageProcessingOptions] = None + ) -> classifications.ClassificationResult: + """Performs image classification on the provided MediaPipe Image. + + Args: + image: MediaPipe Image. + image_processing_options: Options for image processing. + + Returns: + A classification result object that contains a list of classifications. + + Raises: + ValueError: If any of the input arguments is invalid. + RuntimeError: If image classification failed to run. + """ + normalized_rect = self.convert_to_normalized_rect(image_processing_options) + output_packets = self._process_image_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image), + _NORM_RECT_STREAM_NAME: + packet_creator.create_proto(normalized_rect.to_pb2()) + }) + + classification_result_proto = classifications_pb2.ClassificationResult() + classification_result_proto.CopyFrom( + packet_getter.get_proto( + output_packets[_CLASSIFICATION_RESULT_OUT_STREAM_NAME])) + + return classifications.ClassificationResult([ + classifications.Classifications.create_from_pb2(classification) + for classification in classification_result_proto.classifications + ]) + + def classify_for_video( + self, + image: image_module.Image, + timestamp_ms: int, + image_processing_options: Optional[_ImageProcessingOptions] = None + ) -> classifications.ClassificationResult: + """Performs image classification on the provided video frames. + + Only use this method when the ImageClassifier is created with the video + running mode. It's required to provide the video frame's timestamp (in + milliseconds) along with the video frame. The input timestamps should be + monotonically increasing for adjacent calls of this method. + + Args: + image: MediaPipe Image. + timestamp_ms: The timestamp of the input video frame in milliseconds. + image_processing_options: Options for image processing. + + Returns: + A classification result object that contains a list of classifications. + + Raises: + ValueError: If any of the input arguments is invalid. + RuntimeError: If image classification failed to run. + """ + normalized_rect = self.convert_to_normalized_rect(image_processing_options) + output_packets = self._process_video_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND), + _NORM_RECT_STREAM_NAME: + packet_creator.create_proto(normalized_rect.to_pb2()).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND) + }) + + classification_result_proto = classifications_pb2.ClassificationResult() + classification_result_proto.CopyFrom( + packet_getter.get_proto( + output_packets[_CLASSIFICATION_RESULT_OUT_STREAM_NAME])) + + return classifications.ClassificationResult([ + classifications.Classifications.create_from_pb2(classification) + for classification in classification_result_proto.classifications + ]) + + def classify_async( + self, + image: image_module.Image, + timestamp_ms: int, + image_processing_options: Optional[_ImageProcessingOptions] = None + ) -> None: + """Sends live image data (an Image with a unique timestamp) to perform image classification. + + Only use this method when the ImageClassifier is created with the live + stream running mode. The input timestamps should be monotonically increasing + for adjacent calls of this method. This method will return immediately after + the input image is accepted. The results will be available via the + `result_callback` provided in the `ImageClassifierOptions`. The + `classify_async` method is designed to process live stream data such as + camera input. To lower the overall latency, image classifier may drop the + input images if needed. In other words, it's not guaranteed to have output + per input image. + + The `result_callback` provides: + - A classification result object that contains a list of classifications. + - The input image that the image classifier runs on. + - The input timestamp in milliseconds. + + Args: + image: MediaPipe Image. + timestamp_ms: The timestamp of the input image in milliseconds. + image_processing_options: Options for image processing. + + Raises: + ValueError: If the current input timestamp is smaller than what the image + classifier has already processed. + """ + normalized_rect = self.convert_to_normalized_rect(image_processing_options) + self._send_live_stream_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND), + _NORM_RECT_STREAM_NAME: + packet_creator.create_proto(normalized_rect.to_pb2()).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND) + }) diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py new file mode 100644 index 00000000..ebfeca1b --- /dev/null +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -0,0 +1,251 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MediaPipe image segmenter task.""" + +import dataclasses +import enum +from typing import Callable, List, Mapping, Optional + +from mediapipe.python import packet_creator +from mediapipe.python import packet_getter +from mediapipe.python._framework_bindings import image as image_module +from mediapipe.python._framework_bindings import packet +from mediapipe.tasks.cc.components.proto import segmenter_options_pb2 +from mediapipe.tasks.cc.vision.image_segmenter.proto import image_segmenter_options_pb2 +from mediapipe.tasks.python.core import base_options as base_options_module +from mediapipe.tasks.python.core import task_info as task_info_module +from mediapipe.tasks.python.core.optional_dependencies import doc_controls +from mediapipe.tasks.python.vision.core import base_vision_task_api +from mediapipe.tasks.python.vision.core import vision_task_running_mode + +_BaseOptions = base_options_module.BaseOptions +_SegmenterOptionsProto = segmenter_options_pb2.SegmenterOptions +_ImageSegmenterOptionsProto = image_segmenter_options_pb2.ImageSegmenterOptions +_RunningMode = vision_task_running_mode.VisionTaskRunningMode +_TaskInfo = task_info_module.TaskInfo + +_SEGMENTATION_OUT_STREAM_NAME = 'segmented_mask_out' +_SEGMENTATION_TAG = 'GROUPED_SEGMENTATION' +_IMAGE_IN_STREAM_NAME = 'image_in' +_IMAGE_OUT_STREAM_NAME = 'image_out' +_IMAGE_TAG = 'IMAGE' +_TASK_GRAPH_NAME = 'mediapipe.tasks.vision.ImageSegmenterGraph' +_MICRO_SECONDS_PER_MILLISECOND = 1000 + + +class OutputType(enum.Enum): + UNSPECIFIED = 0 + CATEGORY_MASK = 1 + CONFIDENCE_MASK = 2 + + +class Activation(enum.Enum): + NONE = 0 + SIGMOID = 1 + SOFTMAX = 2 + + +@dataclasses.dataclass +class ImageSegmenterOptions: + """Options for the image segmenter task. + + Attributes: + base_options: Base options for the image segmenter task. + running_mode: The running mode of the task. Default to the image mode. Image + segmenter task has three running modes: 1) The image mode for segmenting + objects on single image inputs. 2) The video mode for segmenting objects + on the decoded frames of a video. 3) The live stream mode for segmenting + objects on a live stream of input data, such as from camera. + output_type: The output mask type allows specifying the type of + post-processing to perform on the raw model results. + activation: Activation function to apply to input tensor. + result_callback: The user-defined result callback for processing live stream + data. The result callback should only be specified when the running mode + is set to the live stream mode. + """ + base_options: _BaseOptions + running_mode: _RunningMode = _RunningMode.IMAGE + output_type: Optional[OutputType] = OutputType.CATEGORY_MASK + activation: Optional[Activation] = Activation.NONE + result_callback: Optional[Callable[ + [List[image_module.Image], image_module.Image, int], None]] = None + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _ImageSegmenterOptionsProto: + """Generates an ImageSegmenterOptions protobuf object.""" + base_options_proto = self.base_options.to_pb2() + base_options_proto.use_stream_mode = False if self.running_mode == _RunningMode.IMAGE else True + segmenter_options_proto = _SegmenterOptionsProto( + output_type=self.output_type.value, activation=self.activation.value) + return _ImageSegmenterOptionsProto( + base_options=base_options_proto, + segmenter_options=segmenter_options_proto) + + +class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): + """Class that performs image segmentation on images.""" + + @classmethod + def create_from_model_path(cls, model_path: str) -> 'ImageSegmenter': + """Creates an `ImageSegmenter` object from a TensorFlow Lite model and the default `ImageSegmenterOptions`. + + Note that the created `ImageSegmenter` instance is in image mode, for + performing image segmentation on single image inputs. + + Args: + model_path: Path to the model. + + Returns: + `ImageSegmenter` object that's created from the model file and the default + `ImageSegmenterOptions`. + + Raises: + ValueError: If failed to create `ImageSegmenter` object from the provided + file such as invalid file path. + RuntimeError: If other types of error occurred. + """ + base_options = _BaseOptions(model_asset_path=model_path) + options = ImageSegmenterOptions( + base_options=base_options, running_mode=_RunningMode.IMAGE) + return cls.create_from_options(options) + + @classmethod + def create_from_options(cls, + options: ImageSegmenterOptions) -> 'ImageSegmenter': + """Creates the `ImageSegmenter` object from image segmenter options. + + Args: + options: Options for the image segmenter task. + + Returns: + `ImageSegmenter` object that's created from `options`. + + Raises: + ValueError: If failed to create `ImageSegmenter` object from + `ImageSegmenterOptions` such as missing the model. + RuntimeError: If other types of error occurred. + """ + + def packets_callback(output_packets: Mapping[str, packet.Packet]): + if output_packets[_IMAGE_OUT_STREAM_NAME].is_empty(): + return + segmentation_result = packet_getter.get_image_list( + output_packets[_SEGMENTATION_OUT_STREAM_NAME]) + image = packet_getter.get_image(output_packets[_IMAGE_OUT_STREAM_NAME]) + timestamp = output_packets[_SEGMENTATION_OUT_STREAM_NAME].timestamp + options.result_callback(segmentation_result, image, + timestamp.value // _MICRO_SECONDS_PER_MILLISECOND) + + task_info = _TaskInfo( + task_graph=_TASK_GRAPH_NAME, + input_streams=[':'.join([_IMAGE_TAG, _IMAGE_IN_STREAM_NAME])], + output_streams=[ + ':'.join([_SEGMENTATION_TAG, _SEGMENTATION_OUT_STREAM_NAME]), + ':'.join([_IMAGE_TAG, _IMAGE_OUT_STREAM_NAME]) + ], + task_options=options) + return cls( + task_info.generate_graph_config( + enable_flow_limiting=options.running_mode == + _RunningMode.LIVE_STREAM), options.running_mode, + packets_callback if options.result_callback else None) + + def segment(self, image: image_module.Image) -> List[image_module.Image]: + """Performs the actual segmentation task on the provided MediaPipe Image. + + Args: + image: MediaPipe Image. + + Returns: + If the output_type is CATEGORY_MASK, the returned vector of images is + per-category segmented image mask. + If the output_type is CONFIDENCE_MASK, the returned vector of images + contains only one confidence image mask. A segmentation result object that + contains a list of segmentation masks as images. + + Raises: + ValueError: If any of the input arguments is invalid. + RuntimeError: If image segmentation failed to run. + """ + output_packets = self._process_image_data( + {_IMAGE_IN_STREAM_NAME: packet_creator.create_image(image)}) + segmentation_result = packet_getter.get_image_list( + output_packets[_SEGMENTATION_OUT_STREAM_NAME]) + return segmentation_result + + def segment_for_video(self, image: image_module.Image, + timestamp_ms: int) -> List[image_module.Image]: + """Performs segmentation on the provided video frames. + + Only use this method when the ImageSegmenter is created with the video + running mode. It's required to provide the video frame's timestamp (in + milliseconds) along with the video frame. The input timestamps should be + monotonically increasing for adjacent calls of this method. + + Args: + image: MediaPipe Image. + timestamp_ms: The timestamp of the input video frame in milliseconds. + + Returns: + If the output_type is CATEGORY_MASK, the returned vector of images is + per-category segmented image mask. + If the output_type is CONFIDENCE_MASK, the returned vector of images + contains only one confidence image mask. A segmentation result object that + contains a list of segmentation masks as images. + + Raises: + ValueError: If any of the input arguments is invalid. + RuntimeError: If image segmentation failed to run. + """ + output_packets = self._process_video_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND) + }) + segmentation_result = packet_getter.get_image_list( + output_packets[_SEGMENTATION_OUT_STREAM_NAME]) + return segmentation_result + + def segment_async(self, image: image_module.Image, timestamp_ms: int) -> None: + """Sends live image data (an Image with a unique timestamp) to perform image segmentation. + + Only use this method when the ImageSegmenter is created with the live stream + running mode. The input timestamps should be monotonically increasing for + adjacent calls of this method. This method will return immediately after the + input image is accepted. The results will be available via the + `result_callback` provided in the `ImageSegmenterOptions`. The + `segment_async` method is designed to process live stream data such as + camera input. To lower the overall latency, image segmenter may drop the + input images if needed. In other words, it's not guaranteed to have output + per input image. + + The `result_callback` prvoides: + - A segmentation result object that contains a list of segmentation masks + as images. + - The input image that the image segmenter runs on. + - The input timestamp in milliseconds. + + Args: + image: MediaPipe Image. + timestamp_ms: The timestamp of the input image in milliseconds. + + Raises: + ValueError: If the current input timestamp is smaller than what the image + segmenter has already processed. + """ + self._send_live_stream_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND) + }) diff --git a/mediapipe/tasks/python/vision/object_detector.py b/mediapipe/tasks/python/vision/object_detector.py index a50e5586..f6177cda 100644 --- a/mediapipe/tasks/python/vision/object_detector.py +++ b/mediapipe/tasks/python/vision/object_detector.py @@ -20,7 +20,6 @@ from mediapipe.python import packet_creator from mediapipe.python import packet_getter from mediapipe.python._framework_bindings import image as image_module from mediapipe.python._framework_bindings import packet as packet_module -from mediapipe.python._framework_bindings import task_runner as task_runner_module from mediapipe.tasks.cc.vision.object_detector.proto import object_detector_options_pb2 from mediapipe.tasks.python.components.containers import detections as detections_module from mediapipe.tasks.python.core import base_options as base_options_module @@ -33,7 +32,6 @@ _BaseOptions = base_options_module.BaseOptions _ObjectDetectorOptionsProto = object_detector_options_pb2.ObjectDetectorOptions _RunningMode = running_mode_module.VisionTaskRunningMode _TaskInfo = task_info_module.TaskInfo -_TaskRunner = task_runner_module.TaskRunner _DETECTIONS_OUT_STREAM_NAME = 'detections_out' _DETECTIONS_TAG = 'DETECTIONS' diff --git a/mediapipe/tasks/testdata/metadata/BUILD b/mediapipe/tasks/testdata/metadata/BUILD index 9f50368b..8ed6e1ca 100644 --- a/mediapipe/tasks/testdata/metadata/BUILD +++ b/mediapipe/tasks/testdata/metadata/BUILD @@ -29,6 +29,8 @@ mediapipe_files(srcs = [ "mobile_object_classifier_v0_2_3-metadata-no-name.tflite", "mobilenet_v1_0.25_224_1_default_1.tflite", "mobilenet_v2_1.0_224_quant.tflite", + "mobilenet_v2_1.0_224_quant_without_metadata.tflite", + "mobilenet_v2_1.0_224_without_metadata.tflite", ]) exports_files([ @@ -48,6 +50,10 @@ exports_files([ "score_calibration.txt", "score_calibration_file_meta.json", "score_calibration_tensor_meta.json", + "score_thresholding_meta.json", + "labels.txt", + "mobilenet_v2_1.0_224.json", + "mobilenet_v2_1.0_224_quant.json", ]) filegroup( @@ -59,6 +65,8 @@ filegroup( "mobile_object_classifier_v0_2_3-metadata-no-name.tflite", "mobilenet_v1_0.25_224_1_default_1.tflite", "mobilenet_v2_1.0_224_quant.tflite", + "mobilenet_v2_1.0_224_quant_without_metadata.tflite", + "mobilenet_v2_1.0_224_without_metadata.tflite", ], ) @@ -78,8 +86,12 @@ filegroup( "input_image_tensor_float_meta.json", "input_image_tensor_uint8_meta.json", "input_image_tensor_unsupported_meta.json", + "labels.txt", + "mobilenet_v2_1.0_224.json", + "mobilenet_v2_1.0_224_quant.json", "score_calibration.txt", "score_calibration_file_meta.json", "score_calibration_tensor_meta.json", + "score_thresholding_meta.json", ], ) diff --git a/mediapipe/tasks/testdata/metadata/labels.txt b/mediapipe/tasks/testdata/metadata/labels.txt new file mode 100644 index 00000000..fe811239 --- /dev/null +++ b/mediapipe/tasks/testdata/metadata/labels.txt @@ -0,0 +1,1001 @@ +background +tench +goldfish +great white shark +tiger shark +hammerhead +electric ray +stingray +cock +hen +ostrich +brambling +goldfinch +house finch +junco +indigo bunting +robin +bulbul +jay +magpie +chickadee +water ouzel +kite +bald eagle +vulture +great grey owl +European fire salamander +common newt +eft +spotted salamander +axolotl +bullfrog +tree frog +tailed frog +loggerhead +leatherback turtle +mud turtle +terrapin +box turtle +banded gecko +common iguana +American chameleon +whiptail +agama +frilled lizard +alligator lizard +Gila monster +green lizard +African chameleon +Komodo dragon +African crocodile +American alligator +triceratops +thunder snake +ringneck snake +hognose snake +green snake +king snake +garter snake +water snake +vine snake +night snake +boa constrictor +rock python +Indian cobra +green mamba +sea snake +horned viper +diamondback +sidewinder +trilobite +harvestman +scorpion +black and gold garden spider +barn spider +garden spider +black widow +tarantula +wolf spider +tick +centipede +black grouse +ptarmigan +ruffed grouse +prairie chicken +peacock +quail +partridge +African grey +macaw +sulphur-crested cockatoo +lorikeet +coucal +bee eater +hornbill +hummingbird +jacamar +toucan +drake +red-breasted merganser +goose +black swan +tusker +echidna +platypus +wallaby +koala +wombat +jellyfish +sea anemone +brain coral +flatworm +nematode +conch +snail +slug +sea slug +chiton +chambered nautilus +Dungeness crab +rock crab +fiddler crab +king crab +American lobster +spiny lobster +crayfish +hermit crab +isopod +white stork +black stork +spoonbill +flamingo +little blue heron +American egret +bittern +crane +limpkin +European gallinule +American coot +bustard +ruddy turnstone +red-backed sandpiper +redshank +dowitcher +oystercatcher +pelican +king penguin +albatross +grey whale +killer whale +dugong +sea lion +Chihuahua +Japanese spaniel +Maltese dog +Pekinese +Shih-Tzu +Blenheim spaniel +papillon +toy terrier +Rhodesian ridgeback +Afghan hound +basset +beagle +bloodhound +bluetick +black-and-tan coonhound +Walker hound +English foxhound +redbone +borzoi +Irish wolfhound +Italian greyhound +whippet +Ibizan hound +Norwegian elkhound +otterhound +Saluki +Scottish deerhound +Weimaraner +Staffordshire bullterrier +American Staffordshire terrier +Bedlington terrier +Border terrier +Kerry blue terrier +Irish terrier +Norfolk terrier +Norwich terrier +Yorkshire terrier +wire-haired fox terrier +Lakeland terrier +Sealyham terrier +Airedale +cairn +Australian terrier +Dandie Dinmont +Boston bull +miniature schnauzer +giant schnauzer +standard schnauzer +Scotch terrier +Tibetan terrier +silky terrier +soft-coated wheaten terrier +West Highland white terrier +Lhasa +flat-coated retriever +curly-coated retriever +golden retriever +Labrador retriever +Chesapeake Bay retriever +German short-haired pointer +vizsla +English setter +Irish setter +Gordon setter +Brittany spaniel +clumber +English springer +Welsh springer spaniel +cocker spaniel +Sussex spaniel +Irish water spaniel +kuvasz +schipperke +groenendael +malinois +briard +kelpie +komondor +Old English sheepdog +Shetland sheepdog +collie +Border collie +Bouvier des Flandres +Rottweiler +German shepherd +Doberman +miniature pinscher +Greater Swiss Mountain dog +Bernese mountain dog +Appenzeller +EntleBucher +boxer +bull mastiff +Tibetan mastiff +French bulldog +Great Dane +Saint Bernard +Eskimo dog +malamute +Siberian husky +dalmatian +affenpinscher +basenji +pug +Leonberg +Newfoundland +Great Pyrenees +Samoyed +Pomeranian +chow +keeshond +Brabancon griffon +Pembroke +Cardigan +toy poodle +miniature poodle +standard poodle +Mexican hairless +timber wolf +white wolf +red wolf +coyote +dingo +dhole +African hunting dog +hyena +red fox +kit fox +Arctic fox +grey fox +tabby +tiger cat +Persian cat +Siamese cat +Egyptian cat +cougar +lynx +leopard +snow leopard +jaguar +lion +tiger +cheetah +brown bear +American black bear +ice bear +sloth bear +mongoose +meerkat +tiger beetle +ladybug +ground beetle +long-horned beetle +leaf beetle +dung beetle +rhinoceros beetle +weevil +fly +bee +ant +grasshopper +cricket +walking stick +cockroach +mantis +cicada +leafhopper +lacewing +dragonfly +damselfly +admiral +ringlet +monarch +cabbage butterfly +sulphur butterfly +lycaenid +starfish +sea urchin +sea cucumber +wood rabbit +hare +Angora +hamster +porcupine +fox squirrel +marmot +beaver +guinea pig +sorrel +zebra +hog +wild boar +warthog +hippopotamus +ox +water buffalo +bison +ram +bighorn +ibex +hartebeest +impala +gazelle +Arabian camel +llama +weasel +mink +polecat +black-footed ferret +otter +skunk +badger +armadillo +three-toed sloth +orangutan +gorilla +chimpanzee +gibbon +siamang +guenon +patas +baboon +macaque +langur +colobus +proboscis monkey +marmoset +capuchin +howler monkey +titi +spider monkey +squirrel monkey +Madagascar cat +indri +Indian elephant +African elephant +lesser panda +giant panda +barracouta +eel +coho +rock beauty +anemone fish +sturgeon +gar +lionfish +puffer +abacus +abaya +academic gown +accordion +acoustic guitar +aircraft carrier +airliner +airship +altar +ambulance +amphibian +analog clock +apiary +apron +ashcan +assault rifle +backpack +bakery +balance beam +balloon +ballpoint +Band Aid +banjo +bannister +barbell +barber chair +barbershop +barn +barometer +barrel +barrow +baseball +basketball +bassinet +bassoon +bathing cap +bath towel +bathtub +beach wagon +beacon +beaker +bearskin +beer bottle +beer glass +bell cote +bib +bicycle-built-for-two +bikini +binder +binoculars +birdhouse +boathouse +bobsled +bolo tie +bonnet +bookcase +bookshop +bottlecap +bow +bow tie +brass +brassiere +breakwater +breastplate +broom +bucket +buckle +bulletproof vest +bullet train +butcher shop +cab +caldron +candle +cannon +canoe +can opener +cardigan +car mirror +carousel +carpenter's kit +carton +car wheel +cash machine +cassette +cassette player +castle +catamaran +CD player +cello +cellular telephone +chain +chainlink fence +chain mail +chain saw +chest +chiffonier +chime +china cabinet +Christmas stocking +church +cinema +cleaver +cliff dwelling +cloak +clog +cocktail shaker +coffee mug +coffeepot +coil +combination lock +computer keyboard +confectionery +container ship +convertible +corkscrew +cornet +cowboy boot +cowboy hat +cradle +crane +crash helmet +crate +crib +Crock Pot +croquet ball +crutch +cuirass +dam +desk +desktop computer +dial telephone +diaper +digital clock +digital watch +dining table +dishrag +dishwasher +disk brake +dock +dogsled +dome +doormat +drilling platform +drum +drumstick +dumbbell +Dutch oven +electric fan +electric guitar +electric locomotive +entertainment center +envelope +espresso maker +face powder +feather boa +file +fireboat +fire engine +fire screen +flagpole +flute +folding chair +football helmet +forklift +fountain +fountain pen +four-poster +freight car +French horn +frying pan +fur coat +garbage truck +gasmask +gas pump +goblet +go-kart +golf ball +golfcart +gondola +gong +gown +grand piano +greenhouse +grille +grocery store +guillotine +hair slide +hair spray +half track +hammer +hamper +hand blower +hand-held computer +handkerchief +hard disc +harmonica +harp +harvester +hatchet +holster +home theater +honeycomb +hook +hoopskirt +horizontal bar +horse cart +hourglass +iPod +iron +jack-o'-lantern +jean +jeep +jersey +jigsaw puzzle +jinrikisha +joystick +kimono +knee pad +knot +lab coat +ladle +lampshade +laptop +lawn mower +lens cap +letter opener +library +lifeboat +lighter +limousine +liner +lipstick +Loafer +lotion +loudspeaker +loupe +lumbermill +magnetic compass +mailbag +mailbox +maillot +maillot +manhole cover +maraca +marimba +mask +matchstick +maypole +maze +measuring cup +medicine chest +megalith +microphone +microwave +military uniform +milk can +minibus +miniskirt +minivan +missile +mitten +mixing bowl +mobile home +Model T +modem +monastery +monitor +moped +mortar +mortarboard +mosque +mosquito net +motor scooter +mountain bike +mountain tent +mouse +mousetrap +moving van +muzzle +nail +neck brace +necklace +nipple +notebook +obelisk +oboe +ocarina +odometer +oil filter +organ +oscilloscope +overskirt +oxcart +oxygen mask +packet +paddle +paddlewheel +padlock +paintbrush +pajama +palace +panpipe +paper towel +parachute +parallel bars +park bench +parking meter +passenger car +patio +pay-phone +pedestal +pencil box +pencil sharpener +perfume +Petri dish +photocopier +pick +pickelhaube +picket fence +pickup +pier +piggy bank +pill bottle +pillow +ping-pong ball +pinwheel +pirate +pitcher +plane +planetarium +plastic bag +plate rack +plow +plunger +Polaroid camera +pole +police van +poncho +pool table +pop bottle +pot +potter's wheel +power drill +prayer rug +printer +prison +projectile +projector +puck +punching bag +purse +quill +quilt +racer +racket +radiator +radio +radio telescope +rain barrel +recreational vehicle +reel +reflex camera +refrigerator +remote control +restaurant +revolver +rifle +rocking chair +rotisserie +rubber eraser +rugby ball +rule +running shoe +safe +safety pin +saltshaker +sandal +sarong +sax +scabbard +scale +school bus +schooner +scoreboard +screen +screw +screwdriver +seat belt +sewing machine +shield +shoe shop +shoji +shopping basket +shopping cart +shovel +shower cap +shower curtain +ski +ski mask +sleeping bag +slide rule +sliding door +slot +snorkel +snowmobile +snowplow +soap dispenser +soccer ball +sock +solar dish +sombrero +soup bowl +space bar +space heater +space shuttle +spatula +speedboat +spider web +spindle +sports car +spotlight +stage +steam locomotive +steel arch bridge +steel drum +stethoscope +stole +stone wall +stopwatch +stove +strainer +streetcar +stretcher +studio couch +stupa +submarine +suit +sundial +sunglass +sunglasses +sunscreen +suspension bridge +swab +sweatshirt +swimming trunks +swing +switch +syringe +table lamp +tank +tape player +teapot +teddy +television +tennis ball +thatch +theater curtain +thimble +thresher +throne +tile roof +toaster +tobacco shop +toilet seat +torch +totem pole +tow truck +toyshop +tractor +trailer truck +tray +trench coat +tricycle +trimaran +tripod +triumphal arch +trolleybus +trombone +tub +turnstile +typewriter keyboard +umbrella +unicycle +upright +vacuum +vase +vault +velvet +vending machine +vestment +viaduct +violin +volleyball +waffle iron +wall clock +wallet +wardrobe +warplane +washbasin +washer +water bottle +water jug +water tower +whiskey jug +whistle +wig +window screen +window shade +Windsor tie +wine bottle +wing +wok +wooden spoon +wool +worm fence +wreck +yawl +yurt +web site +comic book +crossword puzzle +street sign +traffic light +book jacket +menu +plate +guacamole +consomme +hot pot +trifle +ice cream +ice lolly +French loaf +bagel +pretzel +cheeseburger +hotdog +mashed potato +head cabbage +broccoli +cauliflower +zucchini +spaghetti squash +acorn squash +butternut squash +cucumber +artichoke +bell pepper +cardoon +mushroom +Granny Smith +strawberry +orange +lemon +fig +pineapple +banana +jackfruit +custard apple +pomegranate +hay +carbonara +chocolate sauce +dough +meat loaf +pizza +potpie +burrito +red wine +espresso +cup +eggnog +alp +bubble +cliff +coral reef +geyser +lakeside +promontory +sandbar +seashore +valley +volcano +ballplayer +groom +scuba diver +rapeseed +daisy +yellow lady's slipper +corn +acorn +hip +buckeye +coral fungus +agaric +gyromitra +stinkhorn +earthstar +hen-of-the-woods +bolete +ear +toilet tissue diff --git a/mediapipe/tasks/testdata/metadata/mobilenet_v2_1.0_224.json b/mediapipe/tasks/testdata/metadata/mobilenet_v2_1.0_224.json new file mode 100644 index 00000000..6f01f9f0 --- /dev/null +++ b/mediapipe/tasks/testdata/metadata/mobilenet_v2_1.0_224.json @@ -0,0 +1,82 @@ +{ + "name": "ImageClassifier", + "description": "Identify the most prominent object in the image from a known set of categories.", + "subgraph_metadata": [ + { + "input_tensor_metadata": [ + { + "name": "image", + "description": "Input image to be processed.", + "content": { + "content_properties_type": "ImageProperties", + "content_properties": { + "color_space": "RGB" + } + }, + "process_units": [ + { + "options_type": "NormalizationOptions", + "options": { + "mean": [ + 127.5 + ], + "std": [ + 127.5 + ] + } + } + ], + "stats": { + "max": [ + 1.0 + ], + "min": [ + -1.0 + ] + } + } + ], + "output_tensor_metadata": [ + { + "name": "score", + "description": "Score of the labels respectively.", + "content": { + "content_properties_type": "FeatureProperties", + "content_properties": { + } + }, + "process_units": [ + { + "options_type": "ScoreCalibrationOptions", + "options": { + "score_transformation": "LOG", + "default_score": 0.2 + } + } + ], + "stats": { + "max": [ + 1.0 + ], + "min": [ + 0.0 + ] + }, + "associated_files": [ + { + "name": "labels.txt", + "description": "Labels for categories that the model can recognize.", + "type": "TENSOR_AXIS_LABELS" + }, + { + "name": "score_calibration.txt", + "description": "Contains sigmoid-based score calibration parameters. The main purposes of score calibration is to make scores across classes comparable, so that a common threshold can be used for all output classes.", + "type": "TENSOR_AXIS_SCORE_CALIBRATION" + } + ] + } + ] + } + ], + "min_parser_version": "1.0.0" +} diff --git a/mediapipe/tasks/testdata/metadata/mobilenet_v2_1.0_224_quant.json b/mediapipe/tasks/testdata/metadata/mobilenet_v2_1.0_224_quant.json new file mode 100644 index 00000000..e2ba42e3 --- /dev/null +++ b/mediapipe/tasks/testdata/metadata/mobilenet_v2_1.0_224_quant.json @@ -0,0 +1,82 @@ +{ + "name": "ImageClassifier", + "description": "Identify the most prominent object in the image from a known set of categories.", + "subgraph_metadata": [ + { + "input_tensor_metadata": [ + { + "name": "image", + "description": "Input image to be processed.", + "content": { + "content_properties_type": "ImageProperties", + "content_properties": { + "color_space": "RGB" + } + }, + "process_units": [ + { + "options_type": "NormalizationOptions", + "options": { + "mean": [ + 127.5 + ], + "std": [ + 127.5 + ] + } + } + ], + "stats": { + "max": [ + 255.0 + ], + "min": [ + 0.0 + ] + } + } + ], + "output_tensor_metadata": [ + { + "name": "score", + "description": "Score of the labels respectively.", + "content": { + "content_properties_type": "FeatureProperties", + "content_properties": { + } + }, + "process_units": [ + { + "options_type": "ScoreCalibrationOptions", + "options": { + "score_transformation": "LOG", + "default_score": 0.2 + } + } + ], + "stats": { + "max": [ + 255.0 + ], + "min": [ + 0.0 + ] + }, + "associated_files": [ + { + "name": "labels.txt", + "description": "Labels for categories that the model can recognize.", + "type": "TENSOR_AXIS_LABELS" + }, + { + "name": "score_calibration.txt", + "description": "Contains sigmoid-based score calibration parameters. The main purposes of score calibration is to make scores across classes comparable, so that a common threshold can be used for all output classes.", + "type": "TENSOR_AXIS_SCORE_CALIBRATION" + } + ] + } + ] + } + ], + "min_parser_version": "1.0.0" +} diff --git a/mediapipe/tasks/testdata/metadata/score_thresholding_meta.json b/mediapipe/tasks/testdata/metadata/score_thresholding_meta.json new file mode 100644 index 00000000..d67a1aae --- /dev/null +++ b/mediapipe/tasks/testdata/metadata/score_thresholding_meta.json @@ -0,0 +1,14 @@ +{ + "subgraph_metadata": [ + { + "input_process_units": [ + { + "options_type": "ScoreThresholdingOptions", + "options": { + "global_score_threshold": 0.5 + } + } + ] + } + ] +} diff --git a/mediapipe/tasks/testdata/text/BUILD b/mediapipe/tasks/testdata/text/BUILD index 6cce5ae4..14999a03 100644 --- a/mediapipe/tasks/testdata/text/BUILD +++ b/mediapipe/tasks/testdata/text/BUILD @@ -76,9 +76,10 @@ filegroup( filegroup( name = "text_classifier_models", - srcs = glob([ - "test_model_text_classifier*.tflite", - ]), + srcs = [ + "test_model_text_classifier_bool_output.tflite", + "test_model_text_classifier_with_regex_tokenizer.tflite", + ], ) filegroup( diff --git a/mediapipe/tasks/testdata/vision/BUILD b/mediapipe/tasks/testdata/vision/BUILD index 0532458a..e23c4a66 100644 --- a/mediapipe/tasks/testdata/vision/BUILD +++ b/mediapipe/tasks/testdata/vision/BUILD @@ -28,6 +28,8 @@ mediapipe_files(srcs = [ "burger_rotated.jpg", "cat.jpg", "cat_mask.jpg", + "cat_rotated.jpg", + "cat_rotated_mask.jpg", "cats_and_dogs.jpg", "cats_and_dogs_no_resizing.jpg", "cats_and_dogs_rotated.jpg", @@ -35,10 +37,11 @@ mediapipe_files(srcs = [ "coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.tflite", "coco_ssd_mobilenet_v1_1.0_quant_2018_06_29_with_dummy_score_calibration.tflite", "deeplabv3.tflite", - "hand_landmark.task", + "fist.jpg", "hand_landmark_full.tflite", "hand_landmark_lite.tflite", "left_hands.jpg", + "left_hands_rotated.jpg", "mobilenet_v1_0.25_192_quantized_1_default_1.tflite", "mobilenet_v1_0.25_224_1_default_1.tflite", "mobilenet_v1_0.25_224_1_metadata_1.tflite", @@ -52,7 +55,9 @@ mediapipe_files(srcs = [ "multi_objects_rotated.jpg", "palm_detection_full.tflite", "pointing_up.jpg", + "pointing_up_rotated.jpg", "right_hands.jpg", + "right_hands_rotated.jpg", "segmentation_golden_rotation0.png", "segmentation_input_rotation0.jpg", "selfie_segm_128_128_3.tflite", @@ -60,14 +65,18 @@ mediapipe_files(srcs = [ "selfie_segm_144_256_3.tflite", "selfie_segm_144_256_3_expected_mask.jpg", "thumb_up.jpg", + "victory.jpg", ]) exports_files( srcs = [ "expected_left_down_hand_landmarks.prototxt", + "expected_left_down_hand_rotated_landmarks.prototxt", "expected_left_up_hand_landmarks.prototxt", + "expected_left_up_hand_rotated_landmarks.prototxt", "expected_right_down_hand_landmarks.prototxt", "expected_right_up_hand_landmarks.prototxt", + "gesture_recognizer.task", ], ) @@ -79,22 +88,29 @@ filegroup( "burger_rotated.jpg", "cat.jpg", "cat_mask.jpg", + "cat_rotated.jpg", + "cat_rotated_mask.jpg", "cats_and_dogs.jpg", "cats_and_dogs_no_resizing.jpg", "cats_and_dogs_rotated.jpg", + "fist.jpg", "hand_landmark_full.tflite", "hand_landmark_lite.tflite", "left_hands.jpg", + "left_hands_rotated.jpg", "mozart_square.jpg", "multi_objects.jpg", "multi_objects_rotated.jpg", "pointing_up.jpg", + "pointing_up_rotated.jpg", "right_hands.jpg", + "right_hands_rotated.jpg", "segmentation_golden_rotation0.png", "segmentation_input_rotation0.jpg", "selfie_segm_128_128_3_expected_mask.jpg", "selfie_segm_144_256_3_expected_mask.jpg", "thumb_up.jpg", + "victory.jpg", ], visibility = [ "//mediapipe/python:__subpackages__", @@ -110,9 +126,9 @@ filegroup( "coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.tflite", "coco_ssd_mobilenet_v1_1.0_quant_2018_06_29_with_dummy_score_calibration.tflite", "deeplabv3.tflite", - "hand_landmark.task", "hand_landmark_full.tflite", "hand_landmark_lite.tflite", + "hand_landmarker.task", "mobilenet_v1_0.25_192_quantized_1_default_1.tflite", "mobilenet_v1_0.25_224_1_default_1.tflite", "mobilenet_v1_0.25_224_1_metadata_1.tflite", @@ -127,16 +143,32 @@ filegroup( ], ) +# Gestures related models. Visible to model_maker. +filegroup( + name = "test_gesture_models", + visibility = [ + "//mediapipe/model_maker:__subpackages__", + "//mediapipe/tasks:internal", + ], +) + filegroup( name = "test_protos", srcs = [ "expected_left_down_hand_landmarks.prototxt", + "expected_left_down_hand_rotated_landmarks.prototxt", "expected_left_up_hand_landmarks.prototxt", + "expected_left_up_hand_rotated_landmarks.prototxt", "expected_right_down_hand_landmarks.prototxt", "expected_right_up_hand_landmarks.prototxt", + "fist_landmarks.pbtxt", "hand_detector_result_one_hand.pbtxt", + "hand_detector_result_one_hand_rotated.pbtxt", "hand_detector_result_two_hands.pbtxt", "pointing_up_landmarks.pbtxt", + "pointing_up_rotated_landmarks.pbtxt", "thumb_up_landmarks.pbtxt", + "thumb_up_rotated_landmarks.pbtxt", + "victory_landmarks.pbtxt", ], ) diff --git a/mediapipe/tasks/testdata/vision/expected_left_down_hand_rotated_landmarks.prototxt b/mediapipe/tasks/testdata/vision/expected_left_down_hand_rotated_landmarks.prototxt new file mode 100644 index 00000000..3cbf8804 --- /dev/null +++ b/mediapipe/tasks/testdata/vision/expected_left_down_hand_rotated_landmarks.prototxt @@ -0,0 +1,84 @@ +landmark { + x: 0.9259716 + y: 0.18969846 +} +landmark { + x: 0.88135517 + y: 0.28856543 +} +landmark { + x: 0.7600651 + y: 0.3578236 +} +landmark { + x: 0.62631166 + y: 0.40490413 +} +landmark { + x: 0.5374573 + y: 0.45170194 +} +landmark { + x: 0.57372385 + y: 0.29924914 +} +landmark { + x: 0.36731184 + y: 0.33081773 +} +landmark { + x: 0.24132833 + y: 0.34759054 +} +landmark { + x: 0.13690609 + y: 0.35727677 +} +landmark { + x: 0.5535803 + y: 0.2398035 +} +landmark { + x: 0.31834763 + y: 0.24999242 +} +landmark { + x: 0.16748133 + y: 0.25625145 +} +landmark { + x: 0.050747424 + y: 0.25991398 +} +landmark { + x: 0.56593156 + y: 0.1867483 +} +landmark { + x: 0.3543046 + y: 0.17923892 +} +landmark { + x: 0.21360746 + y: 0.17454882 +} +landmark { + x: 0.11110917 + y: 0.17232567 +} +landmark { + x: 0.5948908 + y: 0.14024714 +} +landmark { + x: 0.42692152 + y: 0.11949824 +} +landmark { + x: 0.32239118 + y: 0.106370345 +} +landmark { + x: 0.23672739 + y: 0.09432885 +} diff --git a/mediapipe/tasks/testdata/vision/expected_left_up_hand_rotated_landmarks.prototxt b/mediapipe/tasks/testdata/vision/expected_left_up_hand_rotated_landmarks.prototxt new file mode 100644 index 00000000..42eccbcc --- /dev/null +++ b/mediapipe/tasks/testdata/vision/expected_left_up_hand_rotated_landmarks.prototxt @@ -0,0 +1,84 @@ +landmark { + x: 0.06676084 + y: 0.8095678 +} +landmark { + x: 0.11359626 + y: 0.71148247 +} +landmark { + x: 0.23572624 + y: 0.6414506 +} +landmark { + x: 0.37323278 + y: 0.5959156 +} +landmark { + x: 0.46243322 + y: 0.55125874 +} +landmark { + x: 0.4205411 + y: 0.69531494 +} +landmark { + x: 0.62798893 + y: 0.66715276 +} +landmark { + x: 0.7568023 + y: 0.65208924 +} +landmark { + x: 0.86370826 + y: 0.6437276 +} +landmark { + x: 0.445136 + y: 0.75394773 +} +landmark { + x: 0.6787485 + y: 0.745853 +} +landmark { + x: 0.8290694 + y: 0.7412988 +} +landmark { + x: 0.94454145 + y: 0.7384017 +} +landmark { + x: 0.43516788 + y: 0.8082166 +} +landmark { + x: 0.6459554 + y: 0.81768996 +} +landmark { + x: 0.7875173 + y: 0.825062 +} +landmark { + x: 0.89249825 + y: 0.82850707 +} +landmark { + x: 0.40665048 + y: 0.8567925 +} +landmark { + x: 0.57228816 + y: 0.8802181 +} +landmark { + x: 0.6762071 + y: 0.8941581 +} +landmark { + x: 0.76453924 + y: 0.90583205 +} diff --git a/mediapipe/tasks/testdata/vision/fist_landmarks.pbtxt b/mediapipe/tasks/testdata/vision/fist_landmarks.pbtxt new file mode 100644 index 00000000..a24358c3 --- /dev/null +++ b/mediapipe/tasks/testdata/vision/fist_landmarks.pbtxt @@ -0,0 +1,223 @@ +classifications { + classification { + score: 1.0 + label: "Left" + display_name: "Left" + } +} + +landmarks { + landmark { + x: 0.47709703 + y: 0.66129065 + z: -3.3540672e-07 + } + landmark { + x: 0.6125982 + y: 0.5578249 + z: -0.041392017 + } + landmark { + x: 0.71123487 + y: 0.4316616 + z: -0.064544134 + } + landmark { + x: 0.6836403 + y: 0.3199585 + z: -0.08752567 + } + landmark { + x: 0.5593274 + y: 0.3206453 + z: -0.09880819 + } + landmark { + x: 0.60828537 + y: 0.3068749 + z: -0.014799656 + } + landmark { + x: 0.62940764 + y: 0.21414441 + z: -0.06007311 + } + landmark { + x: 0.6244353 + y: 0.32872596 + z: -0.08326768 + } + landmark { + x: 0.60784453 + y: 0.3684796 + z: -0.09658983 + } + landmark { + x: 0.5156504 + y: 0.32194698 + z: -0.021699267 + } + landmark { + x: 0.52931 + y: 0.24767634 + z: -0.062571 + } + landmark { + x: 0.5484773 + y: 0.3805329 + z: -0.07028895 + } + landmark { + x: 0.54428184 + y: 0.3881125 + z: -0.07458326 + } + landmark { + x: 0.43159598 + y: 0.34918433 + z: -0.037482508 + } + landmark { + x: 0.4486106 + y: 0.27649382 + z: -0.08174769 + } + landmark { + x: 0.47723144 + y: 0.3964985 + z: -0.06496752 + } + landmark { + x: 0.46794242 + y: 0.4082967 + z: -0.04897496 + } + landmark { + x: 0.34826216 + y: 0.37813392 + z: -0.057438444 + } + landmark { + x: 0.3861837 + y: 0.32820183 + z: -0.07282783 + } + landmark { + x: 0.41143674 + y: 0.39734486 + z: -0.047633167 + } + landmark { + x: 0.39401984 + y: 0.41149133 + z: -0.029640475 + } +} + +world_landmarks { + landmark { + x: -0.008604452 + y: 0.08165767 + z: 0.0061365655 + } + landmark { + x: 0.027301773 + y: 0.061905317 + z: -0.00872007 + } + landmark { + x: 0.049898714 + y: 0.035359327 + z: -0.016682662 + } + landmark { + x: 0.050297678 + y: 0.005200807 + z: -0.028928496 + } + landmark { + x: 0.015639625 + y: -0.0063155442 + z: -0.03174634 + } + landmark { + x: 0.029161729 + y: -0.0024596984 + z: 0.0011553494 + } + landmark { + x: 0.034491 + y: -0.017581237 + z: -0.020781275 + } + landmark { + x: 0.034020264 + y: -0.0059247985 + z: -0.02573838 + } + landmark { + x: 0.02867364 + y: 0.011137734 + z: -0.009430941 + } + landmark { + x: 0.0015385814 + y: -0.004778851 + z: 0.0056454404 + } + landmark { + x: 0.010490709 + y: -0.019680617 + z: -0.027034117 + } + landmark { + x: 0.0132071925 + y: 0.0071370844 + z: -0.034802448 + } + landmark { + x: 0.0139978565 + y: 0.011672501 + z: -0.0040006908 + } + landmark { + x: -0.019919239 + y: -0.0006897822 + z: -0.0003317799 + } + landmark { + x: -0.01088193 + y: -0.008502296 + z: -0.02873486 + } + landmark { + x: -0.005327127 + y: 0.012745364 + z: -0.034153957 + } + landmark { + x: -0.0027040644 + y: 0.02167169 + z: -0.011669062 + } + landmark { + x: -0.038813893 + y: 0.011925209 + z: -0.0076287366 + } + landmark { + x: -0.030842202 + y: 0.0010964936 + z: -0.022697516 + } + landmark { + x: -0.01829514 + y: 0.013929318 + z: -0.032819964 + } + landmark { + x: -0.024175374 + y: 0.022456694 + z: -0.02357186 + } +} diff --git a/mediapipe/tasks/testdata/vision/gesture_recognizer_with_custom_classifier.task b/mediapipe/tasks/testdata/vision/gesture_recognizer_with_custom_classifier.task new file mode 100644 index 00000000..3c1da7b3 Binary files /dev/null and b/mediapipe/tasks/testdata/vision/gesture_recognizer_with_custom_classifier.task differ diff --git a/mediapipe/tasks/testdata/vision/hand_detector_result_one_hand_rotated.pbtxt b/mediapipe/tasks/testdata/vision/hand_detector_result_one_hand_rotated.pbtxt new file mode 100644 index 00000000..cec4d616 --- /dev/null +++ b/mediapipe/tasks/testdata/vision/hand_detector_result_one_hand_rotated.pbtxt @@ -0,0 +1,33 @@ +detections { + label: "Palm" + score: 0.97115 + location_data { + format: RELATIVE_BOUNDING_BOX + relative_bounding_box { + xmin: 0.5198178 + ymin: 0.6467485 + width: 0.42467535 + height: 0.22546273 + } + } +} +detections { + label: "Palm" + score: 0.96701413 + location_data { + format: RELATIVE_BOUNDING_BOX + relative_bounding_box { + xmin: 0.024490356 + ymin: 0.12620124 + width: 0.43832153 + height: 0.23269764 + } + } +} +hand_rects { + x_center: 0.5760683 + y_center: 0.6829921 + height: 0.5862031 + width: 1.1048855 + rotation: -0.8250832 +} diff --git a/mediapipe/tasks/testdata/vision/hand_gesture_recognizer_with_custom_classifier.task b/mediapipe/tasks/testdata/vision/hand_gesture_recognizer_with_custom_classifier.task new file mode 100644 index 00000000..1390ca88 Binary files /dev/null and b/mediapipe/tasks/testdata/vision/hand_gesture_recognizer_with_custom_classifier.task differ diff --git a/mediapipe/tasks/testdata/vision/hand_landmark.task b/mediapipe/tasks/testdata/vision/hand_landmarker.task similarity index 99% rename from mediapipe/tasks/testdata/vision/hand_landmark.task rename to mediapipe/tasks/testdata/vision/hand_landmarker.task index b6eedf32..1ae9f7f6 100644 Binary files a/mediapipe/tasks/testdata/vision/hand_landmark.task and b/mediapipe/tasks/testdata/vision/hand_landmarker.task differ diff --git a/mediapipe/tasks/testdata/vision/pointing_up_rotated_landmarks.pbtxt b/mediapipe/tasks/testdata/vision/pointing_up_rotated_landmarks.pbtxt new file mode 100644 index 00000000..65bb11bc --- /dev/null +++ b/mediapipe/tasks/testdata/vision/pointing_up_rotated_landmarks.pbtxt @@ -0,0 +1,223 @@ +classifications { + classification { + score: 1.0 + label: "Left" + display_name: "Left" + } +} + +landmarks { + landmark { + x: 0.25546086 + y: 0.47584262 + z: 1.835341e-07 + } + landmark { + x: 0.3363011 + y: 0.54135 + z: -0.041144375 + } + landmark { + x: 0.4375146 + y: 0.57881975 + z: -0.06807727 + } + landmark { + x: 0.49603376 + y: 0.5263966 + z: -0.09387612 + } + landmark { + x: 0.5022822 + y: 0.4413827 + z: -0.1189948 + } + landmark { + x: 0.5569452 + y: 0.4724485 + z: -0.05138246 + } + landmark { + x: 0.6687125 + y: 0.47918057 + z: -0.09121969 + } + landmark { + x: 0.73666537 + y: 0.48318353 + z: -0.11703273 + } + landmark { + x: 0.7998315 + y: 0.4741413 + z: -0.1386424 + } + landmark { + x: 0.5244063 + y: 0.39292705 + z: -0.061040796 + } + landmark { + x: 0.57215345 + y: 0.41514704 + z: -0.11967233 + } + landmark { + x: 0.4724468 + y: 0.45553637 + z: -0.13287684 + } + landmark { + x: 0.43794966 + y: 0.45210314 + z: -0.13210714 + } + landmark { + x: 0.47838163 + y: 0.33329 + z: -0.07421263 + } + landmark { + x: 0.51081127 + y: 0.35479474 + z: -0.13596693 + } + landmark { + x: 0.42433846 + y: 0.40486792 + z: -0.121291734 + } + landmark { + x: 0.40280548 + y: 0.39977497 + z: -0.09928809 + } + landmark { + x: 0.42269367 + y: 0.2798249 + z: -0.09064263 + } + landmark { + x: 0.45849988 + y: 0.3069861 + z: -0.12894689 + } + landmark { + x: 0.40754712 + y: 0.35153976 + z: -0.109160855 + } + landmark { + x: 0.38855004 + y: 0.3467068 + z: -0.08820164 + } +} + +world_landmarks { + landmark { + x: -0.08568013 + y: 0.016593203 + z: 0.036527164 + } + landmark { + x: -0.0565372 + y: 0.041761592 + z: 0.019493781 + } + landmark { + x: -0.031365488 + y: 0.05031186 + z: 0.0025481891 + } + landmark { + x: -0.008534161 + y: 0.04286737 + z: -0.024755282 + } + landmark { + x: -0.0047254 + y: 0.015748458 + z: -0.035581928 + } + landmark { + x: 0.013083893 + y: 0.024668094 + z: 0.0035934823 + } + landmark { + x: 0.04149521 + y: 0.024621274 + z: -0.0030611698 + } + landmark { + x: 0.06257473 + y: 0.025388625 + z: -0.010340984 + } + landmark { + x: 0.08009179 + y: 0.023082614 + z: -0.03162942 + } + landmark { + x: 0.006135068 + y: 0.000696786 + z: 0.0048212176 + } + landmark { + x: 0.01678449 + y: 0.0067061195 + z: -0.029920919 + } + landmark { + x: -0.008948593 + y: 0.016808286 + z: -0.03755109 + } + landmark { + x: -0.01789449 + y: 0.0153161455 + z: -0.012059977 + } + landmark { + x: -0.0061980113 + y: -0.017872887 + z: -0.002366997 + } + landmark { + x: -0.004643807 + y: -0.0108282855 + z: -0.034515083 + } + landmark { + x: -0.027603384 + y: 0.003529715 + z: -0.033665676 + } + landmark { + x: -0.035679806 + y: 0.0038255951 + z: -0.008094264 + } + landmark { + x: -0.02957782 + y: -0.031701155 + z: -0.008180461 + } + landmark { + x: -0.020741666 + y: -0.02506058 + z: -0.026839724 + } + landmark { + x: -0.0310834 + y: -0.009496164 + z: -0.032422185 + } + landmark { + x: -0.037420202 + y: -0.012883307 + z: -0.017971724 + } +} diff --git a/mediapipe/tasks/testdata/vision/thumb_up_rotated_landmarks.pbtxt b/mediapipe/tasks/testdata/vision/thumb_up_rotated_landmarks.pbtxt new file mode 100644 index 00000000..3636e2e4 --- /dev/null +++ b/mediapipe/tasks/testdata/vision/thumb_up_rotated_landmarks.pbtxt @@ -0,0 +1,223 @@ +classifications { + classification { + score: 1.0 + label: "Left" + display_name: "Left" + } +} + +landmarks { + landmark { + x: 0.3283601 + y: 0.63773525 + z: -3.2280354e-07 + } + landmark { + x: 0.46280807 + y: 0.6339767 + z: -0.06408348 + } + landmark { + x: 0.5831279 + y: 0.57430106 + z: -0.08583106 + } + landmark { + x: 0.6689471 + y: 0.49959752 + z: -0.09886064 + } + landmark { + x: 0.74378216 + y: 0.47357544 + z: -0.09680563 + } + landmark { + x: 0.5233122 + y: 0.41020474 + z: -0.038088404 + } + landmark { + x: 0.5296913 + y: 0.3372598 + z: -0.08874837 + } + landmark { + x: 0.49039274 + y: 0.43994758 + z: -0.102315836 + } + landmark { + x: 0.4824569 + y: 0.47969607 + z: -0.1030014 + } + landmark { + x: 0.4451338 + y: 0.39520803 + z: -0.02177739 + } + landmark { + x: 0.4410001 + y: 0.34107083 + z: -0.07294245 + } + landmark { + x: 0.4162798 + y: 0.46102384 + z: -0.07746907 + } + landmark { + x: 0.43492994 + y: 0.47154287 + z: -0.07404131 + } + landmark { + x: 0.37671578 + y: 0.39535576 + z: -0.016277775 + } + landmark { + x: 0.36978847 + y: 0.34265152 + z: -0.07346253 + } + landmark { + x: 0.3559884 + y: 0.44905427 + z: -0.057693005 + } + landmark { + x: 0.37711847 + y: 0.46414754 + z: -0.03662908 + } + landmark { + x: 0.3142985 + y: 0.3942253 + z: -0.0152847925 + } + landmark { + x: 0.30000874 + y: 0.35543376 + z: -0.046002634 + } + landmark { + x: 0.30002704 + y: 0.42357764 + z: -0.032671776 + } + landmark { + x: 0.31079838 + y: 0.44218025 + z: -0.016200554 + } +} + +world_landmarks { + landmark { + x: -0.030687196 + y: 0.0678545 + z: 0.051061403 + } + landmark { + x: 0.0047719833 + y: 0.06330968 + z: 0.018945374 + } + landmark { + x: 0.039799504 + y: 0.054109577 + z: 0.007930638 + } + landmark { + x: 0.069374144 + y: 0.035063196 + z: 2.2522348e-05 + } + landmark { + x: 0.087818466 + y: 0.018390425 + z: 0.004055788 + } + landmark { + x: 0.02810654 + y: 0.0043561812 + z: -0.0038672548 + } + landmark { + x: 0.025270049 + y: -0.0039896416 + z: -0.032991238 + } + landmark { + x: 0.020414166 + y: 0.006768506 + z: -0.032724563 + } + landmark { + x: 0.016415983 + y: 0.024563588 + z: -0.0058115427 + } + landmark { + x: 0.0038743173 + y: -0.0044466974 + z: 0.0024876352 + } + landmark { + x: 0.0041790796 + y: -0.0115309935 + z: -0.03532454 + } + landmark { + x: -0.0016900161 + y: 0.015519895 + z: -0.03596156 + } + landmark { + x: 0.004309217 + y: 0.01917039 + z: 0.003907912 + } + landmark { + x: -0.016969737 + y: -0.005584497 + z: 0.0034258277 + } + landmark { + x: -0.016737012 + y: -0.01159037 + z: -0.02876696 + } + landmark { + x: -0.018165365 + y: 0.01376111 + z: -0.026835402 + } + landmark { + x: -0.012430167 + y: 0.02064222 + z: -0.00087265146 + } + landmark { + x: -0.043247573 + y: 0.0011161827 + z: 0.0056269006 + } + landmark { + x: -0.038128495 + y: -0.011477032 + z: -0.016374081 + } + landmark { + x: -0.034920715 + y: 0.005510211 + z: -0.029714659 + } + landmark { + x: -0.03815982 + y: 0.011989757 + z: -0.014853194 + } +} diff --git a/mediapipe/tasks/testdata/vision/victory_landmarks.pbtxt b/mediapipe/tasks/testdata/vision/victory_landmarks.pbtxt new file mode 100644 index 00000000..7a704ee3 --- /dev/null +++ b/mediapipe/tasks/testdata/vision/victory_landmarks.pbtxt @@ -0,0 +1,223 @@ +classifications { + classification { + score: 1.0 + label: "Left" + display_name: "Left" + } +} + +landmarks { + landmark { + x: 0.5164316 + y: 0.804093 + z: 8.7653416e-07 + } + landmark { + x: 0.6063608 + y: 0.7111354 + z: -0.044089418 + } + landmark { + x: 0.6280186 + y: 0.588498 + z: -0.062358405 + } + landmark { + x: 0.5265348 + y: 0.52083343 + z: -0.08526791 + } + landmark { + x: 0.4243384 + y: 0.4993468 + z: -0.1077741 + } + landmark { + x: 0.5605667 + y: 0.4489705 + z: -0.016151091 + } + landmark { + x: 0.5766643 + y: 0.32260323 + z: -0.049342215 + } + landmark { + x: 0.5795845 + y: 0.24180722 + z: -0.07323826 + } + landmark { + x: 0.5827511 + y: 0.16940045 + z: -0.09069163 + } + landmark { + x: 0.4696163 + y: 0.4599558 + z: -0.032168437 + } + landmark { + x: 0.44361597 + y: 0.31689578 + z: -0.075698614 + } + landmark { + x: 0.42695498 + y: 0.22273324 + z: -0.10819675 + } + landmark { + x: 0.40697217 + y: 0.14279765 + z: -0.12666894 + } + landmark { + x: 0.39543492 + y: 0.50612336 + z: -0.055138163 + } + landmark { + x: 0.3618012 + y: 0.4388296 + z: -0.1298119 + } + landmark { + x: 0.4154368 + y: 0.52674913 + z: -0.1463017 + } + landmark { + x: 0.44916254 + y: 0.59442246 + z: -0.13470782 + } + landmark { + x: 0.33178204 + y: 0.5731769 + z: -0.08103096 + } + landmark { + x: 0.3092102 + y: 0.5040002 + z: -0.13258384 + } + landmark { + x: 0.35576707 + y: 0.5576498 + z: -0.12714732 + } + landmark { + x: 0.393444 + y: 0.6118667 + z: -0.11102459 + } +} + +world_landmarks { + landmark { + x: 0.01299962 + y: 0.09162361 + z: 0.011185312 + } + landmark { + x: 0.03726317 + y: 0.0638103 + z: -0.010005756 + } + landmark { + x: 0.03975261 + y: 0.03712649 + z: -0.02906275 + } + landmark { + x: 0.018798776 + y: 0.012429599 + z: -0.048737116 + } + landmark { + x: -0.0128555335 + y: 0.001022811 + z: -0.044505004 + } + landmark { + x: 0.025658218 + y: -0.008031519 + z: -0.0058278795 + } + landmark { + x: 0.028017294 + y: -0.038120236 + z: -0.010376478 + } + landmark { + x: 0.030067094 + y: -0.059907563 + z: -0.014568218 + } + landmark { + x: 0.027284538 + y: -0.07803874 + z: -0.032692235 + } + landmark { + x: 0.0013260426 + y: -0.005039873 + z: 0.005567288 + } + landmark { + x: -0.002380834 + y: -0.044605374 + z: -0.0038231965 + } + landmark { + x: -0.009240147 + y: -0.066279344 + z: -0.02161214 + } + landmark { + x: -0.0092535615 + y: -0.08933755 + z: -0.037401434 + } + landmark { + x: -0.01751284 + y: 0.0037118336 + z: 0.0047480655 + } + landmark { + x: -0.02195602 + y: -0.010006189 + z: -0.02371484 + } + landmark { + x: -0.012851426 + y: 0.008346066 + z: -0.037721373 + } + landmark { + x: -0.00018795021 + y: 0.026816685 + z: -0.03732748 + } + landmark { + x: -0.034864448 + y: 0.022316 + z: -0.0002774651 + } + landmark { + x: -0.035896845 + y: 0.01066218 + z: -0.017325373 + } + landmark { + x: -0.02358637 + y: 0.018667895 + z: -0.028403495 + } + landmark { + x: -0.013704676 + y: 0.033456434 + z: -0.02595728 + } +} diff --git a/mediapipe/util/resource_util_android.cc b/mediapipe/util/resource_util_android.cc index b18354d5..1e970f21 100644 --- a/mediapipe/util/resource_util_android.cc +++ b/mediapipe/util/resource_util_android.cc @@ -82,7 +82,8 @@ absl::StatusOr PathToResourceAsFile(const std::string& path) { // If that fails, assume it was a relative path, and try just the base name. { const size_t last_slash_idx = path.find_last_of("\\/"); - CHECK_NE(last_slash_idx, std::string::npos); // Make sure it's a path. + RET_CHECK(last_slash_idx != std::string::npos) + << path << " doesn't have a slash in it"; // Make sure it's a path. auto base_name = path.substr(last_slash_idx + 1); auto status_or_path = PathToResourceAsFileInternal(base_name); if (status_or_path.ok()) { diff --git a/mediapipe/util/resource_util_apple.cc b/mediapipe/util/resource_util_apple.cc index c812dcb5..f6471834 100644 --- a/mediapipe/util/resource_util_apple.cc +++ b/mediapipe/util/resource_util_apple.cc @@ -71,7 +71,8 @@ absl::StatusOr PathToResourceAsFile(const std::string& path) { // If that fails, assume it was a relative path, and try just the base name. { const size_t last_slash_idx = path.find_last_of("\\/"); - CHECK_NE(last_slash_idx, std::string::npos); // Make sure it's a path. + RET_CHECK(last_slash_idx != std::string::npos) + << path << " doesn't have a slash in it"; // Make sure it's a path. auto base_name = path.substr(last_slash_idx + 1); auto status_or_path = PathToResourceAsFileInternal(base_name); if (status_or_path.ok()) { diff --git a/mediapipe/util/tflite/BUILD b/mediapipe/util/tflite/BUILD index 9d37b60a..e9b8bfa0 100644 --- a/mediapipe/util/tflite/BUILD +++ b/mediapipe/util/tflite/BUILD @@ -84,33 +84,28 @@ cc_library( "//conditions:default": ["tflite_gpu_runner.h"], }), deps = select({ - "//mediapipe:ios": [], - "//mediapipe:macos": [], - "//conditions:default": [ - "@com_google_absl//absl/strings", - "//mediapipe/framework/port:ret_check", - "//mediapipe/framework/port:status", - "//mediapipe/framework/port:statusor", - "@org_tensorflow//tensorflow/lite:framework", - "@org_tensorflow//tensorflow/lite/delegates/gpu:api", - "@org_tensorflow//tensorflow/lite/delegates/gpu/common:model", - "@org_tensorflow//tensorflow/lite/delegates/gpu/common:model_builder", - "@org_tensorflow//tensorflow/lite/delegates/gpu/gl:api2", - ], - "//mediapipe:android": [ - "@com_google_absl//absl/strings", - "//mediapipe/framework/port:ret_check", - "//mediapipe/framework/port:status", - "//mediapipe/framework/port:statusor", - "@org_tensorflow//tensorflow/lite:framework", - "@org_tensorflow//tensorflow/lite/delegates/gpu:api", - "@org_tensorflow//tensorflow/lite/delegates/gpu/cl:api", - "@org_tensorflow//tensorflow/lite/delegates/gpu/common:model", - "@org_tensorflow//tensorflow/lite/delegates/gpu/common:model_builder", - "@org_tensorflow//tensorflow/lite/delegates/gpu/gl:api2", - ], - }) + [ + "//mediapipe:ios": [], + "//mediapipe:macos": [], + "//conditions:default": [ + "@com_google_absl//absl/strings", + "//mediapipe/framework/port:ret_check", + "//mediapipe/framework/port:status", + "//mediapipe/framework/port:statusor", + "@org_tensorflow//tensorflow/lite:framework", + "@org_tensorflow//tensorflow/lite/delegates/gpu:api", + "@org_tensorflow//tensorflow/lite/delegates/gpu/common:model", + "@org_tensorflow//tensorflow/lite/delegates/gpu/common:model_builder", + "@org_tensorflow//tensorflow/lite/delegates/gpu/gl:api2", + ], + }) + + select({ + "//mediapipe:android": [ + "@org_tensorflow//tensorflow/lite/delegates/gpu/cl:api", + ], + "//conditions:default": [], + }) + [ "@com_google_absl//absl/status", + "//mediapipe/framework:port", "@org_tensorflow//tensorflow/lite/core/api", ], ) diff --git a/mediapipe/util/tflite/tflite_gpu_runner.cc b/mediapipe/util/tflite/tflite_gpu_runner.cc index 4c422835..4e40975c 100644 --- a/mediapipe/util/tflite/tflite_gpu_runner.cc +++ b/mediapipe/util/tflite/tflite_gpu_runner.cc @@ -34,7 +34,7 @@ // This code should be enabled as soon as TensorFlow version, which mediapipe // uses, will include this module. -#ifdef __ANDROID__ +#if defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) #include "tensorflow/lite/delegates/gpu/cl/api.h" #endif @@ -82,7 +82,7 @@ ObjectDef GetSSBOObjectDef(int channels) { return gpu_object_def; } -#ifdef __ANDROID__ +#if defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) cl::InferenceOptions GetClInferenceOptions(const InferenceOptions& options) { cl::InferenceOptions result{}; @@ -106,7 +106,7 @@ absl::Status VerifyShapes(const std::vector& actual, return absl::OkStatus(); } -#endif // __ANDROID__ +#endif // defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) } // namespace @@ -225,7 +225,7 @@ absl::Status TFLiteGPURunner::InitializeOpenGL( absl::Status TFLiteGPURunner::InitializeOpenCL( std::unique_ptr* builder) { -#ifdef __ANDROID__ +#if defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) cl::InferenceEnvironmentOptions env_options; if (!serialized_binary_cache_.empty()) { env_options.serialized_binary_cache = serialized_binary_cache_; @@ -254,11 +254,12 @@ absl::Status TFLiteGPURunner::InitializeOpenCL( return absl::OkStatus(); #else - return mediapipe::UnimplementedError("Currently only Android is supported"); -#endif // __ANDROID__ + return mediapipe::UnimplementedError( + "Currently only Android & ChromeOS are supported"); +#endif // defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) } -#ifdef __ANDROID__ +#if defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) absl::Status TFLiteGPURunner::InitializeOpenCLFromSerializedModel( std::unique_ptr* builder) { @@ -283,7 +284,7 @@ absl::StatusOr> TFLiteGPURunner::GetSerializedModel() { return serialized_model; } -#endif // __ANDROID__ +#endif // defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) } // namespace gpu } // namespace tflite diff --git a/mediapipe/util/tflite/tflite_gpu_runner.h b/mediapipe/util/tflite/tflite_gpu_runner.h index 88d3914f..dfbc8d65 100644 --- a/mediapipe/util/tflite/tflite_gpu_runner.h +++ b/mediapipe/util/tflite/tflite_gpu_runner.h @@ -20,6 +20,7 @@ #include #include "absl/status/status.h" +#include "mediapipe/framework/port.h" #include "mediapipe/framework/port/status.h" #include "mediapipe/framework/port/statusor.h" #include "tensorflow/lite/core/api/op_resolver.h" @@ -28,9 +29,9 @@ #include "tensorflow/lite/delegates/gpu/gl/api2.h" #include "tensorflow/lite/model.h" -#ifdef __ANDROID__ +#if defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) #include "tensorflow/lite/delegates/gpu/cl/api.h" -#endif // __ANDROID__ +#endif // defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) namespace tflite { namespace gpu { @@ -83,7 +84,7 @@ class TFLiteGPURunner { return output_shape_from_model_; } -#ifdef __ANDROID__ +#if defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) void SetSerializedBinaryCache(std::vector&& cache) { serialized_binary_cache_ = std::move(cache); } @@ -98,26 +99,26 @@ class TFLiteGPURunner { } absl::StatusOr> GetSerializedModel(); -#endif // __ANDROID__ +#endif // defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) private: absl::Status InitializeOpenGL(std::unique_ptr* builder); absl::Status InitializeOpenCL(std::unique_ptr* builder); -#ifdef __ANDROID__ +#if defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) absl::Status InitializeOpenCLFromSerializedModel( std::unique_ptr* builder); -#endif // __ANDROID__ +#endif // defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) InferenceOptions options_; std::unique_ptr gl_environment_; -#ifdef __ANDROID__ +#if defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) std::unique_ptr cl_environment_; std::vector serialized_binary_cache_; std::vector serialized_model_; bool serialized_model_used_ = false; -#endif // __ANDROID__ +#endif // defined(__ANDROID__) || defined(MEDIAPIPE_CHROMIUMOS) // graph_gl_ is maintained temporarily and becomes invalid after runner_ is // ready diff --git a/mediapipe/util/time_series_util.h b/mediapipe/util/time_series_util.h index a6a5911a..afa66acc 100644 --- a/mediapipe/util/time_series_util.h +++ b/mediapipe/util/time_series_util.h @@ -58,7 +58,7 @@ absl::Status FillTimeSeriesHeaderIfValid(const Packet& header_packet, absl::Status FillMultiStreamTimeSeriesHeaderIfValid( const Packet& header_packet, MultiStreamTimeSeriesHeader* header); -// Returnsabsl::Status::OK iff options contains an extension of type +// Returns absl::Status::OK iff options contains an extension of type // OptionsClass. template absl::Status HasOptionsExtension(const CalculatorOptions& options) { @@ -75,7 +75,7 @@ absl::Status HasOptionsExtension(const CalculatorOptions& options) { return absl::InvalidArgumentError(error_message); } -// Returnsabsl::Status::OK if the shape of 'matrix' is consistent +// Returns absl::Status::OK if the shape of 'matrix' is consistent // with the num_samples and num_channels fields present in 'header'. // The corresponding matrix dimensions of unset header fields are // ignored, so e.g. an empty header (which is not valid according to diff --git a/mediapipe/util/tracking/tracking.cc b/mediapipe/util/tracking/tracking.cc index 7e80cd5c..88ba3980 100644 --- a/mediapipe/util/tracking/tracking.cc +++ b/mediapipe/util/tracking/tracking.cc @@ -1323,10 +1323,9 @@ void MotionBox::GetSpatialGaussWeights(const MotionBoxState& box_state, const float space_sigma_x = std::max( options_.spatial_sigma(), box_state.inlier_width() * inv_box_domain.x() * 0.5f * box_state.prior_weight() / 1.65f); - const float space_sigma_y = options_.spatial_sigma(); - std::max(options_.spatial_sigma(), box_state.inlier_height() * - inv_box_domain.y() * 0.5f * - box_state.prior_weight() / 1.65f); + const float space_sigma_y = std::max( + options_.spatial_sigma(), box_state.inlier_height() * inv_box_domain.y() * + 0.5f * box_state.prior_weight() / 1.65f); *spatial_gauss_x = -0.5f / (space_sigma_x * space_sigma_x); *spatial_gauss_y = -0.5f / (space_sigma_y * space_sigma_y); diff --git a/setup.py b/setup.py index cb304ba7..cda53b1c 100644 --- a/setup.py +++ b/setup.py @@ -121,8 +121,10 @@ def _add_mp_init_files(): mp_dir_init_file = open(MP_DIR_INIT_PY, 'a') mp_dir_init_file.writelines([ '\n', 'from mediapipe.python import *\n', - 'import mediapipe.python.solutions as solutions', '\n\n', - '__version__ = \'{}\''.format(__version__), '\n' + 'import mediapipe.python.solutions as solutions \n', + 'import mediapipe.tasks.python as tasks\n', '\n\n', 'del framework\n', + 'del gpu\n', 'del modules\n', 'del python\n', 'del mediapipe\n', + 'del util\n', '__version__ = \'{}\''.format(__version__), '\n' ]) mp_dir_init_file.close() diff --git a/third_party/com_google_sentencepiece_no_gflag_no_gtest.diff b/third_party/com_google_sentencepiece_no_gflag_no_gtest.diff new file mode 100644 index 00000000..a084d926 --- /dev/null +++ b/third_party/com_google_sentencepiece_no_gflag_no_gtest.diff @@ -0,0 +1,34 @@ +diff --git a/src/BUILD b/src/BUILD +index b4298d2..f3877a3 100644 +--- a/src/BUILD ++++ b/src/BUILD +@@ -71,9 +71,7 @@ cc_library( + ":common", + ":sentencepiece_cc_proto", + ":sentencepiece_model_cc_proto", +- "@com_github_gflags_gflags//:gflags", + "@com_google_glog//:glog", +- "@com_google_googletest//:gtest", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/container:flat_hash_map", +diff --git a/src/normalizer.h b/src/normalizer.h +index c16ac16..2af58be 100644 +--- a/src/normalizer.h ++++ b/src/normalizer.h +@@ -21,7 +21,6 @@ + #include + #include + +-#include "gtest/gtest_prod.h" + #include "absl/strings/string_view.h" + #include "third_party/darts_clone/include/darts.h" + #include "src/common.h" +@@ -97,7 +96,6 @@ class Normalizer { + friend class Builder; + + private: +- FRIEND_TEST(NormalizerTest, EncodeDecodePrecompiledCharsMapTest); + + void Init(); + diff --git a/third_party/external_files.bzl b/third_party/external_files.bzl index b85d9331..e47dc981 100644 --- a/third_party/external_files.bzl +++ b/third_party/external_files.bzl @@ -31,7 +31,7 @@ def external_files(): http_file( name = "com_google_mediapipe_bert_text_classifier_tflite", sha256 = "1e5a550c09bff0a13e61858bcfac7654d7fcc6d42106b4f15e11117695069600", - urls = ["https://storage.googleapis.com/mediapipe-assets/bert_text_classifier.tflite?generation=1663009542017720"], + urls = ["https://storage.googleapis.com/mediapipe-assets/bert_text_classifier.tflite?generation=1666144699858747"], ) http_file( @@ -46,12 +46,6 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/BUILD?generation=1661875663693976"], ) - http_file( - name = "com_google_mediapipe_BUILD_orig", - sha256 = "650df617b3e125e0890f1b8c936cc64c9d975707f57e616b6430fc667ce315d4", - urls = ["https://storage.googleapis.com/mediapipe-assets/BUILD.orig?generation=1665609930388174"], - ) - http_file( name = "com_google_mediapipe_burger_crop_jpg", sha256 = "8f58de573f0bf59a49c3d86cfabb9ad4061481f574aa049177e8da3963dddc50", @@ -82,6 +76,18 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/cat_mask.jpg?generation=1661875677203533"], ) + http_file( + name = "com_google_mediapipe_cat_rotated_jpg", + sha256 = "b78cee5ad14c9f36b1c25d103db371d81ca74d99030063c46a38e80bb8f38649", + urls = ["https://storage.googleapis.com/mediapipe-assets/cat_rotated.jpg?generation=1666304165042123"], + ) + + http_file( + name = "com_google_mediapipe_cat_rotated_mask_jpg", + sha256 = "f336973e7621d602f2ebc9a6ab1c62d8502272d391713f369d3b99541afda861", + urls = ["https://storage.googleapis.com/mediapipe-assets/cat_rotated_mask.jpg?generation=1666304167148173"], + ) + http_file( name = "com_google_mediapipe_cats_and_dogs_jpg", sha256 = "a2eaa7ad3a1aae4e623dd362a5f737e8a88d122597ecd1a02b3e1444db56df9c", @@ -127,7 +133,7 @@ def external_files(): http_file( name = "com_google_mediapipe_coco_ssd_mobilenet_v1_1_0_quant_2018_06_29_tflite", sha256 = "61d598093ed03ed41aa47c3a39a28ac01e960d6a810a5419b9a5016a1e9c469b", - urls = ["https://storage.googleapis.com/mediapipe-assets/coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.tflite?generation=1661875702588267"], + urls = ["https://storage.googleapis.com/mediapipe-assets/coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.tflite?generation=1666144700870810"], ) http_file( @@ -151,7 +157,7 @@ def external_files(): http_file( name = "com_google_mediapipe_dummy_gesture_recognizer_task", sha256 = "18e54586bda33300d459ca140cd045f6daf43d897224ba215a16db3423eae18e", - urls = ["https://storage.googleapis.com/mediapipe-assets/dummy_gesture_recognizer.task?generation=1665524417056146"], + urls = ["https://storage.googleapis.com/mediapipe-assets/dummy_gesture_recognizer.task?generation=1665707319890725"], ) http_file( @@ -166,12 +172,24 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/expected_left_down_hand_landmarks.prototxt?generation=1661875720230540"], ) + http_file( + name = "com_google_mediapipe_expected_left_down_hand_rotated_landmarks_prototxt", + sha256 = "c4dfdcc2e4cd366eb5f8ad227be94049eb593e3a528564611094687912463687", + urls = ["https://storage.googleapis.com/mediapipe-assets/expected_left_down_hand_rotated_landmarks.prototxt?generation=1666629474155924"], + ) + http_file( name = "com_google_mediapipe_expected_left_up_hand_landmarks_prototxt", sha256 = "1353ba617c4f048083618587cd23a8a22115f634521c153d4e1bd1ebd4f49dd7", urls = ["https://storage.googleapis.com/mediapipe-assets/expected_left_up_hand_landmarks.prototxt?generation=1661875726008879"], ) + http_file( + name = "com_google_mediapipe_expected_left_up_hand_rotated_landmarks_prototxt", + sha256 = "7fb2d33cf69d2da50952a45bad0c0618f30859e608958fee95948a6e0de63ccb", + urls = ["https://storage.googleapis.com/mediapipe-assets/expected_left_up_hand_rotated_landmarks.prototxt?generation=1666629476401757"], + ) + http_file( name = "com_google_mediapipe_expected_right_down_hand_landmarks_prototxt", sha256 = "f281b745175aaa7f458def6cf4c89521fb56302dd61a05642b3b4a4f237ffaa3", @@ -226,6 +244,18 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/feature_tensor_meta.json?generation=1665422818797346"], ) + http_file( + name = "com_google_mediapipe_fist_jpg", + sha256 = "43fa1cabf3f90d574accc9a56986e2ee48638ce59fc65af1846487f73bb2ef24", + urls = ["https://storage.googleapis.com/mediapipe-assets/fist.jpg?generation=1666999359066679"], + ) + + http_file( + name = "com_google_mediapipe_fist_landmarks_pbtxt", + sha256 = "76d6489e6163211ce5e9080e51983165bb9b24ff50146cc7487bd629f011c598", + urls = ["https://storage.googleapis.com/mediapipe-assets/fist_landmarks.pbtxt?generation=1666999360561864"], + ) + http_file( name = "com_google_mediapipe_general_meta_json", sha256 = "b95363e4bae89b9c2af484498312aaad4efc7ff57c7eadcc4e5e7adca641445f", @@ -250,12 +280,24 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/hand_detector_result_one_hand.pbtxt?generation=1662745351291628"], ) + http_file( + name = "com_google_mediapipe_hand_detector_result_one_hand_rotated_pbtxt", + sha256 = "555079c274ea91699757a0b9888c9993a8ab450069103b1bcd4ebb805a8e023c", + urls = ["https://storage.googleapis.com/mediapipe-assets/hand_detector_result_one_hand_rotated.pbtxt?generation=1666629478777955"], + ) + http_file( name = "com_google_mediapipe_hand_detector_result_two_hands_pbtxt", sha256 = "2589cb08b0ee027dc24649fe597adcfa2156a21d12ea2480f83832714ebdf95f", urls = ["https://storage.googleapis.com/mediapipe-assets/hand_detector_result_two_hands.pbtxt?generation=1662745353586157"], ) + http_file( + name = "com_google_mediapipe_hand_landmarker_task", + sha256 = "2ed44f10872e87a5834b9b1130fb9ada30e107af2c6fcc4562ad788aca4e7bc4", + urls = ["https://storage.googleapis.com/mediapipe-assets/hand_landmarker.task?generation=1666153732577904"], + ) + http_file( name = "com_google_mediapipe_hand_landmark_full_tflite", sha256 = "11c272b891e1a99ab034208e23937a8008388cf11ed2a9d776ed3d01d0ba00e3", @@ -269,9 +311,9 @@ def external_files(): ) http_file( - name = "com_google_mediapipe_hand_landmark_task", - sha256 = "dd830295598e48e6bbbdf22fd9e69538fa07768106cd9ceb04d5462ca7e38c95", - urls = ["https://storage.googleapis.com/mediapipe-assets/hand_landmark.task?generation=1665707323647357"], + name = "com_google_mediapipe_hand_landmark_tflite", + sha256 = "bad88ac1fd144f034e00f075afcade4f3a21d0d09c41bee8dd50504dacd70efd", + urls = ["https://storage.googleapis.com/mediapipe-assets/hand_landmark.tflite?generation=1666153735814956"], ) http_file( @@ -346,12 +388,24 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/knift_labelmap.txt?generation=1661875792821628"], ) + http_file( + name = "com_google_mediapipe_labels_txt", + sha256 = "536feacc519de3d418de26b2effb4d75694a8c4c0063e36499a46fa8061e2da9", + urls = ["https://storage.googleapis.com/mediapipe-assets/labels.txt?generation=1665988394538324"], + ) + http_file( name = "com_google_mediapipe_left_hands_jpg", sha256 = "4b5134daa4cb60465535239535f9f74c2842aba3aa5fd30bf04ef5678f93d87f", urls = ["https://storage.googleapis.com/mediapipe-assets/left_hands.jpg?generation=1661875796949017"], ) + http_file( + name = "com_google_mediapipe_left_hands_rotated_jpg", + sha256 = "8609c6202bca43a99bbf23fa8e687e49fa525e89481152e4c0987f46d60d7931", + urls = ["https://storage.googleapis.com/mediapipe-assets/left_hands_rotated.jpg?generation=1666037068103465"], + ) + http_file( name = "com_google_mediapipe_mobilebert_embedding_with_metadata_tflite", sha256 = "fa47142dcc6f446168bc672f2df9605b6da5d0c0d6264e9be62870282365b95c", @@ -424,18 +478,42 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/mobilenet_v1_0.25_224_quant_without_subgraph_metadata.tflite?generation=1661875836078124"], ) + http_file( + name = "com_google_mediapipe_mobilenet_v2_1_0_224_json", + sha256 = "94613ea9539a20a3352604004be6d4d64d4d76250bc9042fcd8685c9a8498517", + urls = ["https://storage.googleapis.com/mediapipe-assets/mobilenet_v2_1.0_224.json?generation=1666633416316646"], + ) + + http_file( + name = "com_google_mediapipe_mobilenet_v2_1_0_224_quant_json", + sha256 = "3703eadcf838b65bbc2b2aa11dbb1f1bc654c7a09a7aba5ca75a26096484a8ac", + urls = ["https://storage.googleapis.com/mediapipe-assets/mobilenet_v2_1.0_224_quant.json?generation=1666633418665507"], + ) + http_file( name = "com_google_mediapipe_mobilenet_v2_1_0_224_quant_tflite", sha256 = "f08d447cde49b4e0446428aa921aff0a14ea589fa9c5817b31f83128e9a43c1d", urls = ["https://storage.googleapis.com/mediapipe-assets/mobilenet_v2_1.0_224_quant.tflite?generation=1664340173966530"], ) + http_file( + name = "com_google_mediapipe_mobilenet_v2_1_0_224_quant_without_metadata_tflite", + sha256 = "f08d447cde49b4e0446428aa921aff0a14ea589fa9c5817b31f83128e9a43c1d", + urls = ["https://storage.googleapis.com/mediapipe-assets/mobilenet_v2_1.0_224_quant_without_metadata.tflite?generation=1665988405130772"], + ) + http_file( name = "com_google_mediapipe_mobilenet_v2_1_0_224_tflite", sha256 = "ff5cb7f9e62c92ebdad971f8a98aa6b3106d82a64587a7787c6a385c9e791339", urls = ["https://storage.googleapis.com/mediapipe-assets/mobilenet_v2_1.0_224.tflite?generation=1661875840611150"], ) + http_file( + name = "com_google_mediapipe_mobilenet_v2_1_0_224_without_metadata_tflite", + sha256 = "9f3bc29e38e90842a852bfed957dbf5e36f2d97a91dd17736b1e5c0aca8d3303", + urls = ["https://storage.googleapis.com/mediapipe-assets/mobilenet_v2_1.0_224_without_metadata.tflite?generation=1665988408360823"], + ) + http_file( name = "com_google_mediapipe_mobilenet_v3_small_100_224_embedder_tflite", sha256 = "f7b9a563cb803bdcba76e8c7e82abde06f5c7a8e67b5e54e43e23095dfe79a78", @@ -451,7 +529,7 @@ def external_files(): http_file( name = "com_google_mediapipe_mobile_object_labeler_v1_tflite", sha256 = "9400671e04685f5277edd3052a311cc51533de9da94255c52ebde1e18484c77c", - urls = ["https://storage.googleapis.com/mediapipe-assets/mobile_object_labeler_v1.tflite?generation=1661875846924538"], + urls = ["https://storage.googleapis.com/mediapipe-assets/mobile_object_labeler_v1.tflite?generation=1666144701839813"], ) http_file( @@ -544,6 +622,18 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/pointing_up_landmarks.pbtxt?generation=1665174976408451"], ) + http_file( + name = "com_google_mediapipe_pointing_up_rotated_jpg", + sha256 = "50ff66f50281207072a038e5bb6648c43f4aacbfb8204a4d2591868756aaeff1", + urls = ["https://storage.googleapis.com/mediapipe-assets/pointing_up_rotated.jpg?generation=1666037072219697"], + ) + + http_file( + name = "com_google_mediapipe_pointing_up_rotated_landmarks_pbtxt", + sha256 = "5ec37218d8b613436f5c10121dc689bf9ee69af0656a6ccf8c2e3e8b652e2ad6", + urls = ["https://storage.googleapis.com/mediapipe-assets/pointing_up_rotated_landmarks.pbtxt?generation=1666629486774022"], + ) + http_file( name = "com_google_mediapipe_pose_detection_tflite", sha256 = "a63c614bef30d35947f13be361820b1e4e3bec9cfeebf4d11216a18373108e85", @@ -580,6 +670,12 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/right_hands.jpg?generation=1661875908672404"], ) + http_file( + name = "com_google_mediapipe_right_hands_rotated_jpg", + sha256 = "b3bdf692f0d54b86c8b67e6d1286dd0078fbe6e9dfcd507b187e3bd8b398c0f9", + urls = ["https://storage.googleapis.com/mediapipe-assets/right_hands_rotated.jpg?generation=1666037076873345"], + ) + http_file( name = "com_google_mediapipe_score_calibration_file_meta_json", sha256 = "6a3c305620371f662419a496f75be5a10caebca7803b1e99d8d5d22ba51cda94", @@ -598,6 +694,12 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/score_calibration.txt?generation=1665422847392804"], ) + http_file( + name = "com_google_mediapipe_score_thresholding_meta_json", + sha256 = "7bb74f21c2d7f0237675ed7c09d7b7afd3507c8373f51dc75fa0507852f6ee19", + urls = ["https://storage.googleapis.com/mediapipe-assets/score_thresholding_meta.json?generation=1667273953630766"], + ) + http_file( name = "com_google_mediapipe_segmentation_golden_rotation0_png", sha256 = "9ee993919b753118928ba2d14f7c5c83a6cfc23355e6943dac4ad81eedd73069", @@ -724,6 +826,12 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/thumb_up_landmarks.pbtxt?generation=1665174979747784"], ) + http_file( + name = "com_google_mediapipe_thumb_up_rotated_landmarks_pbtxt", + sha256 = "6645bbd98ea7f90b3e1ba297e16ea5280847fc5bf5400726d98c282f6c597257", + urls = ["https://storage.googleapis.com/mediapipe-assets/thumb_up_rotated_landmarks.pbtxt?generation=1666629489421733"], + ) + http_file( name = "com_google_mediapipe_two_heads_16000_hz_mono_wav", sha256 = "a291a9c22c39bba30138a26915e154a96286ba6ca3b413053123c504a58cce3b", @@ -748,6 +856,18 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/universal_sentence_encoder_qa_with_metadata.tflite?generation=1665445919252005"], ) + http_file( + name = "com_google_mediapipe_victory_jpg", + sha256 = "84cb8853e3df614e0cb5c93a25e3e2f38ea5e4f92fd428ee7d867ed3479d5764", + urls = ["https://storage.googleapis.com/mediapipe-assets/victory.jpg?generation=1666999364225126"], + ) + + http_file( + name = "com_google_mediapipe_victory_landmarks_pbtxt", + sha256 = "b25ab4f222674489f543afb6454396ecbc1437a7ae6213dbf0553029ae939ab0", + urls = ["https://storage.googleapis.com/mediapipe-assets/victory_landmarks.pbtxt?generation=1666999366036622"], + ) + http_file( name = "com_google_mediapipe_vocab_for_regex_tokenizer_txt", sha256 = "b1134b10927a53ce4224bbc30ccf075c9969c94ebf40c368966d1dcf445ca923",