talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 33adfdf31f3a5cbf9edc07ee1ea583e95080bdc5
This commit is contained in:
MediaPipe Team
2021-06-24 17:55:26 -04:00
committed by chuoling
parent b544a314b3
commit 139237092f
151 changed files with 4023 additions and 1001 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/core/BUILD
+5 -5
View File
@@ -933,8 +933,8 @@ cc_test(
)
cc_library(
name = "split_normalized_landmark_list_calculator",
srcs = ["split_normalized_landmark_list_calculator.cc"],
name = "split_landmarks_calculator",
srcs = ["split_landmarks_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
":split_vector_calculator_cc_proto",
@@ -948,10 +948,10 @@ cc_library(
)
cc_test(
name = "split_normalized_landmark_list_calculator_test",
srcs = ["split_normalized_landmark_list_calculator_test.cc"],
name = "split_landmarks_calculator_test",
srcs = ["split_landmarks_calculator_test.cc"],
deps = [
":split_normalized_landmark_list_calculator",
":split_landmarks_calculator",
":split_vector_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
mediapipe/calculators/core/split_normalized_landmark_list_calculator.cc → mediapipe/calculators/core/split_landmarks_calculator.cc
+29 -23
View File
@@ -12,8 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MEDIAPIPE_CALCULATORS_CORE_SPLIT_NORMALIZED_LANDMARK_LIST_CALCULATOR_H_ // NOLINT
#define MEDIAPIPE_CALCULATORS_CORE_SPLIT_NORMALIZED_LANDMARK_LIST_CALCULATOR_H_ // NOLINT
#ifndef MEDIAPIPE_CALCULATORS_CORE_SPLIT_LANDMARKS_CALCULATOR_H_ // NOLINT
#define MEDIAPIPE_CALCULATORS_CORE_SPLIT_LANDMARKS_CALCULATOR_H_ // NOLINT
#include "mediapipe/calculators/core/split_vector_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
@@ -24,29 +24,30 @@
namespace mediapipe {
// Splits an input packet with NormalizedLandmarkList into
// multiple NormalizedLandmarkList output packets using the [begin, end) ranges
// Splits an input packet with LandmarkListType into
// multiple LandmarkListType output packets using the [begin, end) ranges
// specified in SplitVectorCalculatorOptions. If the option "element_only" is
// set to true, all ranges should be of size 1 and all outputs will be elements
// of type NormalizedLandmark. If "element_only" is false, ranges can be
// non-zero in size and all outputs will be of type NormalizedLandmarkList.
// of type LandmarkType. If "element_only" is false, ranges can be
// non-zero in size and all outputs will be of type LandmarkListType.
// If the option "combine_outputs" is set to true, only one output stream can be
// specified and all ranges of elements will be combined into one
// NormalizedLandmarkList.
class SplitNormalizedLandmarkListCalculator : public CalculatorBase {
// LandmarkListType.
template <typename LandmarkType, typename LandmarkListType>
class SplitLandmarksCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
RET_CHECK(cc->Inputs().NumEntries() == 1);
RET_CHECK(cc->Outputs().NumEntries() != 0);
cc->Inputs().Index(0).Set<NormalizedLandmarkList>();
cc->Inputs().Index(0).Set<LandmarkListType>();
const auto& options =
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
if (options.combine_outputs()) {
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1);
cc->Outputs().Index(0).Set<NormalizedLandmarkList>();
cc->Outputs().Index(0).Set<LandmarkListType>();
for (int i = 0; i < options.ranges_size() - 1; ++i) {
for (int j = i + 1; j < options.ranges_size(); ++j) {
const auto& range_0 = options.ranges(i);
@@ -81,9 +82,9 @@ class SplitNormalizedLandmarkListCalculator : public CalculatorBase {
return absl::InvalidArgumentError(
"Since element_only is true, all ranges should be of size 1.");
}
cc->Outputs().Index(i).Set<NormalizedLandmark>();
cc->Outputs().Index(i).Set<LandmarkType>();
} else {
cc->Outputs().Index(i).Set<NormalizedLandmarkList>();
cc->Outputs().Index(i).Set<LandmarkListType>();
}
}
}
@@ -110,40 +111,39 @@ class SplitNormalizedLandmarkListCalculator : public CalculatorBase {
}
absl::Status Process(CalculatorContext* cc) override {
const NormalizedLandmarkList& input =
cc->Inputs().Index(0).Get<NormalizedLandmarkList>();
const LandmarkListType& input =
cc->Inputs().Index(0).Get<LandmarkListType>();
RET_CHECK_GE(input.landmark_size(), max_range_end_)
<< "Max range end " << max_range_end_ << " exceeds landmarks size "
<< input.landmark_size();
if (combine_outputs_) {
NormalizedLandmarkList output;
LandmarkListType output;
for (int i = 0; i < ranges_.size(); ++i) {
for (int j = ranges_[i].first; j < ranges_[i].second; ++j) {
const NormalizedLandmark& input_landmark = input.landmark(j);
const LandmarkType& input_landmark = input.landmark(j);
*output.add_landmark() = input_landmark;
}
}
RET_CHECK_EQ(output.landmark_size(), total_elements_);
cc->Outputs().Index(0).AddPacket(
MakePacket<NormalizedLandmarkList>(output).At(cc->InputTimestamp()));
MakePacket<LandmarkListType>(output).At(cc->InputTimestamp()));
} else {
if (element_only_) {
for (int i = 0; i < ranges_.size(); ++i) {
cc->Outputs().Index(i).AddPacket(
MakePacket<NormalizedLandmark>(input.landmark(ranges_[i].first))
MakePacket<LandmarkType>(input.landmark(ranges_[i].first))
.At(cc->InputTimestamp()));
}
} else {
for (int i = 0; i < ranges_.size(); ++i) {
NormalizedLandmarkList output;
LandmarkListType output;
for (int j = ranges_[i].first; j < ranges_[i].second; ++j) {
const NormalizedLandmark& input_landmark = input.landmark(j);
const LandmarkType& input_landmark = input.landmark(j);
*output.add_landmark() = input_landmark;
}
cc->Outputs().Index(i).AddPacket(
MakePacket<NormalizedLandmarkList>(output).At(
cc->InputTimestamp()));
MakePacket<LandmarkListType>(output).At(cc->InputTimestamp()));
}
}
}
@@ -159,9 +159,15 @@ class SplitNormalizedLandmarkListCalculator : public CalculatorBase {
bool combine_outputs_ = false;
};
typedef SplitLandmarksCalculator<NormalizedLandmark, NormalizedLandmarkList>
SplitNormalizedLandmarkListCalculator;
REGISTER_CALCULATOR(SplitNormalizedLandmarkListCalculator);
typedef SplitLandmarksCalculator<Landmark, LandmarkList>
SplitLandmarkListCalculator;
REGISTER_CALCULATOR(SplitLandmarkListCalculator);
} // namespace mediapipe
// NOLINTNEXTLINE
#endif // MEDIAPIPE_CALCULATORS_CORE_SPLIT_NORMALIZED_LANDMARK_LIST_CALCULATOR_H_
#endif // MEDIAPIPE_CALCULATORS_CORE_SPLIT_LANDMARKS_CALCULATOR_H_
mediapipe/calculators/core/split_normalized_landmark_list_calculator_test.cc → mediapipe/calculators/core/split_landmarks_calculator_test.cc
View File
mediapipe/calculators/image/BUILD
+59
View File
@@ -80,6 +80,16 @@ mediapipe_proto_library(
],
)
mediapipe_proto_library(
name = "segmentation_smoothing_calculator_proto",
srcs = ["segmentation_smoothing_calculator.proto"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
],
)
cc_library(
name = "color_convert_calculator",
srcs = ["color_convert_calculator.cc"],
@@ -602,3 +612,52 @@ cc_test(
"//mediapipe/framework/port:parse_text_proto",
],
)
cc_library(
name = "segmentation_smoothing_calculator",
srcs = ["segmentation_smoothing_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
":segmentation_smoothing_calculator_cc_proto",
"//mediapipe/framework:calculator_options_cc_proto",
"//mediapipe/framework/formats:image_format_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:image_frame_opencv",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:image_opencv",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:opencv_core",
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:vector",
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gl_quad_renderer",
"//mediapipe/gpu:shader_util",
],
}),
alwayslink = 1,
)
cc_test(
name = "segmentation_smoothing_calculator_test",
srcs = ["segmentation_smoothing_calculator_test.cc"],
deps = [
":image_clone_calculator",
":image_clone_calculator_cc_proto",
":segmentation_smoothing_calculator",
":segmentation_smoothing_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework/deps:file_path",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:image_opencv",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:opencv_imgcodecs",
"//mediapipe/framework/port:opencv_imgproc",
"//mediapipe/framework/port:parse_text_proto",
],
)
mediapipe/calculators/image/segmentation_smoothing_calculator.cc
+429
View File
@@ -0,0 +1,429 @@
// Copyright 2021 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 <algorithm>
#include <memory>
#include "mediapipe/calculators/image/segmentation_smoothing_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_options.pb.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_frame_opencv.h"
#include "mediapipe/framework/formats/image_opencv.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/opencv_core_inc.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/vector.h"
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#include "mediapipe/gpu/shader_util.h"
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
namespace {
constexpr char kCurrentMaskTag[] = "MASK";
constexpr char kPreviousMaskTag[] = "MASK_PREVIOUS";
constexpr char kOutputMaskTag[] = "MASK_SMOOTHED";
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
} // namespace
// A calculator for mixing two segmentation masks together,
// based on an uncertantity probability estimate.
//
// Inputs:
// MASK - Image containing the new/current mask.
// [ImageFormat::VEC32F1, or
// GpuBufferFormat::kBGRA32/kRGB24/kGrayHalf16/kGrayFloat32]
// MASK_PREVIOUS - Image containing previous mask.
// [Same format as MASK_CURRENT]
// * If input channels is >1, only the first channel (R) is used as the mask.
//
// Output:
// MASK_SMOOTHED - Blended mask.
// [Same format as MASK_CURRENT]
// * The resulting filtered mask will be stored in R channel,
// and duplicated in A if 4 channels.
//
// Options:
// combine_with_previous_ratio - Amount of previous to blend with current.
//
// Example:
// node {
// calculator: "SegmentationSmoothingCalculator"
// input_stream: "MASK:mask"
// input_stream: "MASK_PREVIOUS:mask_previous"
// output_stream: "MASK_SMOOTHED:mask_smoothed"
// options: {
// [mediapipe.SegmentationSmoothingCalculatorOptions.ext] {
// combine_with_previous_ratio: 0.9
// }
// }
// }
//
class SegmentationSmoothingCalculator : public CalculatorBase {
public:
SegmentationSmoothingCalculator() = default;
static absl::Status GetContract(CalculatorContract* cc);
// From Calculator.
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
absl::Status Close(CalculatorContext* cc) override;
private:
absl::Status RenderGpu(CalculatorContext* cc);
absl::Status RenderCpu(CalculatorContext* cc);
absl::Status GlSetup(CalculatorContext* cc);
void GlRender(CalculatorContext* cc);
float combine_with_previous_ratio_;
bool gpu_initialized_ = false;
#if !MEDIAPIPE_DISABLE_GPU
mediapipe::GlCalculatorHelper gpu_helper_;
GLuint program_ = 0;
#endif // !MEDIAPIPE_DISABLE_GPU
};
REGISTER_CALCULATOR(SegmentationSmoothingCalculator);
absl::Status SegmentationSmoothingCalculator::GetContract(
CalculatorContract* cc) {
CHECK_GE(cc->Inputs().NumEntries(), 1);
cc->Inputs().Tag(kCurrentMaskTag).Set<Image>();
cc->Inputs().Tag(kPreviousMaskTag).Set<Image>();
cc->Outputs().Tag(kOutputMaskTag).Set<Image>();
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
absl::Status SegmentationSmoothingCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
auto options =
cc->Options<mediapipe::SegmentationSmoothingCalculatorOptions>();
combine_with_previous_ratio_ = options.combine_with_previous_ratio();
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
absl::Status SegmentationSmoothingCalculator::Process(CalculatorContext* cc) {
if (cc->Inputs().Tag(kCurrentMaskTag).IsEmpty()) {
return absl::OkStatus();
}
if (cc->Inputs().Tag(kPreviousMaskTag).IsEmpty()) {
// Pass through current image if previous is not available.
cc->Outputs()
.Tag(kOutputMaskTag)
.AddPacket(cc->Inputs().Tag(kCurrentMaskTag).Value());
return absl::OkStatus();
}
// Run on GPU if incoming data is on GPU.
const bool use_gpu = cc->Inputs().Tag(kCurrentMaskTag).Get<Image>().UsesGpu();
if (use_gpu) {
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, cc]() -> absl::Status {
if (!gpu_initialized_) {
MP_RETURN_IF_ERROR(GlSetup(cc));
gpu_initialized_ = true;
}
MP_RETURN_IF_ERROR(RenderGpu(cc));
return absl::OkStatus();
}));
#else
return absl::InternalError("GPU processing is disabled.");
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
MP_RETURN_IF_ERROR(RenderCpu(cc));
}
return absl::OkStatus();
}
absl::Status SegmentationSmoothingCalculator::Close(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
gpu_helper_.RunInGlContext([this] {
if (program_) glDeleteProgram(program_);
program_ = 0;
});
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
absl::Status SegmentationSmoothingCalculator::RenderCpu(CalculatorContext* cc) {
// Setup source images.
const auto& current_frame = cc->Inputs().Tag(kCurrentMaskTag).Get<Image>();
const cv::Mat current_mat = mediapipe::formats::MatView(&current_frame);
RET_CHECK_EQ(current_mat.type(), CV_32FC1)
<< "Only 1-channel float input image is supported.";
const auto& previous_frame = cc->Inputs().Tag(kPreviousMaskTag).Get<Image>();
const cv::Mat previous_mat = mediapipe::formats::MatView(&previous_frame);
RET_CHECK_EQ(previous_mat.type(), current_mat.type())
<< "Warning: mixing input format types: " << previous_mat.type()
<< " != " << previous_mat.type();
RET_CHECK_EQ(current_mat.rows, previous_mat.rows);
RET_CHECK_EQ(current_mat.cols, previous_mat.cols);
// Setup destination image.
auto output_frame = std::make_shared<ImageFrame>(
current_frame.image_format(), current_mat.cols, current_mat.rows);
cv::Mat output_mat = mediapipe::formats::MatView(output_frame.get());
output_mat.setTo(cv::Scalar(0));
// Blending function.
const auto blending_fn = [&](const float prev_mask_value,
const float new_mask_value) {
/*
* Assume p := new_mask_value
* H(p) := 1 + (p * log(p) + (1-p) * log(1-p)) / log(2)
* uncertainty alpha(p) =
* Clamp(1 - (1 - H(p)) * (1 - H(p)), 0, 1) [squaring the uncertainty]
*
* The following polynomial approximates uncertainty alpha as a function
* of (p + 0.5):
*/
const float c1 = 5.68842;
const float c2 = -0.748699;
const float c3 = -57.8051;
const float c4 = 291.309;
const float c5 = -624.717;
const float t = new_mask_value - 0.5f;
const float x = t * t;
const float uncertainty =
1.0f -
std::min(1.0f, x * (c1 + x * (c2 + x * (c3 + x * (c4 + x * c5)))));
return new_mask_value + (prev_mask_value - new_mask_value) *
(uncertainty * combine_with_previous_ratio_);
};
// Write directly to the first channel of output.
for (int i = 0; i < output_mat.rows; ++i) {
float* out_ptr = output_mat.ptr<float>(i);
const float* curr_ptr = current_mat.ptr<float>(i);
const float* prev_ptr = previous_mat.ptr<float>(i);
for (int j = 0; j < output_mat.cols; ++j) {
const float new_mask_value = curr_ptr[j];
const float prev_mask_value = prev_ptr[j];
out_ptr[j] = blending_fn(prev_mask_value, new_mask_value);
}
}
cc->Outputs()
.Tag(kOutputMaskTag)
.AddPacket(MakePacket<Image>(output_frame).At(cc->InputTimestamp()));
return absl::OkStatus();
}
absl::Status SegmentationSmoothingCalculator::RenderGpu(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
// Setup source textures.
const auto& current_frame = cc->Inputs().Tag(kCurrentMaskTag).Get<Image>();
RET_CHECK(
(current_frame.format() == mediapipe::GpuBufferFormat::kBGRA32 ||
current_frame.format() == mediapipe::GpuBufferFormat::kGrayHalf16 ||
current_frame.format() == mediapipe::GpuBufferFormat::kGrayFloat32 ||
current_frame.format() == mediapipe::GpuBufferFormat::kRGB24))
<< "Only RGBA, RGB, or 1-channel Float input image supported.";
auto current_texture = gpu_helper_.CreateSourceTexture(current_frame);
const auto& previous_frame = cc->Inputs().Tag(kPreviousMaskTag).Get<Image>();
if (previous_frame.format() != current_frame.format()) {
LOG(ERROR) << "Warning: mixing input format types. ";
}
auto previous_texture = gpu_helper_.CreateSourceTexture(previous_frame);
// Setup destination texture.
const int width = current_frame.width(), height = current_frame.height();
auto output_texture = gpu_helper_.CreateDestinationTexture(
width, height, current_frame.format());
// Process shader.
{
gpu_helper_.BindFramebuffer(output_texture);
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, current_texture.name());
glActiveTexture(GL_TEXTURE2);
glBindTexture(GL_TEXTURE_2D, previous_texture.name());
GlRender(cc);
glActiveTexture(GL_TEXTURE2);
glBindTexture(GL_TEXTURE_2D, 0);
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, 0);
}
glFlush();
// Send out image as GPU packet.
auto output_frame = output_texture.GetFrame<Image>();
cc->Outputs()
.Tag(kOutputMaskTag)
.Add(output_frame.release(), cc->InputTimestamp());
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
void SegmentationSmoothingCalculator::GlRender(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
static const GLfloat square_vertices[] = {
-1.0f, -1.0f, // bottom left
1.0f, -1.0f, // bottom right
-1.0f, 1.0f, // top left
1.0f, 1.0f, // top right
};
static const GLfloat texture_vertices[] = {
0.0f, 0.0f, // bottom left
1.0f, 0.0f, // bottom right
0.0f, 1.0f, // top left
1.0f, 1.0f, // top right
};
// program
glUseProgram(program_);
// vertex storage
GLuint vbo[2];
glGenBuffers(2, vbo);
GLuint vao;
glGenVertexArrays(1, &vao);
glBindVertexArray(vao);
// vbo 0
glBindBuffer(GL_ARRAY_BUFFER, vbo[0]);
glBufferData(GL_ARRAY_BUFFER, 4 * 2 * sizeof(GLfloat), square_vertices,
GL_STATIC_DRAW);
glEnableVertexAttribArray(ATTRIB_VERTEX);
glVertexAttribPointer(ATTRIB_VERTEX, 2, GL_FLOAT, 0, 0, nullptr);
// vbo 1
glBindBuffer(GL_ARRAY_BUFFER, vbo[1]);
glBufferData(GL_ARRAY_BUFFER, 4 * 2 * sizeof(GLfloat), texture_vertices,
GL_STATIC_DRAW);
glEnableVertexAttribArray(ATTRIB_TEXTURE_POSITION);
glVertexAttribPointer(ATTRIB_TEXTURE_POSITION, 2, GL_FLOAT, 0, 0, nullptr);
// draw
glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
// cleanup
glDisableVertexAttribArray(ATTRIB_VERTEX);
glDisableVertexAttribArray(ATTRIB_TEXTURE_POSITION);
glBindBuffer(GL_ARRAY_BUFFER, 0);
glBindVertexArray(0);
glDeleteVertexArrays(1, &vao);
glDeleteBuffers(2, vbo);
#endif // !MEDIAPIPE_DISABLE_GPU
}
absl::Status SegmentationSmoothingCalculator::GlSetup(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
};
const GLchar* attr_name[NUM_ATTRIBUTES] = {
"position",
"texture_coordinate",
};
// Shader to blend in previous mask based on computed uncertainty probability.
const std::string frag_src =
absl::StrCat(std::string(mediapipe::kMediaPipeFragmentShaderPreamble),
R"(
DEFAULT_PRECISION(mediump, float)
#ifdef GL_ES
#define fragColor gl_FragColor
#else
out vec4 fragColor;
#endif // defined(GL_ES);
in vec2 sample_coordinate;
uniform sampler2D current_mask;
uniform sampler2D previous_mask;
uniform float combine_with_previous_ratio;
void main() {
vec4 current_pix = texture2D(current_mask, sample_coordinate);
vec4 previous_pix = texture2D(previous_mask, sample_coordinate);
float new_mask_value = current_pix.r;
float prev_mask_value = previous_pix.r;
// Assume p := new_mask_value
// H(p) := 1 + (p * log(p) + (1-p) * log(1-p)) / log(2)
// uncertainty alpha(p) =
// Clamp(1 - (1 - H(p)) * (1 - H(p)), 0, 1) [squaring the uncertainty]
//
// The following polynomial approximates uncertainty alpha as a function
// of (p + 0.5):
const float c1 = 5.68842;
const float c2 = -0.748699;
const float c3 = -57.8051;
const float c4 = 291.309;
const float c5 = -624.717;
float t = new_mask_value - 0.5;
float x = t * t;
float uncertainty =
1.0 - min(1.0, x * (c1 + x * (c2 + x * (c3 + x * (c4 + x * c5)))));
new_mask_value +=
(prev_mask_value - new_mask_value) * (uncertainty * combine_with_previous_ratio);
fragColor = vec4(new_mask_value, 0.0, 0.0, new_mask_value);
}
)");
// Create shader program and set parameters.
mediapipe::GlhCreateProgram(mediapipe::kBasicVertexShader, frag_src.c_str(),
NUM_ATTRIBUTES, (const GLchar**)&attr_name[0],
attr_location, &program_);
RET_CHECK(program_) << "Problem initializing the program.";
glUseProgram(program_);
glUniform1i(glGetUniformLocation(program_, "current_mask"), 1);
glUniform1i(glGetUniformLocation(program_, "previous_mask"), 2);
glUniform1f(glGetUniformLocation(program_, "combine_with_previous_ratio"),
combine_with_previous_ratio_);
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
} // namespace mediapipe
mediapipe/calculators/image/segmentation_smoothing_calculator.proto
+35
View File
@@ -0,0 +1,35 @@
// Copyright 2021 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";
message SegmentationSmoothingCalculatorOptions {
extend CalculatorOptions {
optional SegmentationSmoothingCalculatorOptions ext = 377425128;
}
// How much to blend in previous mask, based on a probability estimate.
// Range: [0-1]
// 0 = Use only current frame (no blending).
// 1 = Blend in the previous mask based on uncertainty estimate.
// With ratio at 1, the uncertainty estimate is trusted completely.
// When uncertainty is high, the previous mask is given higher weight.
// Therefore, if both ratio and uncertainty are 1, only old mask is used.
// A pixel is 'uncertain' if its value is close to the middle (0.5 or 127).
optional float combine_with_previous_ratio = 1 [default = 0.0];
}
mediapipe/calculators/image/segmentation_smoothing_calculator_test.cc
+206
View File
@@ -0,0 +1,206 @@
// Copyright 2018 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/calculators/image/segmentation_smoothing_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_runner.h"
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_opencv.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/opencv_imgcodecs_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"
namespace mediapipe {
namespace {
// 4x4 VEC32F1, center 2x2 block set at ~250
const float mask_data[] = {
0.00, 0.00, 0.00, 0.00, //
0.00, 0.98, 0.98, 0.00, //
0.00, 0.98, 0.98, 0.00, //
0.00, 0.00, 0.00, 0.00, //
};
void RunGraph(Packet curr_packet, Packet prev_packet, bool use_gpu, float ratio,
cv::Mat* result) {
CalculatorGraphConfig graph_config;
if (use_gpu) {
graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(absl::Substitute(
R"pb(
input_stream: "curr_mask"
input_stream: "prev_mask"
output_stream: "new_mask"
node {
calculator: "ImageCloneCalculator"
input_stream: "curr_mask"
output_stream: "curr_mask_gpu"
options: {
[mediapipe.ImageCloneCalculatorOptions.ext] {
output_on_gpu: true
}
}
}
node {
calculator: "ImageCloneCalculator"
input_stream: "prev_mask"
output_stream: "prev_mask_gpu"
options: {
[mediapipe.ImageCloneCalculatorOptions.ext] {
output_on_gpu: true
}
}
}
node {
calculator: "SegmentationSmoothingCalculator"
input_stream: "MASK:curr_mask_gpu"
input_stream: "MASK_PREVIOUS:prev_mask_gpu"
output_stream: "MASK_SMOOTHED:new_mask"
node_options {
[type.googleapis.com/
mediapipe.SegmentationSmoothingCalculatorOptions]: {
combine_with_previous_ratio: $0
}
}
}
)pb",
ratio));
} else {
graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(absl::Substitute(
R"pb(
input_stream: "curr_mask"
input_stream: "prev_mask"
output_stream: "new_mask"
node {
calculator: "SegmentationSmoothingCalculator"
input_stream: "MASK:curr_mask"
input_stream: "MASK_PREVIOUS:prev_mask"
output_stream: "MASK_SMOOTHED:new_mask"
node_options {
[type.googleapis.com/
mediapipe.SegmentationSmoothingCalculatorOptions]: {
combine_with_previous_ratio: $0
}
}
}
)pb",
ratio));
}
std::vector<Packet> output_packets;
tool::AddVectorSink("new_mask", &graph_config, &output_packets);
CalculatorGraph graph(graph_config);
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(
graph.AddPacketToInputStream("curr_mask", curr_packet.At(Timestamp(0))));
MP_ASSERT_OK(
graph.AddPacketToInputStream("prev_mask", prev_packet.At(Timestamp(0))));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(1, output_packets.size());
Image result_image = output_packets[0].Get<Image>();
cv::Mat result_mat = formats::MatView(&result_image);
result_mat.copyTo(*result);
// Fully close graph at end, otherwise calculator+Images are destroyed
// after calling WaitUntilDone().
MP_ASSERT_OK(graph.CloseInputStream("curr_mask"));
MP_ASSERT_OK(graph.CloseInputStream("prev_mask"));
MP_ASSERT_OK(graph.WaitUntilDone());
}
void RunTest(bool use_gpu, float mix_ratio, cv::Mat& test_result) {
cv::Mat mask_mat(cv::Size(4, 4), CV_32FC1, const_cast<float*>(mask_data));
cv::Mat curr_mat = mask_mat;
// 3x3 blur of 250 block produces all pixels '111'.
cv::Mat prev_mat;
cv::blur(mask_mat, prev_mat, cv::Size(3, 3));
Packet curr_packet = MakePacket<Image>(std::make_unique<ImageFrame>(
ImageFormat::VEC32F1, curr_mat.size().width, curr_mat.size().height));
curr_mat.copyTo(formats::MatView(&(curr_packet.Get<Image>())));
Packet prev_packet = MakePacket<Image>(std::make_unique<ImageFrame>(
ImageFormat::VEC32F1, prev_mat.size().width, prev_mat.size().height));
prev_mat.copyTo(formats::MatView(&(prev_packet.Get<Image>())));
cv::Mat result;
RunGraph(curr_packet, prev_packet, use_gpu, mix_ratio, &result);
ASSERT_EQ(curr_mat.rows, result.rows);
ASSERT_EQ(curr_mat.cols, result.cols);
ASSERT_EQ(curr_mat.type(), result.type());
result.copyTo(test_result);
if (mix_ratio == 1.0) {
for (int i = 0; i < 4; ++i) {
for (int j = 0; j < 4; ++j) {
float in = curr_mat.at<float>(i, j);
float out = result.at<float>(i, j);
// Since the input has high value (250), it has low uncertainty.
// So the output should have changed lower (towards prev),
// but not too much.
if (in > 0) EXPECT_NE(in, out);
EXPECT_NEAR(in, out, 3.0 / 255.0);
}
}
} else if (mix_ratio == 0.0) {
for (int i = 0; i < 4; ++i) {
for (int j = 0; j < 4; ++j) {
float in = curr_mat.at<float>(i, j);
float out = result.at<float>(i, j);
EXPECT_EQ(in, out); // Output should match current.
}
}
} else {
LOG(ERROR) << "invalid ratio";
}
}
TEST(SegmentationSmoothingCalculatorTest, TestSmoothing) {
bool use_gpu;
float mix_ratio;
use_gpu = false;
mix_ratio = 0.0;
cv::Mat cpu_0;
RunTest(use_gpu, mix_ratio, cpu_0);
use_gpu = false;
mix_ratio = 1.0;
cv::Mat cpu_1;
RunTest(use_gpu, mix_ratio, cpu_1);
use_gpu = true;
mix_ratio = 1.0;
cv::Mat gpu_1;
RunTest(use_gpu, mix_ratio, gpu_1);
// CPU & GPU should match.
for (int i = 0; i < 4; ++i) {
for (int j = 0; j < 4; ++j) {
float gpu = gpu_1.at<float>(i, j);
float cpu = cpu_1.at<float>(i, j);
EXPECT_EQ(cpu, gpu);
}
}
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/tensor/BUILD
+5 -1
View File
@@ -109,6 +109,8 @@ cc_library(
"//mediapipe/gpu:MPPMetalUtil",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/objc:mediapipe_framework_ios",
"//mediapipe/util/tflite:config",
"@com_google_absl//absl/memory",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate_internal",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
@@ -478,7 +480,6 @@ cc_library(
deps = [
":image_to_tensor_calculator_cc_proto",
":image_to_tensor_converter",
":image_to_tensor_converter_opencv",
":image_to_tensor_utils",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/formats:image",
@@ -494,6 +495,9 @@ cc_library(
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [":image_to_tensor_calculator_gpu_deps"],
}) + select({
"//mediapipe/framework/port:disable_opencv": [],
"//conditions:default": [":image_to_tensor_converter_opencv"],
}),
alwayslink = 1,
)
mediapipe/calculators/tensor/image_to_tensor_calculator.cc
+9 -1
View File
@@ -18,7 +18,6 @@
#include "mediapipe/calculators/tensor/image_to_tensor_calculator.pb.h"
#include "mediapipe/calculators/tensor/image_to_tensor_converter.h"
#include "mediapipe/calculators/tensor/image_to_tensor_converter_opencv.h"
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
@@ -33,6 +32,10 @@
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/gpu/gpu_origin.pb.h"
#if !MEDIAPIPE_DISABLE_OPENCV
#include "mediapipe/calculators/tensor/image_to_tensor_converter_opencv.h"
#endif
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/gpu/gpu_buffer.h"
@@ -301,8 +304,13 @@ class ImageToTensorCalculator : public Node {
}
} else {
if (!cpu_converter_) {
#if !MEDIAPIPE_DISABLE_OPENCV
ASSIGN_OR_RETURN(cpu_converter_,
CreateOpenCvConverter(cc, GetBorderMode()));
#else
LOG(FATAL) << "Cannot create image to tensor opencv converter since "
"MEDIAPIPE_DISABLE_OPENCV is defined.";
#endif // !MEDIAPIPE_DISABLE_OPENCV
}
}
return absl::OkStatus();
mediapipe/calculators/tensor/image_to_tensor_converter_gl_buffer.cc
+2 -3
View File
@@ -312,7 +312,7 @@ class GlProcessor : public ImageToTensorConverter {
return absl::OkStatus();
}));
return std::move(tensor);
return tensor;
}
~GlProcessor() override {
@@ -338,8 +338,7 @@ CreateImageToGlBufferTensorConverter(CalculatorContext* cc,
auto result = absl::make_unique<GlProcessor>();
MP_RETURN_IF_ERROR(result->Init(cc, input_starts_at_bottom, border_mode));
// Simply "return std::move(result)" failed to build on macOS with bazel.
return std::unique_ptr<ImageToTensorConverter>(std::move(result));
return result;
}
} // namespace mediapipe
mediapipe/calculators/tensor/image_to_tensor_converter_gl_texture.cc
+1 -3
View File
@@ -334,9 +334,7 @@ CreateImageToGlTextureTensorConverter(CalculatorContext* cc,
BorderMode border_mode) {
auto result = absl::make_unique<GlProcessor>();
MP_RETURN_IF_ERROR(result->Init(cc, input_starts_at_bottom, border_mode));
// Simply "return std::move(result)" failed to build on macOS with bazel.
return std::unique_ptr<ImageToTensorConverter>(std::move(result));
return result;
}
} // namespace mediapipe
mediapipe/calculators/tensor/image_to_tensor_converter_metal.cc
+2 -3
View File
@@ -383,7 +383,7 @@ class MetalProcessor : public ImageToTensorConverter {
tflite::gpu::HW(output_dims.height, output_dims.width),
command_buffer, buffer_view.buffer()));
[command_buffer commit];
return std::move(tensor);
return tensor;
}
}
@@ -399,8 +399,7 @@ absl::StatusOr<std::unique_ptr<ImageToTensorConverter>> CreateMetalConverter(
auto result = absl::make_unique<MetalProcessor>();
MP_RETURN_IF_ERROR(result->Init(cc, border_mode));
// Simply "return std::move(result)" failed to build on macOS with bazel.
return std::unique_ptr<ImageToTensorConverter>(std::move(result));
return result;
}
} // namespace mediapipe
mediapipe/calculators/tensor/image_to_tensor_converter_opencv.cc
+2 -5
View File
@@ -103,7 +103,7 @@ class OpenCvProcessor : public ImageToTensorConverter {
GetValueRangeTransformation(kInputImageRangeMin, kInputImageRangeMax,
range_min, range_max));
transformed.convertTo(dst, CV_32FC3, transform.scale, transform.offset);
return std::move(tensor);
return tensor;
}
private:
@@ -114,10 +114,7 @@ class OpenCvProcessor : public ImageToTensorConverter {
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>> CreateOpenCvConverter(
CalculatorContext* cc, BorderMode border_mode) {
// Simply "return absl::make_unique<OpenCvProcessor>()" failed to build on
// macOS with bazel.
return std::unique_ptr<ImageToTensorConverter>(
absl::make_unique<OpenCvProcessor>(border_mode));
return absl::make_unique<OpenCvProcessor>(border_mode);
}
} // namespace mediapipe
mediapipe/calculators/tensor/testdata/face_detection_test.pbtxt
Vendored
+1 -1
View File
@@ -4,7 +4,7 @@ output_stream: "detections"
# Subgraph that detects faces.
node {
calculator: "FaceDetectionFrontCpu"
calculator: "FaceDetectionShortRangeCpu"
input_stream: "IMAGE:image"
output_stream: "DETECTIONS:detections"
}
mediapipe/calculators/tensorflow/tensorflow_inference_calculator.cc
+1 -1
View File
@@ -490,7 +490,7 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
<< keyed_tensors.first;
}
} else {
// Pad by replicating the first tens or, then ignore the values.
// Pad by replicating the first tensor, then ignore the values.
keyed_tensors.second.resize(options_.batch_size());
std::fill(keyed_tensors.second.begin() +
inference_state->batch_timestamps_.size(),
mediapipe/calculators/util/BUILD
+39
View File
@@ -840,6 +840,20 @@ cc_test(
],
)
cc_library(
name = "world_landmark_projection_calculator",
srcs = ["world_landmark_projection_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
],
alwayslink = 1,
)
mediapipe_proto_library(
name = "landmarks_smoothing_calculator_proto",
srcs = ["landmarks_smoothing_calculator.proto"],
@@ -894,6 +908,31 @@ cc_library(
alwayslink = 1,
)
mediapipe_proto_library(
name = "visibility_copy_calculator_proto",
srcs = ["visibility_copy_calculator.proto"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
],
)
cc_library(
name = "visibility_copy_calculator",
srcs = ["visibility_copy_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
":visibility_copy_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:timestamp",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/port:ret_check",
"@com_google_absl//absl/algorithm:container",
],
alwayslink = 1,
)
cc_library(
name = "landmarks_to_floats_calculator",
srcs = ["landmarks_to_floats_calculator.cc"],
mediapipe/calculators/util/annotation_overlay_calculator.cc
+9
View File
@@ -272,6 +272,15 @@ absl::Status AnnotationOverlayCalculator::Open(CalculatorContext* cc) {
}
absl::Status AnnotationOverlayCalculator::Process(CalculatorContext* cc) {
if (cc->Inputs().HasTag(kGpuBufferTag) &&
cc->Inputs().Tag(kGpuBufferTag).IsEmpty()) {
return absl::OkStatus();
}
if (cc->Inputs().HasTag(kImageFrameTag) &&
cc->Inputs().Tag(kImageFrameTag).IsEmpty()) {
return absl::OkStatus();
}
// Initialize render target, drawn with OpenCV.
std::unique_ptr<cv::Mat> image_mat;
ImageFormat::Format target_format;
mediapipe/calculators/util/detections_to_rects_calculator.cc
+3
View File
@@ -203,6 +203,9 @@ absl::Status DetectionsToRectsCalculator::Process(CalculatorContext* cc) {
cc->Inputs().Tag(kDetectionsTag).IsEmpty()) {
return absl::OkStatus();
}
if (rotate_ && !HasTagValue(cc, kImageSizeTag)) {
return absl::OkStatus();
}
std::vector<Detection> detections;
if (cc->Inputs().HasTag(kDetectionTag)) {
mediapipe/calculators/util/rect_transformation_calculator.cc
+4 -4
View File
@@ -130,8 +130,8 @@ absl::Status RectTransformationCalculator::Process(CalculatorContext* cc) {
}
cc->Outputs().Index(0).Add(output_rects.release(), cc->InputTimestamp());
}
if (cc->Inputs().HasTag(kNormRectTag) &&
!cc->Inputs().Tag(kNormRectTag).IsEmpty()) {
if (HasTagValue(cc->Inputs(), kNormRectTag) &&
HasTagValue(cc->Inputs(), kImageSizeTag)) {
auto rect = cc->Inputs().Tag(kNormRectTag).Get<NormalizedRect>();
const auto& image_size =
cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
@@ -139,8 +139,8 @@ absl::Status RectTransformationCalculator::Process(CalculatorContext* cc) {
cc->Outputs().Index(0).AddPacket(
MakePacket<NormalizedRect>(rect).At(cc->InputTimestamp()));
}
if (cc->Inputs().HasTag(kNormRectsTag) &&
!cc->Inputs().Tag(kNormRectsTag).IsEmpty()) {
if (HasTagValue(cc->Inputs(), kNormRectsTag) &&
HasTagValue(cc->Inputs(), kImageSizeTag)) {
auto rects =
cc->Inputs().Tag(kNormRectsTag).Get<std::vector<NormalizedRect>>();
const auto& image_size =
mediapipe/calculators/video/motion_analysis_calculator.cc
+1 -1
View File
@@ -549,7 +549,7 @@ absl::Status MotionAnalysisCalculator::Process(CalculatorContext* cc) {
timestamp_buffer_.push_back(timestamp);
++frame_idx_;
VLOG_EVERY_N(0, 100) << "Analyzed frame " << frame_idx_;
VLOG_EVERY_N(1, 100) << "Analyzed frame " << frame_idx_;
// Buffer input frames only if visualization is requested.
if (visualize_output_ || video_output_) {