Project import generated by Copybara.
GitOrigin-RevId: 08c2016a4df5aef571b464a4d4491f38c6b2af10
This commit is contained in:
@@ -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];
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user