diff --git a/mediapipe/calculators/util/BUILD b/mediapipe/calculators/util/BUILD index 22e6b073..b6f50b84 100644 --- a/mediapipe/calculators/util/BUILD +++ b/mediapipe/calculators/util/BUILD @@ -1285,12 +1285,14 @@ cc_library( srcs = ["flat_color_image_calculator.cc"], deps = [ ":flat_color_image_calculator_cc_proto", + "//mediapipe/framework:calculator_contract", "//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/framework/port:ret_check", "//mediapipe/util:color_cc_proto", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", diff --git a/mediapipe/calculators/util/flat_color_image_calculator.cc b/mediapipe/calculators/util/flat_color_image_calculator.cc index 71d3582c..f3b9c184 100644 --- a/mediapipe/calculators/util/flat_color_image_calculator.cc +++ b/mediapipe/calculators/util/flat_color_image_calculator.cc @@ -15,14 +15,13 @@ #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_contract.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 { @@ -32,6 +31,7 @@ namespace { using ::mediapipe::api2::Input; using ::mediapipe::api2::Node; using ::mediapipe::api2::Output; +using ::mediapipe::api2::SideOutput; } // namespace // A calculator for generating an image filled with a single color. @@ -45,7 +45,8 @@ using ::mediapipe::api2::Output; // // Outputs: // IMAGE (Image) -// Image filled with the requested color. +// Image filled with the requested color. Can be either an output_stream +// or an output_side_packet. // // Example useage: // node { @@ -68,9 +69,10 @@ class FlatColorImageCalculator : public Node { public: static constexpr Input::Optional kInImage{"IMAGE"}; static constexpr Input::Optional kInColor{"COLOR"}; - static constexpr Output kOutImage{"IMAGE"}; + static constexpr Output::Optional kOutImage{"IMAGE"}; + static constexpr SideOutput::Optional kOutSideImage{"IMAGE"}; - MEDIAPIPE_NODE_CONTRACT(kInImage, kInColor, kOutImage); + MEDIAPIPE_NODE_CONTRACT(kInImage, kInColor, kOutImage, kOutSideImage); static absl::Status UpdateContract(CalculatorContract* cc) { const auto& options = cc->Options(); @@ -81,6 +83,13 @@ class FlatColorImageCalculator : public Node { RET_CHECK(kInColor(cc).IsConnected() ^ options.has_color()) << "Either set COLOR input stream, or set through options"; + RET_CHECK(kOutImage(cc).IsConnected() ^ kOutSideImage(cc).IsConnected()) + << "Set IMAGE either as output stream, or as output side packet"; + + RET_CHECK(!kOutSideImage(cc).IsConnected() || + (options.has_output_height() && options.has_output_width())) + << "Set size through options, when setting IMAGE as output side packet"; + return absl::OkStatus(); } @@ -88,6 +97,9 @@ class FlatColorImageCalculator : public Node { absl::Status Process(CalculatorContext* cc) override; private: + std::optional> CreateOutputFrame( + CalculatorContext* cc); + bool use_dimension_from_option_ = false; bool use_color_from_option_ = false; }; @@ -96,10 +108,31 @@ MEDIAPIPE_REGISTER_NODE(FlatColorImageCalculator); absl::Status FlatColorImageCalculator::Open(CalculatorContext* cc) { use_dimension_from_option_ = !kInImage(cc).IsConnected(); use_color_from_option_ = !kInColor(cc).IsConnected(); + + if (!kOutImage(cc).IsConnected()) { + std::optional> output_frame = + CreateOutputFrame(cc); + if (output_frame.has_value()) { + kOutSideImage(cc).Set(Image(output_frame.value())); + } + } return absl::OkStatus(); } absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) { + if (kOutImage(cc).IsConnected()) { + std::optional> output_frame = + CreateOutputFrame(cc); + if (output_frame.has_value()) { + kOutImage(cc).Send(Image(output_frame.value())); + } + } + + return absl::OkStatus(); +} + +std::optional> +FlatColorImageCalculator::CreateOutputFrame(CalculatorContext* cc) { const auto& options = cc->Options(); int output_height = -1; @@ -112,7 +145,7 @@ absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) { output_height = input_image.height(); output_width = input_image.width(); } else { - return absl::OkStatus(); + return std::nullopt; } Color color; @@ -121,7 +154,7 @@ absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) { } else if (!kInColor(cc).IsEmpty()) { color = kInColor(cc).Get(); } else { - return absl::OkStatus(); + return std::nullopt; } auto output_frame = std::make_shared(ImageFormat::SRGB, @@ -130,9 +163,7 @@ absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) { output_mat.setTo(cv::Scalar(color.r(), color.g(), color.b())); - kOutImage(cc).Send(Image(output_frame)); - - return absl::OkStatus(); + return output_frame; } } // namespace mediapipe diff --git a/mediapipe/calculators/util/flat_color_image_calculator_test.cc b/mediapipe/calculators/util/flat_color_image_calculator_test.cc index 53c6de1b..c09064bf 100644 --- a/mediapipe/calculators/util/flat_color_image_calculator_test.cc +++ b/mediapipe/calculators/util/flat_color_image_calculator_test.cc @@ -113,6 +113,35 @@ TEST(FlatColorImageCalculatorTest, SpecifyDimensionThroughOptions) { } } +TEST(FlatColorImageCalculatorTest, ProducesOutputSidePacket) { + CalculatorRunner runner(R"pb( + calculator: "FlatColorImageCalculator" + output_side_packet: "IMAGE:out_packet" + options { + [mediapipe.FlatColorImageCalculatorOptions.ext] { + output_width: 1 + output_height: 1 + color: { + r: 100, + g: 200, + b: 255, + } + } + } + )pb"); + + MP_ASSERT_OK(runner.Run()); + + const auto& image = runner.OutputSidePackets().Tag(kImageTag).Get(); + EXPECT_EQ(image.width(), 1); + EXPECT_EQ(image.height(), 1); + auto image_frame = image.GetImageFrameSharedPtr(); + const uint8_t* 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, FailureMissingDimension) { CalculatorRunner runner(R"pb( calculator: "FlatColorImageCalculator" @@ -206,5 +235,56 @@ TEST(FlatColorImageCalculatorTest, FailureDuplicateColor) { HasSubstr("Either set COLOR input stream")); } +TEST(FlatColorImageCalculatorTest, FailureDuplicateOutputs) { + CalculatorRunner runner(R"pb( + calculator: "FlatColorImageCalculator" + output_stream: "IMAGE:out_image" + output_side_packet: "IMAGE:out_packet" + options { + [mediapipe.FlatColorImageCalculatorOptions.ext] { + output_width: 1 + output_height: 1 + color: { + r: 100, + g: 200, + b: 255, + } + } + } + )pb"); + + ASSERT_THAT( + runner.Run().message(), + HasSubstr("Set IMAGE either as output stream, or as output side packet")); +} + +TEST(FlatColorImageCalculatorTest, FailureSettingInputImageOnOutputSidePacket) { + CalculatorRunner runner(R"pb( + calculator: "FlatColorImageCalculator" + input_stream: "IMAGE:image" + output_side_packet: "IMAGE:out_packet" + 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))); + } + ASSERT_THAT(runner.Run().message(), + HasSubstr("Set size through options, when setting IMAGE as " + "output side packet")); +} + } // namespace } // namespace mediapipe