From ae05ad04b3ae43d475ccb2868e23f1418fea8746 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Mon, 10 May 2021 21:52:16 -0700 Subject: [PATCH] Project import generated by Copybara. GitOrigin-RevId: 016275ca4057540b2370ed4531dbc81eb92caae2 --- docs/solutions/pose.md | 2 +- mediapipe/calculators/image/BUILD | 30 +++++ .../image/image_clone_calculator.cc | 125 ++++++++++++++++++ .../image/image_clone_calculator.proto | 28 ++++ mediapipe/gpu/BUILD | 1 + mediapipe/gpu/gl_scaler_calculator.cc | 33 ++++- 6 files changed, 211 insertions(+), 8 deletions(-) create mode 100644 mediapipe/calculators/image/image_clone_calculator.cc create mode 100644 mediapipe/calculators/image/image_clone_calculator.proto diff --git a/docs/solutions/pose.md b/docs/solutions/pose.md index c2c96293..feed2ad3 100644 --- a/docs/solutions/pose.md +++ b/docs/solutions/pose.md @@ -399,4 +399,4 @@ on how to build MediaPipe examples. * [Python Colab](https://mediapipe.page.link/pose_py_colab) [`mAP`]: https://cocodataset.org/#keypoints-eval -[`PCK@0.2`]: https\://github.com/cbsudux/Human-Pose-Estimation-101 +[`PCK@0.2`]: https://github.com/cbsudux/Human-Pose-Estimation-101 diff --git a/mediapipe/calculators/image/BUILD b/mediapipe/calculators/image/BUILD index c92b9a2a..39f81c04 100644 --- a/mediapipe/calculators/image/BUILD +++ b/mediapipe/calculators/image/BUILD @@ -405,6 +405,36 @@ cc_library( alwayslink = 1, ) +mediapipe_proto_library( + name = "image_clone_calculator_proto", + srcs = ["image_clone_calculator.proto"], + visibility = ["//visibility:public"], + deps = [ + "//mediapipe/framework:calculator_options_proto", + "//mediapipe/framework:calculator_proto", + ], +) + +cc_library( + name = "image_clone_calculator", + srcs = ["image_clone_calculator.cc"], + visibility = ["//visibility:public"], + deps = [ + ":image_clone_calculator_cc_proto", + "//mediapipe/framework/api2:node", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework/formats:image", + "//mediapipe/framework/port:ret_check", + "//mediapipe/framework/port:status", + ] + select({ + "//mediapipe/gpu:disable_gpu": [], + "//conditions:default": [ + "//mediapipe/gpu:gl_calculator_helper", + ], + }), + alwayslink = 1, +) + cc_library( name = "image_properties_calculator", srcs = ["image_properties_calculator.cc"], diff --git a/mediapipe/calculators/image/image_clone_calculator.cc b/mediapipe/calculators/image/image_clone_calculator.cc new file mode 100644 index 00000000..107c42b9 --- /dev/null +++ b/mediapipe/calculators/image/image_clone_calculator.cc @@ -0,0 +1,125 @@ +// Copyright 2021 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 "mediapipe/calculators/image/image_clone_calculator.pb.h" +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/image.h" +#include "mediapipe/framework/port/status.h" + +#if !MEDIAPIPE_DISABLE_GPU +#include "mediapipe/gpu/gl_calculator_helper.h" +#endif // !MEDIAPIPE_DISABLE_GPU + +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 + +// Clones an input image and makes sure in the output clone the pixel data are +// stored on the target storage (CPU vs GPU) specified in the calculator option. +// +// The clone shares ownership of the input pixel data on the existing storage. +// If the target storage is diffrent from the existing one, then the data is +// further copied there. +// +// Example usage: +// node { +// calculator: "ImageCloneCalculator" +// input_stream: "input" +// output_stream: "output" +// options: { +// [mediapipe.ImageCloneCalculatorOptions.ext] { +// output_on_gpu: true +// } +// } +// } +class ImageCloneCalculator : public Node { + public: + static constexpr Input kIn{""}; + static constexpr Output kOut{""}; + + MEDIAPIPE_NODE_CONTRACT(kIn, kOut); + + static absl::Status UpdateContract(CalculatorContract* cc) { +#if MEDIAPIPE_DISABLE_GPU + if (cc->Options().output_on_gpu()) { + return absl::UnimplementedError( + "GPU processing is disabled in build flags"); + } +#else + MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc)); +#endif // MEDIAPIPE_DISABLE_GPU + return absl::OkStatus(); + } + + absl::Status Open(CalculatorContext* cc) override { + const auto& options = cc->Options(); + output_on_gpu_ = options.output_on_gpu(); +#if !MEDIAPIPE_DISABLE_GPU + MP_RETURN_IF_ERROR(gpu_helper_.Open(cc)); +#endif // !MEDIAPIPE_DISABLE_GPU + return absl::OkStatus(); + } + + absl::Status Process(CalculatorContext* cc) override { + std::unique_ptr output; + const auto& input = *kIn(cc); + if (input.UsesGpu()) { +#if !MEDIAPIPE_DISABLE_GPU + // Create an output Image that co-owns the underlying texture buffer as + // the input Image. + output = std::make_unique(input.GetGpuBuffer()); +#endif // !MEDIAPIPE_DISABLE_GPU + } else { + // Make a copy of the input packet to co-own the input Image. + mediapipe::Packet* packet_copy_ptr = + new mediapipe::Packet(kIn(cc).packet()); + // Create an output Image that (co-)owns a new ImageFrame that points to + // the same pixel data as the input Image and also owns the packet + // copy. As a result, the output Image indirectly co-owns the input + // Image. This ensures a correct life span of the shared pixel data. + output = std::make_unique(std::make_unique( + input.image_format(), input.width(), input.height(), input.step(), + const_cast(input.GetImageFrameSharedPtr()->PixelData()), + [packet_copy_ptr](uint8*) { delete packet_copy_ptr; })); + } + + if (output_on_gpu_) { +#if !MEDIAPIPE_DISABLE_GPU + gpu_helper_.RunInGlContext([&output]() { output->ConvertToGpu(); }); +#endif // !MEDIAPIPE_DISABLE_GPU + } else { + output->ConvertToCpu(); + } + kOut(cc).Send(std::move(output)); + + return absl::OkStatus(); + } + + private: + bool output_on_gpu_; +#if !MEDIAPIPE_DISABLE_GPU + mediapipe::GlCalculatorHelper gpu_helper_; +#endif // !MEDIAPIPE_DISABLE_GPU +}; +MEDIAPIPE_REGISTER_NODE(ImageCloneCalculator); + +} // namespace api2 +} // namespace mediapipe diff --git a/mediapipe/calculators/image/image_clone_calculator.proto b/mediapipe/calculators/image/image_clone_calculator.proto new file mode 100644 index 00000000..6fa05dcc --- /dev/null +++ b/mediapipe/calculators/image/image_clone_calculator.proto @@ -0,0 +1,28 @@ +// Copyright 2021 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 ImageCloneCalculatorOptions { + extend CalculatorOptions { + optional ImageCloneCalculatorOptions ext = 372781894; + } + + // Whether the output clone should have pixel data already available on GPU. + optional bool output_on_gpu = 1 [default = false]; +} diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD index 1da8115a..fbcdb4ab 100644 --- a/mediapipe/gpu/BUILD +++ b/mediapipe/gpu/BUILD @@ -753,6 +753,7 @@ cc_library( ":gl_simple_shaders", ":shader_util", "//mediapipe/framework:calculator_framework", + "//mediapipe/framework/formats:image", "//mediapipe/framework/port:ret_check", "//mediapipe/framework/port:status", "//mediapipe/framework/tool:options_util", diff --git a/mediapipe/gpu/gl_scaler_calculator.cc b/mediapipe/gpu/gl_scaler_calculator.cc index e50ef228..61918762 100644 --- a/mediapipe/gpu/gl_scaler_calculator.cc +++ b/mediapipe/gpu/gl_scaler_calculator.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/image.h" #include "mediapipe/framework/port/ret_check.h" #include "mediapipe/framework/port/status.h" #include "mediapipe/framework/tool/options_util.h" @@ -33,6 +34,8 @@ typedef int DimensionsPacketType[2]; namespace mediapipe { +using Image = mediapipe::Image; + // Scales, rotates, horizontal or vertical flips the image. // See GlSimpleCalculatorBase for inputs, outputs and input side packets. // Additional input streams: @@ -99,8 +102,17 @@ REGISTER_CALCULATOR(GlScalerCalculator); // static absl::Status GlScalerCalculator::GetContract(CalculatorContract* cc) { - TagOrIndex(&cc->Inputs(), "VIDEO", 0).Set(); - TagOrIndex(&cc->Outputs(), "VIDEO", 0).Set(); + if (cc->Inputs().HasTag("IMAGE")) { + cc->Inputs().Tag("IMAGE").Set(); + } else { + TagOrIndex(&cc->Inputs(), "VIDEO", 0).Set(); + } + if (cc->Outputs().HasTag("IMAGE")) { + cc->Outputs().Tag("IMAGE").Set(); + } else { + TagOrIndex(&cc->Outputs(), "VIDEO", 0).Set(); + } + if (cc->Inputs().HasTag("ROTATION")) { cc->Inputs().Tag("ROTATION").Set(); } @@ -198,7 +210,10 @@ absl::Status GlScalerCalculator::Process(CalculatorContext* cc) { } return helper_.RunInGlContext([this, cc]() -> absl::Status { - const auto& input = TagOrIndex(cc->Inputs(), "VIDEO", 0).Get(); + const auto& input = + cc->Inputs().HasTag("IMAGE") + ? cc->Inputs().Tag("IMAGE").Get().GetGpuBuffer() + : TagOrIndex(cc->Inputs(), "VIDEO", 0).Get(); QuadRenderer* renderer = nullptr; GlTexture src1; GlTexture src2; @@ -289,10 +304,14 @@ absl::Status GlScalerCalculator::Process(CalculatorContext* cc) { glFlush(); - auto output = dst.GetFrame(); - - TagOrIndex(&cc->Outputs(), "VIDEO", 0) - .Add(output.release(), cc->InputTimestamp()); + if (cc->Outputs().HasTag("IMAGE")) { + auto output = dst.GetFrame(); + cc->Outputs().Tag("IMAGE").Add(output.release(), cc->InputTimestamp()); + } else { + auto output = dst.GetFrame(); + TagOrIndex(&cc->Outputs(), "VIDEO", 0) + .Add(output.release(), cc->InputTimestamp()); + } return absl::OkStatus(); });