diff --git a/mediapipe/calculators/util/BUILD b/mediapipe/calculators/util/BUILD index 6ac60d2c1..710a60d8a 100644 --- a/mediapipe/calculators/util/BUILD +++ b/mediapipe/calculators/util/BUILD @@ -1270,6 +1270,50 @@ cc_library( alwayslink = 1, ) +mediapipe_proto_library( + name = "flat_color_image_calculator_proto", + srcs = ["flat_color_image_calculator.proto"], + deps = [ + "//mediapipe/framework:calculator_options_proto", + "//mediapipe/framework:calculator_proto", + "//mediapipe/util:color_proto", + ], +) + +cc_library( + name = "flat_color_image_calculator", + srcs = ["flat_color_image_calculator.cc"], + deps = [ + ":flat_color_image_calculator_cc_proto", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework/api2:node", + "//mediapipe/framework/formats:image", + "//mediapipe/framework/formats:image_frame", + "//mediapipe/framework/formats:image_frame_opencv", + "//mediapipe/framework/port:opencv_core", + "//mediapipe/util:color_cc_proto", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + ], + alwayslink = 1, +) + +cc_test( + name = "flat_color_image_calculator_test", + srcs = ["flat_color_image_calculator_test.cc"], + deps = [ + ":flat_color_image_calculator", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework:calculator_runner", + "//mediapipe/framework:packet", + "//mediapipe/framework/formats:image", + "//mediapipe/framework/formats:image_frame", + "//mediapipe/framework/port:gtest", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/util:color_cc_proto", + ], +) + cc_library( name = "from_image_calculator", srcs = ["from_image_calculator.cc"], diff --git a/mediapipe/calculators/util/flat_color_image_calculator.cc b/mediapipe/calculators/util/flat_color_image_calculator.cc new file mode 100644 index 000000000..71d3582c5 --- /dev/null +++ b/mediapipe/calculators/util/flat_color_image_calculator.cc @@ -0,0 +1,138 @@ +// 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 + +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" +#include "mediapipe/calculators/util/flat_color_image_calculator.pb.h" +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/image.h" +#include "mediapipe/framework/formats/image_frame.h" +#include "mediapipe/framework/formats/image_frame_opencv.h" +#include "mediapipe/framework/port/opencv_core_inc.h" +#include "mediapipe/util/color.pb.h" + +namespace mediapipe { + +namespace { + +using ::mediapipe::api2::Input; +using ::mediapipe::api2::Node; +using ::mediapipe::api2::Output; +} // namespace + +// A calculator for generating an image filled with a single color. +// +// Inputs: +// IMAGE (Image, optional) +// If provided, the output will have the same size +// COLOR (Color proto, optional) +// Color to paint the output with. Takes precedence over the equivalent +// calculator options. +// +// Outputs: +// IMAGE (Image) +// Image filled with the requested color. +// +// Example useage: +// node { +// calculator: "FlatColorImageCalculator" +// input_stream: "IMAGE:image" +// input_stream: "COLOR:color" +// output_stream: "IMAGE:blank_image" +// options { +// [mediapipe.FlatColorImageCalculatorOptions.ext] { +// color: { +// r: 255 +// g: 255 +// b: 255 +// } +// } +// } +// } + +class FlatColorImageCalculator : public Node { + public: + static constexpr Input::Optional kInImage{"IMAGE"}; + static constexpr Input::Optional kInColor{"COLOR"}; + static constexpr Output kOutImage{"IMAGE"}; + + MEDIAPIPE_NODE_CONTRACT(kInImage, kInColor, kOutImage); + + static absl::Status UpdateContract(CalculatorContract* cc) { + const auto& options = cc->Options(); + + RET_CHECK(kInImage(cc).IsConnected() ^ + (options.has_output_height() || options.has_output_width())) + << "Either set IMAGE input stream, or set through options"; + RET_CHECK(kInColor(cc).IsConnected() ^ options.has_color()) + << "Either set COLOR input stream, or set through options"; + + return absl::OkStatus(); + } + + absl::Status Open(CalculatorContext* cc) override; + absl::Status Process(CalculatorContext* cc) override; + + private: + bool use_dimension_from_option_ = false; + bool use_color_from_option_ = false; +}; +MEDIAPIPE_REGISTER_NODE(FlatColorImageCalculator); + +absl::Status FlatColorImageCalculator::Open(CalculatorContext* cc) { + use_dimension_from_option_ = !kInImage(cc).IsConnected(); + use_color_from_option_ = !kInColor(cc).IsConnected(); + return absl::OkStatus(); +} + +absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) { + const auto& options = cc->Options(); + + int output_height = -1; + int output_width = -1; + if (use_dimension_from_option_) { + output_height = options.output_height(); + output_width = options.output_width(); + } else if (!kInImage(cc).IsEmpty()) { + const Image& input_image = kInImage(cc).Get(); + output_height = input_image.height(); + output_width = input_image.width(); + } else { + return absl::OkStatus(); + } + + Color color; + if (use_color_from_option_) { + color = options.color(); + } else if (!kInColor(cc).IsEmpty()) { + color = kInColor(cc).Get(); + } else { + return absl::OkStatus(); + } + + auto output_frame = std::make_shared(ImageFormat::SRGB, + output_width, output_height); + cv::Mat output_mat = mediapipe::formats::MatView(output_frame.get()); + + output_mat.setTo(cv::Scalar(color.r(), color.g(), color.b())); + + kOutImage(cc).Send(Image(output_frame)); + + return absl::OkStatus(); +} + +} // namespace mediapipe diff --git a/mediapipe/calculators/util/flat_color_image_calculator.proto b/mediapipe/calculators/util/flat_color_image_calculator.proto new file mode 100644 index 000000000..183bc796e --- /dev/null +++ b/mediapipe/calculators/util/flat_color_image_calculator.proto @@ -0,0 +1,32 @@ +// 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. + +syntax = "proto2"; + +package mediapipe; + +import "mediapipe/framework/calculator.proto"; +import "mediapipe/util/color.proto"; + +message FlatColorImageCalculatorOptions { + extend CalculatorOptions { + optional FlatColorImageCalculatorOptions ext = 515548435; + } + + // Output dimensions. + optional int32 output_width = 1; + optional int32 output_height = 2; + // The color to fill with in the output image. + optional Color color = 3; +} diff --git a/mediapipe/calculators/util/flat_color_image_calculator_test.cc b/mediapipe/calculators/util/flat_color_image_calculator_test.cc new file mode 100644 index 000000000..53c6de1b1 --- /dev/null +++ b/mediapipe/calculators/util/flat_color_image_calculator_test.cc @@ -0,0 +1,210 @@ +// 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 + +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/calculator_runner.h" +#include "mediapipe/framework/formats/image.h" +#include "mediapipe/framework/formats/image_frame.h" +#include "mediapipe/framework/packet.h" +#include "mediapipe/framework/port/gmock.h" +#include "mediapipe/framework/port/gtest.h" +#include "mediapipe/framework/port/status_matchers.h" +#include "mediapipe/util/color.pb.h" + +namespace mediapipe { +namespace { + +using ::testing::HasSubstr; + +constexpr char kImageTag[] = "IMAGE"; +constexpr char kColorTag[] = "COLOR"; +constexpr int kImageWidth = 256; +constexpr int kImageHeight = 256; + +TEST(FlatColorImageCalculatorTest, SpecifyColorThroughOptions) { + CalculatorRunner runner(R"pb( + calculator: "FlatColorImageCalculator" + input_stream: "IMAGE:image" + output_stream: "IMAGE:out_image" + options { + [mediapipe.FlatColorImageCalculatorOptions.ext] { + color: { + r: 100, + g: 200, + b: 255, + } + } + } + )pb"); + + auto image_frame = std::make_shared(ImageFormat::SRGB, + kImageWidth, kImageHeight); + + for (int ts = 0; ts < 3; ++ts) { + runner.MutableInputs()->Tag(kImageTag).packets.push_back( + MakePacket(image_frame).At(Timestamp(ts))); + } + MP_ASSERT_OK(runner.Run()); + + const auto& outputs = runner.Outputs().Tag(kImageTag).packets; + ASSERT_EQ(outputs.size(), 3); + + for (const auto& packet : outputs) { + const auto& image = packet.Get(); + EXPECT_EQ(image.width(), kImageWidth); + EXPECT_EQ(image.height(), kImageHeight); + auto image_frame = image.GetImageFrameSharedPtr(); + auto* pixel_data = image_frame->PixelData(); + EXPECT_EQ(pixel_data[0], 100); + EXPECT_EQ(pixel_data[1], 200); + EXPECT_EQ(pixel_data[2], 255); + } +} + +TEST(FlatColorImageCalculatorTest, SpecifyDimensionThroughOptions) { + CalculatorRunner runner(R"pb( + calculator: "FlatColorImageCalculator" + input_stream: "COLOR:color" + output_stream: "IMAGE:out_image" + options { + [mediapipe.FlatColorImageCalculatorOptions.ext] { + output_width: 7, + output_height: 13, + } + } + )pb"); + + Color color; + color.set_r(0); + color.set_g(5); + color.set_b(0); + + for (int ts = 0; ts < 3; ++ts) { + runner.MutableInputs()->Tag(kColorTag).packets.push_back( + MakePacket(color).At(Timestamp(ts))); + } + MP_ASSERT_OK(runner.Run()); + + const auto& outputs = runner.Outputs().Tag(kImageTag).packets; + ASSERT_EQ(outputs.size(), 3); + + for (const auto& packet : outputs) { + const auto& image = packet.Get(); + EXPECT_EQ(image.width(), 7); + EXPECT_EQ(image.height(), 13); + auto image_frame = image.GetImageFrameSharedPtr(); + const uint8_t* pixel_data = image_frame->PixelData(); + EXPECT_EQ(pixel_data[0], 0); + EXPECT_EQ(pixel_data[1], 5); + EXPECT_EQ(pixel_data[2], 0); + } +} + +TEST(FlatColorImageCalculatorTest, FailureMissingDimension) { + CalculatorRunner runner(R"pb( + calculator: "FlatColorImageCalculator" + input_stream: "COLOR:color" + output_stream: "IMAGE:out_image" + )pb"); + + Color color; + color.set_r(0); + color.set_g(5); + color.set_b(0); + + for (int ts = 0; ts < 3; ++ts) { + runner.MutableInputs()->Tag(kColorTag).packets.push_back( + MakePacket(color).At(Timestamp(ts))); + } + ASSERT_THAT(runner.Run().message(), + HasSubstr("Either set IMAGE input stream")); +} + +TEST(FlatColorImageCalculatorTest, FailureMissingColor) { + CalculatorRunner runner(R"pb( + calculator: "FlatColorImageCalculator" + input_stream: "IMAGE:image" + output_stream: "IMAGE:out_image" + )pb"); + + auto image_frame = std::make_shared(ImageFormat::SRGB, + kImageWidth, kImageHeight); + + for (int ts = 0; ts < 3; ++ts) { + runner.MutableInputs()->Tag(kImageTag).packets.push_back( + MakePacket(image_frame).At(Timestamp(ts))); + } + ASSERT_THAT(runner.Run().message(), + HasSubstr("Either set COLOR input stream")); +} + +TEST(FlatColorImageCalculatorTest, FailureDuplicateDimension) { + CalculatorRunner runner(R"pb( + calculator: "FlatColorImageCalculator" + input_stream: "IMAGE:image" + input_stream: "COLOR:color" + output_stream: "IMAGE:out_image" + options { + [mediapipe.FlatColorImageCalculatorOptions.ext] { + output_width: 7, + output_height: 13, + } + } + )pb"); + + auto image_frame = std::make_shared(ImageFormat::SRGB, + kImageWidth, kImageHeight); + + for (int ts = 0; ts < 3; ++ts) { + runner.MutableInputs()->Tag(kImageTag).packets.push_back( + MakePacket(image_frame).At(Timestamp(ts))); + } + ASSERT_THAT(runner.Run().message(), + HasSubstr("Either set IMAGE input stream")); +} + +TEST(FlatColorImageCalculatorTest, FailureDuplicateColor) { + CalculatorRunner runner(R"pb( + calculator: "FlatColorImageCalculator" + input_stream: "IMAGE:image" + input_stream: "COLOR:color" + output_stream: "IMAGE:out_image" + options { + [mediapipe.FlatColorImageCalculatorOptions.ext] { + color: { + r: 100, + g: 200, + b: 255, + } + } + } + )pb"); + + Color color; + color.set_r(0); + color.set_g(5); + color.set_b(0); + + for (int ts = 0; ts < 3; ++ts) { + runner.MutableInputs()->Tag(kColorTag).packets.push_back( + MakePacket(color).At(Timestamp(ts))); + } + ASSERT_THAT(runner.Run().message(), + HasSubstr("Either set COLOR input stream")); +} + +} // namespace +} // namespace mediapipe