Internal change
PiperOrigin-RevId: 516607678
This commit is contained in:
parent
fef8b9cb58
commit
ed3e728bb8
|
@ -1270,6 +1270,50 @@ cc_library(
|
||||||
alwayslink = 1,
|
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(
|
cc_library(
|
||||||
name = "from_image_calculator",
|
name = "from_image_calculator",
|
||||||
srcs = ["from_image_calculator.cc"],
|
srcs = ["from_image_calculator.cc"],
|
||||||
|
|
138
mediapipe/calculators/util/flat_color_image_calculator.cc
Normal file
138
mediapipe/calculators/util/flat_color_image_calculator.cc
Normal file
|
@ -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
|
32
mediapipe/calculators/util/flat_color_image_calculator.proto
Normal file
32
mediapipe/calculators/util/flat_color_image_calculator.proto
Normal file
|
@ -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;
|
||||||
|
}
|
210
mediapipe/calculators/util/flat_color_image_calculator_test.cc
Normal file
210
mediapipe/calculators/util/flat_color_image_calculator_test.cc
Normal file
|
@ -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
|
Loading…
Reference in New Issue
Block a user