diff --git a/mediapipe/calculators/image/BUILD b/mediapipe/calculators/image/BUILD index d627bdc4..7b94a031 100644 --- a/mediapipe/calculators/image/BUILD +++ b/mediapipe/calculators/image/BUILD @@ -156,6 +156,7 @@ cc_library( "//mediapipe/framework/port:opencv_core", "//mediapipe/framework/port:status", "//mediapipe/framework/port:vector", + "//mediapipe/framework/port:opencv_imgproc", ] + select({ "//mediapipe/gpu:disable_gpu": [], "//conditions:default": [ @@ -168,6 +169,25 @@ cc_library( alwayslink = 1, ) +cc_test( + name = "set_alpha_calculator_test", + srcs = ["set_alpha_calculator_test.cc"], + deps = [ + ":set_alpha_calculator", + ":set_alpha_calculator_cc_proto", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework:calculator_runner", + "//mediapipe/framework/formats:image_frame_opencv", + "//mediapipe/framework/formats:rect_cc_proto", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:opencv_core", + "//mediapipe/framework/port:opencv_imgproc", + "//mediapipe/framework/port:parse_text_proto", + "//mediapipe/framework/port:status", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "bilateral_filter_calculator", srcs = ["bilateral_filter_calculator.cc"], diff --git a/mediapipe/calculators/image/set_alpha_calculator.cc b/mediapipe/calculators/image/set_alpha_calculator.cc index ea29278d..e20621e8 100644 --- a/mediapipe/calculators/image/set_alpha_calculator.cc +++ b/mediapipe/calculators/image/set_alpha_calculator.cc @@ -22,6 +22,7 @@ #include "mediapipe/framework/formats/image_frame_opencv.h" #include "mediapipe/framework/port/logging.h" #include "mediapipe/framework/port/opencv_core_inc.h" +#include "mediapipe/framework/port/opencv_imgproc_inc.h" #include "mediapipe/framework/port/status.h" #include "mediapipe/framework/port/vector.h" @@ -53,24 +54,16 @@ enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES }; // range of [0, 1). Only the first channel of Alpha is used. Input & output Mat // must be uchar. template -absl::Status MergeRGBA8Image(const cv::Mat input_mat, const cv::Mat& alpha_mat, - cv::Mat& output_mat) { - RET_CHECK_EQ(input_mat.rows, alpha_mat.rows); - RET_CHECK_EQ(input_mat.cols, alpha_mat.cols); - RET_CHECK_EQ(input_mat.rows, output_mat.rows); - RET_CHECK_EQ(input_mat.cols, output_mat.cols); +absl::Status CopyAlphaImage(const cv::Mat& alpha_mat, cv::Mat& output_mat) { + RET_CHECK_EQ(output_mat.rows, alpha_mat.rows); + RET_CHECK_EQ(output_mat.cols, alpha_mat.cols); for (int i = 0; i < output_mat.rows; ++i) { - const uchar* in_ptr = input_mat.ptr(i); const AlphaType* alpha_ptr = alpha_mat.ptr(i); uchar* out_ptr = output_mat.ptr(i); for (int j = 0; j < output_mat.cols; ++j) { const int out_idx = j * kNumChannelsRGBA; - const int in_idx = j * input_mat.channels(); const int alpha_idx = j * alpha_mat.channels(); - out_ptr[out_idx + 0] = in_ptr[in_idx + 0]; - out_ptr[out_idx + 1] = in_ptr[in_idx + 1]; - out_ptr[out_idx + 2] = in_ptr[in_idx + 2]; if constexpr (std::is_same::value) { out_ptr[out_idx + 3] = alpha_ptr[alpha_idx + 0]; // channel 0 of mask } else { @@ -273,7 +266,7 @@ absl::Status SetAlphaCalculator::RenderCpu(CalculatorContext* cc) { // Setup source image const auto& input_frame = cc->Inputs().Tag(kInputFrameTag).Get(); - const cv::Mat input_mat = mediapipe::formats::MatView(&input_frame); + const cv::Mat input_mat = formats::MatView(&input_frame); if (!(input_mat.type() == CV_8UC3 || input_mat.type() == CV_8UC4)) { LOG(ERROR) << "Only 3 or 4 channel 8-bit input image supported"; } @@ -281,38 +274,38 @@ absl::Status SetAlphaCalculator::RenderCpu(CalculatorContext* cc) { // Setup destination image auto output_frame = absl::make_unique( ImageFormat::SRGBA, input_mat.cols, input_mat.rows); - cv::Mat output_mat = mediapipe::formats::MatView(output_frame.get()); + cv::Mat output_mat = formats::MatView(output_frame.get()); const bool has_alpha_mask = cc->Inputs().HasTag(kInputAlphaTag) && !cc->Inputs().Tag(kInputAlphaTag).IsEmpty(); const bool use_alpha_mask = alpha_value_ < 0 && has_alpha_mask; - // Setup alpha image and Update image in CPU. + // Copy rgb part of the image in CPU + if (input_mat.channels() == 3) { + cv::cvtColor(input_mat, output_mat, cv::COLOR_RGB2RGBA); + } else { + input_mat.copyTo(output_mat); + } + + // Setup alpha image in CPU. if (use_alpha_mask) { const auto& alpha_mask = cc->Inputs().Tag(kInputAlphaTag).Get(); - cv::Mat alpha_mat = mediapipe::formats::MatView(&alpha_mask); + cv::Mat alpha_mat = formats::MatView(&alpha_mask); const bool alpha_is_float = CV_MAT_DEPTH(alpha_mat.type()) == CV_32F; RET_CHECK(alpha_is_float || CV_MAT_DEPTH(alpha_mat.type()) == CV_8U); if (alpha_is_float) { - MP_RETURN_IF_ERROR( - MergeRGBA8Image(input_mat, alpha_mat, output_mat)); + MP_RETURN_IF_ERROR(CopyAlphaImage(alpha_mat, output_mat)); } else { - MP_RETURN_IF_ERROR( - MergeRGBA8Image(input_mat, alpha_mat, output_mat)); + MP_RETURN_IF_ERROR(CopyAlphaImage(alpha_mat, output_mat)); } } else { const uchar alpha_value = std::min(std::max(0.0f, alpha_value_), 255.0f); for (int i = 0; i < output_mat.rows; ++i) { - const uchar* in_ptr = input_mat.ptr(i); uchar* out_ptr = output_mat.ptr(i); for (int j = 0; j < output_mat.cols; ++j) { const int out_idx = j * kNumChannelsRGBA; - const int in_idx = j * input_mat.channels(); - out_ptr[out_idx + 0] = in_ptr[in_idx + 0]; - out_ptr[out_idx + 1] = in_ptr[in_idx + 1]; - out_ptr[out_idx + 2] = in_ptr[in_idx + 2]; out_ptr[out_idx + 3] = alpha_value; // use value from options } } diff --git a/mediapipe/calculators/image/set_alpha_calculator_test.cc b/mediapipe/calculators/image/set_alpha_calculator_test.cc new file mode 100644 index 00000000..cb2352d0 --- /dev/null +++ b/mediapipe/calculators/image/set_alpha_calculator_test.cc @@ -0,0 +1,156 @@ +#include + +#include "mediapipe/calculators/image/set_alpha_calculator.pb.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/calculator_runner.h" +#include "mediapipe/framework/formats/image_frame_opencv.h" +#include "mediapipe/framework/formats/rect.pb.h" +#include "mediapipe/framework/port/gtest.h" +#include "mediapipe/framework/port/opencv_core_inc.h" +#include "mediapipe/framework/port/opencv_imgproc_inc.h" +#include "mediapipe/framework/port/parse_text_proto.h" +#include "mediapipe/framework/port/status_matchers.h" +#include "testing/base/public/benchmark.h" + +namespace mediapipe { + +namespace { + +constexpr int input_width = 100; +constexpr int input_height = 100; + +std::unique_ptr GetInputFrame(int width, int height, int channel) { + const int total_size = width * height * channel; + + ImageFormat::Format image_format; + if (channel == 4) { + image_format = ImageFormat::SRGBA; + } else if (channel == 3) { + image_format = ImageFormat::SRGB; + } else { + image_format = ImageFormat::GRAY8; + } + + auto input_frame = std::make_unique(image_format, width, height, + /*alignment_boundary =*/1); + for (int i = 0; i < total_size; ++i) { + input_frame->MutablePixelData()[i] = i % 256; + } + return input_frame; +} + +// Test SetAlphaCalculator with RGB IMAGE input. +TEST(SetAlphaCalculatorTest, CpuRgb) { + auto calculator_node = ParseTextProtoOrDie( + R"pb( + calculator: "SetAlphaCalculator" + input_stream: "IMAGE:input_frames" + input_stream: "ALPHA:masks" + output_stream: "IMAGE:output_frames" + )pb"); + CalculatorRunner runner(calculator_node); + + // Input frames. + const auto input_frame = GetInputFrame(input_width, input_height, 3); + const auto mask_frame = GetInputFrame(input_width, input_height, 1); + auto input_frame_packet = MakePacket(std::move(*input_frame)); + auto mask_frame_packet = MakePacket(std::move(*mask_frame)); + runner.MutableInputs()->Tag("IMAGE").packets.push_back( + input_frame_packet.At(Timestamp(1))); + runner.MutableInputs()->Tag("ALPHA").packets.push_back( + mask_frame_packet.At(Timestamp(1))); + + MP_ASSERT_OK(runner.Run()); + + const auto& outputs = runner.Outputs(); + EXPECT_EQ(outputs.NumEntries(), 1); + const auto& output_image = outputs.Tag("IMAGE").packets[0].Get(); + + // Generate ground truth (expected_mat). + const auto image = GetInputFrame(input_width, input_height, 3); + const auto input_mat = formats::MatView(image.get()); + const auto mask = GetInputFrame(input_width, input_height, 1); + const auto mask_mat = formats::MatView(mask.get()); + const std::array input_mats = {input_mat, mask_mat}; + cv::Mat expected_mat(input_width, input_height, CV_8UC4); + cv::mixChannels(input_mats, {expected_mat}, {0, 0, 1, 1, 2, 2, 3, 3}); + + cv::Mat output_mat = formats::MatView(&output_image); + double max_diff = cv::norm(expected_mat, output_mat, cv::NORM_INF); + EXPECT_FLOAT_EQ(max_diff, 0); +} // TEST + +// Test SetAlphaCalculator with RGBA IMAGE input. +TEST(SetAlphaCalculatorTest, CpuRgba) { + auto calculator_node = ParseTextProtoOrDie( + R"pb( + calculator: "SetAlphaCalculator" + input_stream: "IMAGE:input_frames" + input_stream: "ALPHA:masks" + output_stream: "IMAGE:output_frames" + )pb"); + CalculatorRunner runner(calculator_node); + + // Input frames. + const auto input_frame = GetInputFrame(input_width, input_height, 4); + const auto mask_frame = GetInputFrame(input_width, input_height, 1); + auto input_frame_packet = MakePacket(std::move(*input_frame)); + auto mask_frame_packet = MakePacket(std::move(*mask_frame)); + runner.MutableInputs()->Tag("IMAGE").packets.push_back( + input_frame_packet.At(Timestamp(1))); + runner.MutableInputs()->Tag("ALPHA").packets.push_back( + mask_frame_packet.At(Timestamp(1))); + + MP_ASSERT_OK(runner.Run()); + + const auto& outputs = runner.Outputs(); + EXPECT_EQ(outputs.NumEntries(), 1); + const auto& output_image = outputs.Tag("IMAGE").packets[0].Get(); + + // Generate ground truth (expected_mat). + const auto image = GetInputFrame(input_width, input_height, 4); + const auto input_mat = formats::MatView(image.get()); + const auto mask = GetInputFrame(input_width, input_height, 1); + const auto mask_mat = formats::MatView(mask.get()); + const std::array input_mats = {input_mat, mask_mat}; + cv::Mat expected_mat(input_width, input_height, CV_8UC4); + cv::mixChannels(input_mats, {expected_mat}, {0, 0, 1, 1, 2, 2, 4, 3}); + + cv::Mat output_mat = formats::MatView(&output_image); + double max_diff = cv::norm(expected_mat, output_mat, cv::NORM_INF); + EXPECT_FLOAT_EQ(max_diff, 0); +} // TEST + +static void BM_SetAlpha3ChannelImage(benchmark::State& state) { + auto calculator_node = ParseTextProtoOrDie( + R"pb( + calculator: "SetAlphaCalculator" + input_stream: "IMAGE:input_frames" + input_stream: "ALPHA:masks" + output_stream: "IMAGE:output_frames" + )pb"); + CalculatorRunner runner(calculator_node); + + // Input frames. + const auto input_frame = GetInputFrame(input_width, input_height, 3); + const auto mask_frame = GetInputFrame(input_width, input_height, 1); + auto input_frame_packet = MakePacket(std::move(*input_frame)); + auto mask_frame_packet = MakePacket(std::move(*mask_frame)); + runner.MutableInputs()->Tag("IMAGE").packets.push_back( + input_frame_packet.At(Timestamp(1))); + runner.MutableInputs()->Tag("ALPHA").packets.push_back( + mask_frame_packet.At(Timestamp(1))); + + MP_ASSERT_OK(runner.Run()); + const auto& outputs = runner.Outputs(); + ASSERT_EQ(1, outputs.NumEntries()); + + for (const auto _ : state) { + MP_ASSERT_OK(runner.Run()); + } +} + +BENCHMARK(BM_SetAlpha3ChannelImage); + +} // namespace +} // namespace mediapipe