Internal change

PiperOrigin-RevId: 516607678

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"],

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 <memory>
+
+#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<Image>::Optional kInImage{"IMAGE"};
+  static constexpr Input<Color>::Optional kInColor{"COLOR"};
+  static constexpr Output<Image> kOutImage{"IMAGE"};
+
+  MEDIAPIPE_NODE_CONTRACT(kInImage, kInColor, kOutImage);
+
+  static absl::Status UpdateContract(CalculatorContract* cc) {
+    const auto& options = cc->Options<FlatColorImageCalculatorOptions>();
+
+    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<FlatColorImageCalculatorOptions>();
+
+  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<ImageFrame>(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

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;
+}

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 <memory>
+
+#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<ImageFrame>(ImageFormat::SRGB,
+                                                  kImageWidth, kImageHeight);
+
+  for (int ts = 0; ts < 3; ++ts) {
+    runner.MutableInputs()->Tag(kImageTag).packets.push_back(
+        MakePacket<Image>(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<Image>();
+    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>(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<Image>();
+    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>(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<ImageFrame>(ImageFormat::SRGB,
+                                                  kImageWidth, kImageHeight);
+
+  for (int ts = 0; ts < 3; ++ts) {
+    runner.MutableInputs()->Tag(kImageTag).packets.push_back(
+        MakePacket<Image>(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<ImageFrame>(ImageFormat::SRGB,
+                                                  kImageWidth, kImageHeight);
+
+  for (int ts = 0; ts < 3; ++ts) {
+    runner.MutableInputs()->Tag(kImageTag).packets.push_back(
+        MakePacket<Image>(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>(color).At(Timestamp(ts)));
+  }
+  ASSERT_THAT(runner.Run().message(),
+              HasSubstr("Either set COLOR input stream"));
+}
+
+}  // namespace
+}  // namespace mediapipe