diff --git a/mediapipe/tasks/cc/vision/utils/BUILD b/mediapipe/tasks/cc/vision/utils/BUILD index 442fd2717..0eb5ba75c 100644 --- a/mediapipe/tasks/cc/vision/utils/BUILD +++ b/mediapipe/tasks/cc/vision/utils/BUILD @@ -111,3 +111,23 @@ cc_test( "//mediapipe/tasks/cc/components/containers:rect", ], ) + +cc_library( + name = "data_renderer", + srcs = ["data_renderer.cc"], + hdrs = ["data_renderer.h"], + deps = [ + "//mediapipe/calculators/util:annotation_overlay_calculator", + "//mediapipe/calculators/util:landmarks_to_render_data_calculator", + "//mediapipe/calculators/util:landmarks_to_render_data_calculator_cc_proto", + "//mediapipe/calculators/util:rect_to_render_data_calculator_cc_proto", + "//mediapipe/calculators/util:rect_to_render_scale_calculator", + "//mediapipe/calculators/util:rect_to_render_scale_calculator_cc_proto", + "//mediapipe/framework/api2:builder", + "//mediapipe/framework/formats:image", + "//mediapipe/framework/formats:landmark_cc_proto", + "//mediapipe/framework/formats:rect_cc_proto", + "//mediapipe/util:render_data_cc_proto", + "@com_google_absl//absl/types:span", + ], +) diff --git a/mediapipe/tasks/cc/vision/utils/data_renderer.cc b/mediapipe/tasks/cc/vision/utils/data_renderer.cc new file mode 100644 index 000000000..aeefbba2f --- /dev/null +++ b/mediapipe/tasks/cc/vision/utils/data_renderer.cc @@ -0,0 +1,88 @@ +/* Copyright 2023 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/tasks/cc/vision/utils/data_renderer.h" + +#include +#include +#include + +#include "absl/types/span.h" +#include "mediapipe/calculators/util/landmarks_to_render_data_calculator.pb.h" +#include "mediapipe/calculators/util/rect_to_render_data_calculator.pb.h" +#include "mediapipe/calculators/util/rect_to_render_scale_calculator.pb.h" +#include "mediapipe/framework/api2/builder.h" +#include "mediapipe/framework/formats/image.h" +#include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/formats/rect.pb.h" +#include "mediapipe/util/render_data.pb.h" + +namespace mediapipe::tasks::vision::utils { + +using ::mediapipe::api2::builder::Graph; +using ::mediapipe::api2::builder::Stream; + +Stream Render(Stream image, + absl::Span> render_data_list, + Graph& graph) { + auto& annotation_overlay = graph.AddNode("AnnotationOverlayCalculator"); + image >> annotation_overlay.In("UIMAGE"); + for (int i = 0; i < render_data_list.size(); ++i) { + render_data_list[i] >> annotation_overlay.In(i); + } + return annotation_overlay.Out("UIMAGE").Cast(); +} + +Stream RenderLandmarks( + Stream landmarks, + std::optional> render_scale, + const mediapipe::LandmarksToRenderDataCalculatorOptions& renderer_options, + Graph& graph) { + auto& landmarks_render = graph.AddNode("LandmarksToRenderDataCalculator"); + landmarks_render + .GetOptions() + .CopyFrom(renderer_options); + landmarks >> landmarks_render.In("NORM_LANDMARKS"); + if (render_scale.has_value()) { + *render_scale >> landmarks_render.In("RENDER_SCALE"); + } + auto render_data = landmarks_render.Out("RENDER_DATA"); + return render_data.Cast(); +} + +Stream GetRenderScale(Stream> image_size, + Stream roi, float multiplier, + Graph& graph) { + auto& to_render_scale = graph.AddNode("RectToRenderScaleCalculator"); + to_render_scale.GetOptions() + .set_multiplier(multiplier); + roi >> to_render_scale.In("NORM_RECT"); + image_size >> to_render_scale.In("IMAGE_SIZE"); + return to_render_scale.Out("RENDER_SCALE").Cast(); +} + +Stream RenderRect( + Stream rect, + const mediapipe::RectToRenderDataCalculatorOptions& renderer_options, + Graph& graph) { + auto& rect_render = graph.AddNode("RectToRenderDataCalculator"); + rect_render.GetOptions() + .CopyFrom(renderer_options); + rect >> rect_render.In("NORM_RECT"); + auto render_data = rect_render.Out("RENDER_DATA"); + return render_data.Cast(); +} + +} // namespace mediapipe::tasks::vision::utils diff --git a/mediapipe/tasks/cc/vision/utils/data_renderer.h b/mediapipe/tasks/cc/vision/utils/data_renderer.h new file mode 100644 index 000000000..f58f94ee8 --- /dev/null +++ b/mediapipe/tasks/cc/vision/utils/data_renderer.h @@ -0,0 +1,69 @@ +/* Copyright 2023 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. +==============================================================================*/ + +#ifndef MEDIAPIPE_TASKS_CC_VISION_UTILS_DATA_RENDERER_H_ +#define MEDIAPIPE_TASKS_CC_VISION_UTILS_DATA_RENDERER_H_ + +#include +#include + +#include "absl/types/span.h" +#include "mediapipe/calculators/util/landmarks_to_render_data_calculator.pb.h" +#include "mediapipe/calculators/util/rect_to_render_data_calculator.pb.h" +#include "mediapipe/framework/api2/builder.h" +#include "mediapipe/framework/formats/image.h" +#include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/formats/rect.pb.h" +#include "mediapipe/util/render_data.pb.h" + +namespace mediapipe::tasks::vision::utils { + +// Adds a node to the provided graph that renders the render_data_list on the +// given image, and returns the rendered image. +api2::builder::Stream Render( + api2::builder::Stream image, + absl::Span> render_data_list, + api2::builder::Graph& graph); + +// Adds a node to the provided graph that infers the render scale from the image +// size and the object RoI. It will give you bigger rendered primitives for +// bigger/closer objects and smaller primitives for smaller/far objects. The +// primitives scale is proportional to `roi_size * multiplier`. +// +// See more details in +// mediapipe/calculators/util/rect_to_render_scale_calculator.cc +api2::builder::Stream GetRenderScale( + api2::builder::Stream> image_size, + api2::builder::Stream roi, float multiplier, + api2::builder::Graph& graph); + +// Adds a node to the provided graph that gets the landmarks render data +// according to the renderer_options. +api2::builder::Stream RenderLandmarks( + api2::builder::Stream landmarks, + std::optional> render_scale, + const mediapipe::LandmarksToRenderDataCalculatorOptions& renderer_options, + api2::builder::Graph& graph); + +// Adds a node to the provided graph that gets the rect render data according to +// the renderer_options. +api2::builder::Stream RenderRect( + api2::builder::Stream rect, + const mediapipe::RectToRenderDataCalculatorOptions& renderer_options, + api2::builder::Graph& graph); + +} // namespace mediapipe::tasks::vision::utils + +#endif // MEDIAPIPE_TASKS_CC_VISION_UTILS_DATA_RENDERER_H_ diff --git a/mediapipe/tasks/cc/vision/utils/data_renderer_test.cc b/mediapipe/tasks/cc/vision/utils/data_renderer_test.cc new file mode 100644 index 000000000..b42c335b2 --- /dev/null +++ b/mediapipe/tasks/cc/vision/utils/data_renderer_test.cc @@ -0,0 +1,133 @@ +/* Copyright 2023 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/tasks/cc/vision/utils/data_renderer.h" + +#include +#include + +#include "absl/types/span.h" +#include "mediapipe/framework/api2/builder.h" +#include "mediapipe/framework/calculator.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/gmock.h" +#include "mediapipe/framework/port/gtest.h" +#include "mediapipe/framework/port/parse_text_proto.h" +#include "mediapipe/util/render_data.pb.h" + +namespace mediapipe::tasks::vision::utils { +namespace { + +using ::mediapipe::CalculatorGraphConfig; +using ::mediapipe::EqualsProto; +using ::mediapipe::NormalizedRect; +using ::mediapipe::api2::builder::Graph; +using ::mediapipe::api2::builder::Stream; + +TEST(DataRenderer, Render) { + Graph graph; + Stream image_in = graph.In("IMAGE").Cast(); + Stream render_data_in = + graph.In("RENDER_DATA").Cast(); + std::vector> render_data_list = {render_data_in}; + Stream image_out = + Render(image_in, absl::Span>(render_data_list), graph); + image_out.SetName("image_out"); + EXPECT_THAT( + graph.GetConfig(), + EqualsProto(mediapipe::ParseTextProtoOrDie(R"pb( + node { + calculator: "AnnotationOverlayCalculator" + input_stream: "__stream_1" + input_stream: "UIMAGE:__stream_0" + output_stream: "UIMAGE:image_out" + } + input_stream: "IMAGE:__stream_0" + input_stream: "RENDER_DATA:__stream_1" + )pb"))); +} + +TEST(DataRenderer, RenderLandmarks) { + Graph graph; + Stream rect = + graph.In("NORM_LANDMARKS").Cast(); + Stream render_data = + RenderLandmarks(rect, std::nullopt, {}, graph); + render_data.SetName("render_data"); + EXPECT_THAT( + graph.GetConfig(), + EqualsProto(mediapipe::ParseTextProtoOrDie(R"pb( + node { + calculator: "LandmarksToRenderDataCalculator" + input_stream: "NORM_LANDMARKS:__stream_0" + output_stream: "RENDER_DATA:render_data" + options { + [mediapipe.LandmarksToRenderDataCalculatorOptions.ext] {} + } + } + input_stream: "NORM_LANDMARKS:__stream_0" + )pb"))); +} + +TEST(DataRenderer, GetRenderScale) { + Graph graph; + Stream> image_size = + graph.In("IMAGE_SIZE").Cast>(); + Stream roi = graph.In("ROI").Cast(); + Stream render_scale = GetRenderScale(image_size, roi, 0.0001, graph); + render_scale.SetName("render_scale"); + EXPECT_THAT( + graph.GetConfig(), + EqualsProto(mediapipe::ParseTextProtoOrDie(R"pb( + node { + calculator: "RectToRenderScaleCalculator" + input_stream: "IMAGE_SIZE:__stream_0" + input_stream: "NORM_RECT:__stream_1" + output_stream: "RENDER_SCALE:render_scale" + options { + [mediapipe.RectToRenderScaleCalculatorOptions.ext] { + multiplier: 0.0001 + } + } + } + input_stream: "IMAGE_SIZE:__stream_0" + input_stream: "ROI:__stream_1" + )pb"))); +} + +TEST(DataRenderer, RenderRect) { + Graph graph; + Stream rect = graph.In("NORM_RECT").Cast(); + Stream render_data = RenderRect(rect, {}, graph); + render_data.SetName("render_data"); + EXPECT_THAT( + graph.GetConfig(), + EqualsProto(mediapipe::ParseTextProtoOrDie(R"pb( + node { + calculator: "RectToRenderDataCalculator" + input_stream: "NORM_RECT:__stream_0" + output_stream: "RENDER_DATA:render_data" + options { + [mediapipe.RectToRenderDataCalculatorOptions.ext] {} + } + } + input_stream: "NORM_RECT:__stream_0" + )pb"))); +} + +} // namespace +} // namespace mediapipe::tasks::vision::utils