internal change
PiperOrigin-RevId: 528719459
This commit is contained in:
parent
3719aaef7e
commit
5b93477589
|
@ -1285,12 +1285,14 @@ cc_library(
|
||||||
srcs = ["flat_color_image_calculator.cc"],
|
srcs = ["flat_color_image_calculator.cc"],
|
||||||
deps = [
|
deps = [
|
||||||
":flat_color_image_calculator_cc_proto",
|
":flat_color_image_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_contract",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/api2:node",
|
"//mediapipe/framework/api2:node",
|
||||||
"//mediapipe/framework/formats:image",
|
"//mediapipe/framework/formats:image",
|
||||||
"//mediapipe/framework/formats:image_frame",
|
"//mediapipe/framework/formats:image_frame",
|
||||||
"//mediapipe/framework/formats:image_frame_opencv",
|
"//mediapipe/framework/formats:image_frame_opencv",
|
||||||
"//mediapipe/framework/port:opencv_core",
|
"//mediapipe/framework/port:opencv_core",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/util:color_cc_proto",
|
"//mediapipe/util:color_cc_proto",
|
||||||
"@com_google_absl//absl/status",
|
"@com_google_absl//absl/status",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
|
|
|
@ -15,14 +15,13 @@
|
||||||
#include <memory>
|
#include <memory>
|
||||||
|
|
||||||
#include "absl/status/status.h"
|
#include "absl/status/status.h"
|
||||||
#include "absl/strings/str_cat.h"
|
|
||||||
#include "mediapipe/calculators/util/flat_color_image_calculator.pb.h"
|
#include "mediapipe/calculators/util/flat_color_image_calculator.pb.h"
|
||||||
#include "mediapipe/framework/api2/node.h"
|
#include "mediapipe/framework/api2/node.h"
|
||||||
|
#include "mediapipe/framework/calculator_contract.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/formats/image.h"
|
#include "mediapipe/framework/formats/image.h"
|
||||||
#include "mediapipe/framework/formats/image_frame.h"
|
#include "mediapipe/framework/formats/image_frame.h"
|
||||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||||
#include "mediapipe/framework/port/opencv_core_inc.h"
|
|
||||||
#include "mediapipe/util/color.pb.h"
|
#include "mediapipe/util/color.pb.h"
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
@ -32,6 +31,7 @@ namespace {
|
||||||
using ::mediapipe::api2::Input;
|
using ::mediapipe::api2::Input;
|
||||||
using ::mediapipe::api2::Node;
|
using ::mediapipe::api2::Node;
|
||||||
using ::mediapipe::api2::Output;
|
using ::mediapipe::api2::Output;
|
||||||
|
using ::mediapipe::api2::SideOutput;
|
||||||
} // namespace
|
} // namespace
|
||||||
|
|
||||||
// A calculator for generating an image filled with a single color.
|
// A calculator for generating an image filled with a single color.
|
||||||
|
@ -45,7 +45,8 @@ using ::mediapipe::api2::Output;
|
||||||
//
|
//
|
||||||
// Outputs:
|
// Outputs:
|
||||||
// IMAGE (Image)
|
// 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:
|
// Example useage:
|
||||||
// node {
|
// node {
|
||||||
|
@ -68,9 +69,10 @@ class FlatColorImageCalculator : public Node {
|
||||||
public:
|
public:
|
||||||
static constexpr Input<Image>::Optional kInImage{"IMAGE"};
|
static constexpr Input<Image>::Optional kInImage{"IMAGE"};
|
||||||
static constexpr Input<Color>::Optional kInColor{"COLOR"};
|
static constexpr Input<Color>::Optional kInColor{"COLOR"};
|
||||||
static constexpr Output<Image> kOutImage{"IMAGE"};
|
static constexpr Output<Image>::Optional kOutImage{"IMAGE"};
|
||||||
|
static constexpr SideOutput<Image>::Optional kOutSideImage{"IMAGE"};
|
||||||
|
|
||||||
MEDIAPIPE_NODE_CONTRACT(kInImage, kInColor, kOutImage);
|
MEDIAPIPE_NODE_CONTRACT(kInImage, kInColor, kOutImage, kOutSideImage);
|
||||||
|
|
||||||
static absl::Status UpdateContract(CalculatorContract* cc) {
|
static absl::Status UpdateContract(CalculatorContract* cc) {
|
||||||
const auto& options = cc->Options<FlatColorImageCalculatorOptions>();
|
const auto& options = cc->Options<FlatColorImageCalculatorOptions>();
|
||||||
|
@ -81,6 +83,13 @@ class FlatColorImageCalculator : public Node {
|
||||||
RET_CHECK(kInColor(cc).IsConnected() ^ options.has_color())
|
RET_CHECK(kInColor(cc).IsConnected() ^ options.has_color())
|
||||||
<< "Either set COLOR input stream, or set through options";
|
<< "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();
|
return absl::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -88,6 +97,9 @@ class FlatColorImageCalculator : public Node {
|
||||||
absl::Status Process(CalculatorContext* cc) override;
|
absl::Status Process(CalculatorContext* cc) override;
|
||||||
|
|
||||||
private:
|
private:
|
||||||
|
std::optional<std::shared_ptr<ImageFrame>> CreateOutputFrame(
|
||||||
|
CalculatorContext* cc);
|
||||||
|
|
||||||
bool use_dimension_from_option_ = false;
|
bool use_dimension_from_option_ = false;
|
||||||
bool use_color_from_option_ = false;
|
bool use_color_from_option_ = false;
|
||||||
};
|
};
|
||||||
|
@ -96,10 +108,31 @@ MEDIAPIPE_REGISTER_NODE(FlatColorImageCalculator);
|
||||||
absl::Status FlatColorImageCalculator::Open(CalculatorContext* cc) {
|
absl::Status FlatColorImageCalculator::Open(CalculatorContext* cc) {
|
||||||
use_dimension_from_option_ = !kInImage(cc).IsConnected();
|
use_dimension_from_option_ = !kInImage(cc).IsConnected();
|
||||||
use_color_from_option_ = !kInColor(cc).IsConnected();
|
use_color_from_option_ = !kInColor(cc).IsConnected();
|
||||||
|
|
||||||
|
if (!kOutImage(cc).IsConnected()) {
|
||||||
|
std::optional<std::shared_ptr<ImageFrame>> output_frame =
|
||||||
|
CreateOutputFrame(cc);
|
||||||
|
if (output_frame.has_value()) {
|
||||||
|
kOutSideImage(cc).Set(Image(output_frame.value()));
|
||||||
|
}
|
||||||
|
}
|
||||||
return absl::OkStatus();
|
return absl::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) {
|
absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) {
|
||||||
|
if (kOutImage(cc).IsConnected()) {
|
||||||
|
std::optional<std::shared_ptr<ImageFrame>> output_frame =
|
||||||
|
CreateOutputFrame(cc);
|
||||||
|
if (output_frame.has_value()) {
|
||||||
|
kOutImage(cc).Send(Image(output_frame.value()));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return absl::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
std::optional<std::shared_ptr<ImageFrame>>
|
||||||
|
FlatColorImageCalculator::CreateOutputFrame(CalculatorContext* cc) {
|
||||||
const auto& options = cc->Options<FlatColorImageCalculatorOptions>();
|
const auto& options = cc->Options<FlatColorImageCalculatorOptions>();
|
||||||
|
|
||||||
int output_height = -1;
|
int output_height = -1;
|
||||||
|
@ -112,7 +145,7 @@ absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) {
|
||||||
output_height = input_image.height();
|
output_height = input_image.height();
|
||||||
output_width = input_image.width();
|
output_width = input_image.width();
|
||||||
} else {
|
} else {
|
||||||
return absl::OkStatus();
|
return std::nullopt;
|
||||||
}
|
}
|
||||||
|
|
||||||
Color color;
|
Color color;
|
||||||
|
@ -121,7 +154,7 @@ absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) {
|
||||||
} else if (!kInColor(cc).IsEmpty()) {
|
} else if (!kInColor(cc).IsEmpty()) {
|
||||||
color = kInColor(cc).Get();
|
color = kInColor(cc).Get();
|
||||||
} else {
|
} else {
|
||||||
return absl::OkStatus();
|
return std::nullopt;
|
||||||
}
|
}
|
||||||
|
|
||||||
auto output_frame = std::make_shared<ImageFrame>(ImageFormat::SRGB,
|
auto output_frame = std::make_shared<ImageFrame>(ImageFormat::SRGB,
|
||||||
|
@ -130,9 +163,7 @@ absl::Status FlatColorImageCalculator::Process(CalculatorContext* cc) {
|
||||||
|
|
||||||
output_mat.setTo(cv::Scalar(color.r(), color.g(), color.b()));
|
output_mat.setTo(cv::Scalar(color.r(), color.g(), color.b()));
|
||||||
|
|
||||||
kOutImage(cc).Send(Image(output_frame));
|
return output_frame;
|
||||||
|
|
||||||
return absl::OkStatus();
|
|
||||||
}
|
}
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|
|
@ -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<Image>();
|
||||||
|
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) {
|
TEST(FlatColorImageCalculatorTest, FailureMissingDimension) {
|
||||||
CalculatorRunner runner(R"pb(
|
CalculatorRunner runner(R"pb(
|
||||||
calculator: "FlatColorImageCalculator"
|
calculator: "FlatColorImageCalculator"
|
||||||
|
@ -206,5 +235,56 @@ TEST(FlatColorImageCalculatorTest, FailureDuplicateColor) {
|
||||||
HasSubstr("Either set COLOR input stream"));
|
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<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("Set size through options, when setting IMAGE as "
|
||||||
|
"output side packet"));
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|
Loading…
Reference in New Issue
Block a user