Project import generated by Copybara.

GitOrigin-RevId: 08c2016a4df5aef571b464a4d4491f38c6b2af10
This commit is contained in:
MediaPipe Team
2021-06-03 17:04:35 -04:00
committed by chuoling
parent ae05ad04b3
commit 8b57bf879b
118 changed files with 3999 additions and 391 deletions
@@ -37,6 +37,22 @@ constexpr char kImageFrameTag[] = "IMAGE";
constexpr char kMaskCpuTag[] = "MASK";
constexpr char kGpuBufferTag[] = "IMAGE_GPU";
constexpr char kMaskGpuTag[] = "MASK_GPU";
inline cv::Vec3b Blend(const cv::Vec3b& color1, const cv::Vec3b& color2,
float weight, int invert_mask,
int adjust_with_luminance) {
weight = (1 - invert_mask) * weight + invert_mask * (1.0f - weight);
float luminance =
(1 - adjust_with_luminance) * 1.0f +
adjust_with_luminance *
(color1[0] * 0.299 + color1[1] * 0.587 + color1[2] * 0.114) / 255;
float mix_value = weight * luminance;
return color1 * (1.0 - mix_value) + color2 * mix_value;
}
} // namespace
namespace mediapipe {
@@ -44,15 +60,14 @@ namespace mediapipe {
// A calculator to recolor a masked area of an image to a specified color.
//
// A mask image is used to specify where to overlay a user defined color.
// The luminance of the input image is used to adjust the blending weight,
// to help preserve image textures.
//
// Inputs:
// One of the following IMAGE tags:
// IMAGE: An ImageFrame input image, RGB or RGBA.
// IMAGE: An ImageFrame input image in ImageFormat::SRGB.
// IMAGE_GPU: A GpuBuffer input image, RGBA.
// One of the following MASK tags:
// MASK: An ImageFrame input mask, Gray, RGB or RGBA.
// MASK: An ImageFrame input mask in ImageFormat::GRAY8, SRGB, SRGBA, or
// VEC32F1
// MASK_GPU: A GpuBuffer input mask, RGBA.
// Output:
// One of the following IMAGE tags:
@@ -98,10 +113,12 @@ class RecolorCalculator : public CalculatorBase {
void GlRender();
bool initialized_ = false;
std::vector<float> color_;
std::vector<uint8> color_;
mediapipe::RecolorCalculatorOptions::MaskChannel mask_channel_;
bool use_gpu_ = false;
bool invert_mask_ = false;
bool adjust_with_luminance_ = false;
#if !MEDIAPIPE_DISABLE_GPU
mediapipe::GlCalculatorHelper gpu_helper_;
GLuint program_ = 0;
@@ -233,11 +250,15 @@ absl::Status RecolorCalculator::RenderCpu(CalculatorContext* cc) {
}
cv::Mat mask_full;
cv::resize(mask_mat, mask_full, input_mat.size());
const cv::Vec3b recolor = {color_[0], color_[1], color_[2]};
auto output_img = absl::make_unique<ImageFrame>(
input_img.Format(), input_mat.cols, input_mat.rows);
cv::Mat output_mat = mediapipe::formats::MatView(output_img.get());
const int invert_mask = invert_mask_ ? 1 : 0;
const int adjust_with_luminance = adjust_with_luminance_ ? 1 : 0;
// From GPU shader:
/*
vec4 weight = texture2D(mask, sample_coordinate);
@@ -249,18 +270,23 @@ absl::Status RecolorCalculator::RenderCpu(CalculatorContext* cc) {
fragColor = mix(color1, color2, mix_value);
*/
for (int i = 0; i < output_mat.rows; ++i) {
for (int j = 0; j < output_mat.cols; ++j) {
float weight = mask_full.at<uchar>(i, j) * (1.0 / 255.0);
cv::Vec3f color1 = input_mat.at<cv::Vec3b>(i, j);
cv::Vec3f color2 = {color_[0], color_[1], color_[2]};
float luminance =
(color1[0] * 0.299 + color1[1] * 0.587 + color1[2] * 0.114) / 255;
float mix_value = weight * luminance;
cv::Vec3b mix_color = color1 * (1.0 - mix_value) + color2 * mix_value;
output_mat.at<cv::Vec3b>(i, j) = mix_color;
if (mask_img.Format() == ImageFormat::VEC32F1) {
for (int i = 0; i < output_mat.rows; ++i) {
for (int j = 0; j < output_mat.cols; ++j) {
const float weight = mask_full.at<float>(i, j);
output_mat.at<cv::Vec3b>(i, j) =
Blend(input_mat.at<cv::Vec3b>(i, j), recolor, weight, invert_mask,
adjust_with_luminance);
}
}
} else {
for (int i = 0; i < output_mat.rows; ++i) {
for (int j = 0; j < output_mat.cols; ++j) {
const float weight = mask_full.at<uchar>(i, j) * (1.0 / 255.0);
output_mat.at<cv::Vec3b>(i, j) =
Blend(input_mat.at<cv::Vec3b>(i, j), recolor, weight, invert_mask,
adjust_with_luminance);
}
}
}
@@ -385,6 +411,9 @@ absl::Status RecolorCalculator::LoadOptions(CalculatorContext* cc) {
color_.push_back(options.color().g());
color_.push_back(options.color().b());
invert_mask_ = options.invert_mask();
adjust_with_luminance_ = options.adjust_with_luminance();
return absl::OkStatus();
}
@@ -435,13 +464,20 @@ absl::Status RecolorCalculator::InitGpu(CalculatorContext* cc) {
uniform sampler2D frame;
uniform sampler2D mask;
uniform vec3 recolor;
uniform float invert_mask;
uniform float adjust_with_luminance;
void main() {
vec4 weight = texture2D(mask, sample_coordinate);
vec4 color1 = texture2D(frame, sample_coordinate);
vec4 color2 = vec4(recolor, 1.0);
float luminance = dot(color1.rgb, vec3(0.299, 0.587, 0.114));
weight = mix(weight, 1.0 - weight, invert_mask);
float luminance = mix(1.0,
dot(color1.rgb, vec3(0.299, 0.587, 0.114)),
adjust_with_luminance);
float mix_value = weight.MASK_COMPONENT * luminance;
fragColor = mix(color1, color2, mix_value);
@@ -458,6 +494,10 @@ absl::Status RecolorCalculator::InitGpu(CalculatorContext* cc) {
glUniform1i(glGetUniformLocation(program_, "mask"), 2);
glUniform3f(glGetUniformLocation(program_, "recolor"), color_[0] / 255.0,
color_[1] / 255.0, color_[2] / 255.0);
glUniform1f(glGetUniformLocation(program_, "invert_mask"),
invert_mask_ ? 1.0f : 0.0f);
glUniform1f(glGetUniformLocation(program_, "adjust_with_luminance"),
adjust_with_luminance_ ? 1.0f : 0.0f);
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
@@ -36,4 +36,11 @@ message RecolorCalculatorOptions {
// Color to blend into input image where mask is > 0.
// The blending is based on the input image luminosity.
optional Color color = 2;
// Swap the meaning of mask values for foreground/background.
optional bool invert_mask = 3 [default = false];
// Whether to use the luminance of the input image to further adjust the
// blending weight, to help preserve image textures.
optional bool adjust_with_luminance = 4 [default = true];
}