cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

internal change

Browse Source 
PiperOrigin-RevId: 528719459
This commit is contained in:
MediaPipe Team 2023-05-02 02:15:03 -07:00 committed by Copybara-Service
parent 3719aaef7e
commit 5b93477589
3 changed files with 123 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mediapipe/calculators/util/BUILD
View File

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

51
mediapipe/calculators/util/flat_color_image_calculator.cc
View File

@ -15,14 +15,13 @@
#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_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<Image>::Optional kInImage{"IMAGE"};
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) {
const auto& options = cc->Options<FlatColorImageCalculatorOptions>();
@ -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<std::shared_ptr<ImageFrame>> 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<std::shared_ptr<ImageFrame>> 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<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>();
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<ImageFrame>(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

80
mediapipe/calculators/util/flat_color_image_calculator_test.cc
View File

@ -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) {
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<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 mediapipe