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];
|
||||
}
|
||||
|
||||
@@ -753,3 +753,76 @@ cc_test(
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
],
|
||||
)
|
||||
|
||||
# Copied from /mediapipe/calculators/tflite/BUILD
|
||||
selects.config_setting_group(
|
||||
name = "gpu_inference_disabled",
|
||||
match_any = [
|
||||
"//mediapipe/gpu:disable_gpu",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_proto_library(
|
||||
name = "tensors_to_segmentation_calculator_proto",
|
||||
srcs = ["tensors_to_segmentation_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_options_proto",
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/gpu:gpu_origin_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tensors_to_segmentation_calculator",
|
||||
srcs = ["tensors_to_segmentation_calculator.cc"],
|
||||
copts = select({
|
||||
"//mediapipe:apple": [
|
||||
"-x objective-c++",
|
||||
"-fobjc-arc", # enable reference-counting
|
||||
],
|
||||
"//conditions:default": [],
|
||||
}),
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":tensors_to_segmentation_calculator_cc_proto",
|
||||
"@com_google_absl//absl/strings:str_format",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/types:span",
|
||||
"//mediapipe/framework/formats:image",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_opencv",
|
||||
"//mediapipe/framework/formats:tensor",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:port",
|
||||
"//mediapipe/util:resource_util",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
"//mediapipe/gpu:gpu_origin_cc_proto",
|
||||
"//mediapipe/framework/port:statusor",
|
||||
] + selects.with_or({
|
||||
"//mediapipe/gpu:disable_gpu": [],
|
||||
"//conditions:default": [
|
||||
"//mediapipe/gpu:gl_calculator_helper",
|
||||
"//mediapipe/gpu:gl_simple_shaders",
|
||||
"//mediapipe/gpu:gpu_buffer",
|
||||
"//mediapipe/gpu:shader_util",
|
||||
],
|
||||
}) + selects.with_or({
|
||||
":gpu_inference_disabled": [],
|
||||
"//mediapipe:ios": [
|
||||
"//mediapipe/gpu:MPPMetalUtil",
|
||||
"//mediapipe/gpu:MPPMetalHelper",
|
||||
],
|
||||
"//conditions:default": [
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_texture",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl/converters:util",
|
||||
],
|
||||
}),
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
@@ -105,6 +105,15 @@ void ConvertAnchorsToRawValues(const std::vector<Anchor>& anchors,
|
||||
// for anchors (e.g. for SSD models) depend on the outputs of the
|
||||
// detection model. The size of anchor tensor must be (num_boxes *
|
||||
// 4).
|
||||
//
|
||||
// Input side packet:
|
||||
// ANCHORS (optional) - The anchors used for decoding the bounding boxes, as a
|
||||
// vector of `Anchor` protos. Not required if post-processing is built-in
|
||||
// the model.
|
||||
// IGNORE_CLASSES (optional) - The list of class ids that should be ignored, as
|
||||
// a vector of integers. It overrides the corresponding field in the
|
||||
// calculator options.
|
||||
//
|
||||
// Output:
|
||||
// DETECTIONS - Result MediaPipe detections.
|
||||
//
|
||||
@@ -132,8 +141,11 @@ class TensorsToDetectionsCalculator : public Node {
|
||||
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
|
||||
static constexpr SideInput<std::vector<Anchor>>::Optional kInAnchors{
|
||||
"ANCHORS"};
|
||||
static constexpr SideInput<std::vector<int>>::Optional kSideInIgnoreClasses{
|
||||
"IGNORE_CLASSES"};
|
||||
static constexpr Output<std::vector<Detection>> kOutDetections{"DETECTIONS"};
|
||||
MEDIAPIPE_NODE_CONTRACT(kInTensors, kInAnchors, kOutDetections);
|
||||
MEDIAPIPE_NODE_CONTRACT(kInTensors, kInAnchors, kSideInIgnoreClasses,
|
||||
kOutDetections);
|
||||
static absl::Status UpdateContract(CalculatorContract* cc);
|
||||
|
||||
absl::Status Open(CalculatorContext* cc) override;
|
||||
@@ -566,8 +578,15 @@ absl::Status TensorsToDetectionsCalculator::LoadOptions(CalculatorContext* cc) {
|
||||
kNumCoordsPerBox,
|
||||
num_coords_);
|
||||
|
||||
for (int i = 0; i < options_.ignore_classes_size(); ++i) {
|
||||
ignore_classes_.insert(options_.ignore_classes(i));
|
||||
if (kSideInIgnoreClasses(cc).IsConnected()) {
|
||||
RET_CHECK(!kSideInIgnoreClasses(cc).IsEmpty());
|
||||
for (int ignore_class : *kSideInIgnoreClasses(cc)) {
|
||||
ignore_classes_.insert(ignore_class);
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < options_.ignore_classes_size(); ++i) {
|
||||
ignore_classes_.insert(options_.ignore_classes(i));
|
||||
}
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
|
||||
@@ -56,7 +56,7 @@ message TensorsToDetectionsCalculatorOptions {
|
||||
// [x_center, y_center, w, h].
|
||||
optional bool reverse_output_order = 14 [default = false];
|
||||
// The ids of classes that should be ignored during decoding the score for
|
||||
// each predicted box.
|
||||
// each predicted box. Can be overridden with IGNORE_CLASSES side packet.
|
||||
repeated int32 ignore_classes = 8;
|
||||
|
||||
optional bool sigmoid_score = 15 [default = false];
|
||||
|
||||
@@ -0,0 +1,885 @@
|
||||
// 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 <vector>
|
||||
|
||||
#include "absl/strings/str_format.h"
|
||||
#include "absl/types/span.h"
|
||||
#include "mediapipe/calculators/tensor/tensors_to_segmentation_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
#include "mediapipe/framework/formats/image_opencv.h"
|
||||
#include "mediapipe/framework/formats/tensor.h"
|
||||
#include "mediapipe/framework/port.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/statusor.h"
|
||||
#include "mediapipe/gpu/gpu_origin.pb.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "mediapipe/gpu/gl_simple_shaders.h"
|
||||
#include "mediapipe/gpu/gpu_buffer.h"
|
||||
#include "mediapipe/gpu/shader_util.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
#include "tensorflow/lite/delegates/gpu/gl/converters/util.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_texture.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
||||
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
|
||||
#if MEDIAPIPE_METAL_ENABLED
|
||||
#import <CoreVideo/CoreVideo.h>
|
||||
#import <Metal/Metal.h>
|
||||
#import <MetalKit/MetalKit.h>
|
||||
|
||||
#import "mediapipe/gpu/MPPMetalHelper.h"
|
||||
#include "mediapipe/gpu/MPPMetalUtil.h"
|
||||
#endif // MEDIAPIPE_METAL_ENABLED
|
||||
|
||||
namespace {
|
||||
constexpr int kWorkgroupSize = 8; // Block size for GPU shader.
|
||||
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
||||
|
||||
// Commonly used to compute the number of blocks to launch in a kernel.
|
||||
int NumGroups(const int size, const int group_size) { // NOLINT
|
||||
return (size + group_size - 1) / group_size;
|
||||
}
|
||||
|
||||
bool CanUseGpu() {
|
||||
#if !MEDIAPIPE_DISABLE_GPU || MEDIAPIPE_METAL_ENABLED
|
||||
// TODO: Configure GPU usage policy in individual calculators.
|
||||
constexpr bool kAllowGpuProcessing = true;
|
||||
return kAllowGpuProcessing;
|
||||
#else
|
||||
return false;
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU || MEDIAPIPE_METAL_ENABLED
|
||||
}
|
||||
|
||||
constexpr char kTensorsTag[] = "TENSORS";
|
||||
constexpr char kOutputSizeTag[] = "OUTPUT_SIZE";
|
||||
constexpr char kMaskTag[] = "MASK";
|
||||
|
||||
absl::StatusOr<std::tuple<int, int, int>> GetHwcFromDims(
|
||||
const std::vector<int>& dims) {
|
||||
if (dims.size() == 3) {
|
||||
return std::make_tuple(dims[0], dims[1], dims[2]);
|
||||
} else if (dims.size() == 4) {
|
||||
// BHWC format check B == 1
|
||||
RET_CHECK_EQ(1, dims[0]) << "Expected batch to be 1 for BHWC heatmap";
|
||||
return std::make_tuple(dims[1], dims[2], dims[3]);
|
||||
} else {
|
||||
RET_CHECK(false) << "Invalid shape for segmentation tensor " << dims.size();
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
using ::tflite::gpu::gl::GlProgram;
|
||||
using ::tflite::gpu::gl::GlShader;
|
||||
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
|
||||
// Converts Tensors from a tflite segmentation model to an image mask.
|
||||
//
|
||||
// Performs optional upscale to OUTPUT_SIZE dimensions if provided,
|
||||
// otherwise the mask is the same size as input tensor.
|
||||
//
|
||||
// If at least one input tensor is already on GPU, processing happens on GPU and
|
||||
// the output mask is also stored on GPU. Otherwise, processing and the output
|
||||
// mask are both on CPU.
|
||||
//
|
||||
// On GPU, the mask is an RGBA image, in both the R & A channels, scaled 0-1.
|
||||
// On CPU, the mask is a ImageFormat::VEC32F1 image, with values scaled 0-1.
|
||||
//
|
||||
//
|
||||
// Inputs:
|
||||
// One of the following TENSORS tags:
|
||||
// TENSORS: Vector of Tensor,
|
||||
// The tensor dimensions are specified in this calculator's options.
|
||||
// OUTPUT_SIZE(optional): std::pair<int, int>,
|
||||
// If provided, the size to upscale mask to.
|
||||
//
|
||||
// Output:
|
||||
// MASK: An Image output mask, RGBA(GPU) / VEC32F1(CPU).
|
||||
//
|
||||
// Options:
|
||||
// See tensors_to_segmentation_calculator.proto
|
||||
//
|
||||
// Usage example:
|
||||
// node {
|
||||
// calculator: "TensorsToSegmentationCalculator"
|
||||
// input_stream: "TENSORS:tensors"
|
||||
// input_stream: "OUTPUT_SIZE:size"
|
||||
// output_stream: "MASK:hair_mask"
|
||||
// node_options: {
|
||||
// [mediapipe.TensorsToSegmentationCalculatorOptions] {
|
||||
// output_layer_index: 1
|
||||
// # gpu_origin: CONVENTIONAL # or TOP_LEFT
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// Currently only OpenGLES 3.1 and CPU backends supported.
|
||||
// TODO Refactor and add support for other backends/platforms.
|
||||
//
|
||||
class TensorsToSegmentationCalculator : public CalculatorBase {
|
||||
public:
|
||||
static absl::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
absl::Status Open(CalculatorContext* cc) override;
|
||||
absl::Status Process(CalculatorContext* cc) override;
|
||||
absl::Status Close(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
absl::Status LoadOptions(CalculatorContext* cc);
|
||||
absl::Status InitGpu(CalculatorContext* cc);
|
||||
absl::Status ProcessGpu(CalculatorContext* cc);
|
||||
absl::Status ProcessCpu(CalculatorContext* cc);
|
||||
void GlRender();
|
||||
|
||||
bool DoesGpuTextureStartAtBottom() {
|
||||
return options_.gpu_origin() != mediapipe::GpuOrigin_Mode_TOP_LEFT;
|
||||
}
|
||||
|
||||
template <class T>
|
||||
absl::Status ApplyActivation(cv::Mat& tensor_mat, cv::Mat* small_mask_mat);
|
||||
|
||||
::mediapipe::TensorsToSegmentationCalculatorOptions options_;
|
||||
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||
GLuint upsample_program_;
|
||||
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
std::unique_ptr<GlProgram> mask_program_31_;
|
||||
#else
|
||||
GLuint mask_program_20_;
|
||||
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
#if MEDIAPIPE_METAL_ENABLED
|
||||
MPPMetalHelper* metal_helper_ = nullptr;
|
||||
id<MTLComputePipelineState> mask_program_;
|
||||
#endif // MEDIAPIPE_METAL_ENABLED
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
};
|
||||
REGISTER_CALCULATOR(TensorsToSegmentationCalculator);
|
||||
|
||||
// static
|
||||
absl::Status TensorsToSegmentationCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
RET_CHECK(!cc->Inputs().GetTags().empty());
|
||||
RET_CHECK(!cc->Outputs().GetTags().empty());
|
||||
|
||||
// Inputs.
|
||||
cc->Inputs().Tag(kTensorsTag).Set<std::vector<Tensor>>();
|
||||
if (cc->Inputs().HasTag(kOutputSizeTag)) {
|
||||
cc->Inputs().Tag(kOutputSizeTag).Set<std::pair<int, int>>();
|
||||
}
|
||||
|
||||
// Outputs.
|
||||
cc->Outputs().Tag(kMaskTag).Set<Image>();
|
||||
|
||||
if (CanUseGpu()) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||
#if MEDIAPIPE_METAL_ENABLED
|
||||
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
|
||||
#endif // MEDIAPIPE_METAL_ENABLED
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status TensorsToSegmentationCalculator::Open(CalculatorContext* cc) {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
bool use_gpu = false;
|
||||
|
||||
if (CanUseGpu()) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
use_gpu = true;
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
#if MEDIAPIPE_METAL_ENABLED
|
||||
metal_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
|
||||
RET_CHECK(metal_helper_);
|
||||
#endif // MEDIAPIPE_METAL_ENABLED
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
|
||||
MP_RETURN_IF_ERROR(LoadOptions(cc));
|
||||
|
||||
if (use_gpu) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
MP_RETURN_IF_ERROR(InitGpu(cc));
|
||||
#else
|
||||
RET_CHECK_FAIL() << "GPU processing disabled.";
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status TensorsToSegmentationCalculator::Process(CalculatorContext* cc) {
|
||||
if (cc->Inputs().Tag(kTensorsTag).IsEmpty()) {
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag(kTensorsTag).Get<std::vector<Tensor>>();
|
||||
|
||||
bool use_gpu = false;
|
||||
if (CanUseGpu()) {
|
||||
// Use GPU processing only if at least one input tensor is already on GPU.
|
||||
for (const auto& tensor : input_tensors) {
|
||||
if (tensor.ready_on_gpu()) {
|
||||
use_gpu = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Validate tensor channels and activation type.
|
||||
{
|
||||
RET_CHECK(!input_tensors.empty());
|
||||
ASSIGN_OR_RETURN(auto hwc, GetHwcFromDims(input_tensors[0].shape().dims));
|
||||
int tensor_channels = std::get<2>(hwc);
|
||||
typedef mediapipe::TensorsToSegmentationCalculatorOptions Options;
|
||||
switch (options_.activation()) {
|
||||
case Options::NONE:
|
||||
RET_CHECK_EQ(tensor_channels, 1);
|
||||
break;
|
||||
case Options::SIGMOID:
|
||||
RET_CHECK_EQ(tensor_channels, 1);
|
||||
break;
|
||||
case Options::SOFTMAX:
|
||||
RET_CHECK_EQ(tensor_channels, 2);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (use_gpu) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, cc]() -> absl::Status {
|
||||
MP_RETURN_IF_ERROR(ProcessGpu(cc));
|
||||
return absl::OkStatus();
|
||||
}));
|
||||
#else
|
||||
RET_CHECK_FAIL() << "GPU processing disabled.";
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
} else {
|
||||
MP_RETURN_IF_ERROR(ProcessCpu(cc));
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status TensorsToSegmentationCalculator::Close(CalculatorContext* cc) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
gpu_helper_.RunInGlContext([this] {
|
||||
if (upsample_program_) glDeleteProgram(upsample_program_);
|
||||
upsample_program_ = 0;
|
||||
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
mask_program_31_.reset();
|
||||
#else
|
||||
if (mask_program_20_) glDeleteProgram(mask_program_20_);
|
||||
mask_program_20_ = 0;
|
||||
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
#if MEDIAPIPE_METAL_ENABLED
|
||||
mask_program_ = nil;
|
||||
#endif // MEDIAPIPE_METAL_ENABLED
|
||||
});
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status TensorsToSegmentationCalculator::ProcessCpu(
|
||||
CalculatorContext* cc) {
|
||||
// Get input streams, and dimensions.
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag(kTensorsTag).Get<std::vector<Tensor>>();
|
||||
ASSIGN_OR_RETURN(auto hwc, GetHwcFromDims(input_tensors[0].shape().dims));
|
||||
auto [tensor_height, tensor_width, tensor_channels] = hwc;
|
||||
int output_width = tensor_width, output_height = tensor_height;
|
||||
if (cc->Inputs().HasTag(kOutputSizeTag)) {
|
||||
const auto& size =
|
||||
cc->Inputs().Tag(kOutputSizeTag).Get<std::pair<int, int>>();
|
||||
output_width = size.first;
|
||||
output_height = size.second;
|
||||
}
|
||||
|
||||
// Create initial working mask.
|
||||
cv::Mat small_mask_mat(cv::Size(tensor_width, tensor_height), CV_32FC1);
|
||||
|
||||
// Wrap input tensor.
|
||||
auto raw_input_tensor = &input_tensors[0];
|
||||
auto raw_input_view = raw_input_tensor->GetCpuReadView();
|
||||
const float* raw_input_data = raw_input_view.buffer<float>();
|
||||
cv::Mat tensor_mat(cv::Size(tensor_width, tensor_height),
|
||||
CV_MAKETYPE(CV_32F, tensor_channels),
|
||||
const_cast<float*>(raw_input_data));
|
||||
|
||||
// Process mask tensor and apply activation function.
|
||||
if (tensor_channels == 2) {
|
||||
MP_RETURN_IF_ERROR(ApplyActivation<cv::Vec2f>(tensor_mat, &small_mask_mat));
|
||||
} else if (tensor_channels == 1) {
|
||||
RET_CHECK(mediapipe::TensorsToSegmentationCalculatorOptions::SOFTMAX !=
|
||||
options_.activation()); // Requires 2 channels.
|
||||
if (mediapipe::TensorsToSegmentationCalculatorOptions::NONE ==
|
||||
options_.activation()) // Pass-through optimization.
|
||||
tensor_mat.copyTo(small_mask_mat);
|
||||
else
|
||||
MP_RETURN_IF_ERROR(ApplyActivation<float>(tensor_mat, &small_mask_mat));
|
||||
} else {
|
||||
RET_CHECK_FAIL() << "Unsupported number of tensor channels "
|
||||
<< tensor_channels;
|
||||
}
|
||||
|
||||
// Send out image as CPU packet.
|
||||
std::shared_ptr<ImageFrame> mask_frame = std::make_shared<ImageFrame>(
|
||||
ImageFormat::VEC32F1, output_width, output_height);
|
||||
std::unique_ptr<Image> output_mask = absl::make_unique<Image>(mask_frame);
|
||||
cv::Mat output_mat = formats::MatView(output_mask.get());
|
||||
// Upsample small mask into output.
|
||||
cv::resize(small_mask_mat, output_mat, cv::Size(output_width, output_height));
|
||||
cc->Outputs().Tag(kMaskTag).Add(output_mask.release(), cc->InputTimestamp());
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
template <class T>
|
||||
absl::Status TensorsToSegmentationCalculator::ApplyActivation(
|
||||
cv::Mat& tensor_mat, cv::Mat* small_mask_mat) {
|
||||
// Configure activation function.
|
||||
const int output_layer_index = options_.output_layer_index();
|
||||
typedef mediapipe::TensorsToSegmentationCalculatorOptions Options;
|
||||
const auto activation_fn = [&](const cv::Vec2f& mask_value) {
|
||||
float new_mask_value = 0;
|
||||
// TODO consider moving switch out of the loop,
|
||||
// and also avoid float/Vec2f casting.
|
||||
switch (options_.activation()) {
|
||||
case Options::NONE: {
|
||||
new_mask_value = mask_value[0];
|
||||
break;
|
||||
}
|
||||
case Options::SIGMOID: {
|
||||
const float pixel0 = mask_value[0];
|
||||
new_mask_value = 1.0 / (std::exp(-pixel0) + 1.0);
|
||||
break;
|
||||
}
|
||||
case Options::SOFTMAX: {
|
||||
const float pixel0 = mask_value[0];
|
||||
const float pixel1 = mask_value[1];
|
||||
const float max_pixel = std::max(pixel0, pixel1);
|
||||
const float min_pixel = std::min(pixel0, pixel1);
|
||||
const float softmax_denom =
|
||||
/*exp(max_pixel - max_pixel)=*/1.0f +
|
||||
std::exp(min_pixel - max_pixel);
|
||||
new_mask_value = std::exp(mask_value[output_layer_index] - max_pixel) /
|
||||
softmax_denom;
|
||||
break;
|
||||
}
|
||||
}
|
||||
return new_mask_value;
|
||||
};
|
||||
|
||||
// Process mask tensor.
|
||||
for (int i = 0; i < tensor_mat.rows; ++i) {
|
||||
for (int j = 0; j < tensor_mat.cols; ++j) {
|
||||
const T& input_pix = tensor_mat.at<T>(i, j);
|
||||
const float mask_value = activation_fn(input_pix);
|
||||
small_mask_mat->at<float>(i, j) = mask_value;
|
||||
}
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
// Steps:
|
||||
// 1. receive tensor
|
||||
// 2. process segmentation tensor into small mask
|
||||
// 3. upsample small mask into output mask to be same size as input image
|
||||
absl::Status TensorsToSegmentationCalculator::ProcessGpu(
|
||||
CalculatorContext* cc) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
// Get input streams, and dimensions.
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag(kTensorsTag).Get<std::vector<Tensor>>();
|
||||
ASSIGN_OR_RETURN(auto hwc, GetHwcFromDims(input_tensors[0].shape().dims));
|
||||
auto [tensor_height, tensor_width, tensor_channels] = hwc;
|
||||
int output_width = tensor_width, output_height = tensor_height;
|
||||
if (cc->Inputs().HasTag(kOutputSizeTag)) {
|
||||
const auto& size =
|
||||
cc->Inputs().Tag(kOutputSizeTag).Get<std::pair<int, int>>();
|
||||
output_width = size.first;
|
||||
output_height = size.second;
|
||||
}
|
||||
|
||||
// Create initial working mask texture.
|
||||
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
tflite::gpu::gl::GlTexture small_mask_texture;
|
||||
#else
|
||||
mediapipe::GlTexture small_mask_texture;
|
||||
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
|
||||
// Run shader, process mask tensor.
|
||||
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
{
|
||||
MP_RETURN_IF_ERROR(CreateReadWriteRgbaImageTexture(
|
||||
tflite::gpu::DataType::UINT8, // GL_RGBA8
|
||||
{tensor_width, tensor_height}, &small_mask_texture));
|
||||
|
||||
const int output_index = 0;
|
||||
glBindImageTexture(output_index, small_mask_texture.id(), 0, GL_FALSE, 0,
|
||||
GL_WRITE_ONLY, GL_RGBA8);
|
||||
|
||||
auto read_view = input_tensors[0].GetOpenGlBufferReadView();
|
||||
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, read_view.name());
|
||||
|
||||
const tflite::gpu::uint3 workgroups = {
|
||||
NumGroups(tensor_width, kWorkgroupSize),
|
||||
NumGroups(tensor_height, kWorkgroupSize), 1};
|
||||
|
||||
glUseProgram(mask_program_31_->id());
|
||||
glUniform2i(glGetUniformLocation(mask_program_31_->id(), "out_size"),
|
||||
tensor_width, tensor_height);
|
||||
|
||||
MP_RETURN_IF_ERROR(mask_program_31_->Dispatch(workgroups));
|
||||
}
|
||||
#elif MEDIAPIPE_METAL_ENABLED
|
||||
{
|
||||
id<MTLCommandBuffer> command_buffer = [metal_helper_ commandBuffer];
|
||||
command_buffer.label = @"SegmentationKernel";
|
||||
id<MTLComputeCommandEncoder> command_encoder =
|
||||
[command_buffer computeCommandEncoder];
|
||||
[command_encoder setComputePipelineState:mask_program_];
|
||||
|
||||
auto read_view = input_tensors[0].GetMtlBufferReadView(command_buffer);
|
||||
[command_encoder setBuffer:read_view.buffer() offset:0 atIndex:0];
|
||||
|
||||
mediapipe::GpuBuffer small_mask_buffer = [metal_helper_
|
||||
mediapipeGpuBufferWithWidth:tensor_width
|
||||
height:tensor_height
|
||||
format:mediapipe::GpuBufferFormat::kBGRA32];
|
||||
id<MTLTexture> small_mask_texture_metal =
|
||||
[metal_helper_ metalTextureWithGpuBuffer:small_mask_buffer];
|
||||
[command_encoder setTexture:small_mask_texture_metal atIndex:1];
|
||||
|
||||
unsigned int out_size[] = {static_cast<unsigned int>(tensor_width),
|
||||
static_cast<unsigned int>(tensor_height)};
|
||||
[command_encoder setBytes:&out_size length:sizeof(out_size) atIndex:2];
|
||||
|
||||
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, kWorkgroupSize, 1);
|
||||
MTLSize threadgroups =
|
||||
MTLSizeMake(NumGroups(tensor_width, kWorkgroupSize),
|
||||
NumGroups(tensor_height, kWorkgroupSize), 1);
|
||||
[command_encoder dispatchThreadgroups:threadgroups
|
||||
threadsPerThreadgroup:threads_per_group];
|
||||
[command_encoder endEncoding];
|
||||
[command_buffer commit];
|
||||
|
||||
small_mask_texture = gpu_helper_.CreateSourceTexture(small_mask_buffer);
|
||||
}
|
||||
#else
|
||||
{
|
||||
small_mask_texture = gpu_helper_.CreateDestinationTexture(
|
||||
tensor_width, tensor_height,
|
||||
mediapipe::GpuBufferFormat::kBGRA32); // actually GL_RGBA8
|
||||
|
||||
// Go through CPU if not already texture 2D (no direct conversion yet).
|
||||
// Tensor::GetOpenGlTexture2dReadView() doesn't automatically convert types.
|
||||
if (!input_tensors[0].ready_as_opengl_texture_2d()) {
|
||||
(void)input_tensors[0].GetCpuReadView();
|
||||
}
|
||||
|
||||
auto read_view = input_tensors[0].GetOpenGlTexture2dReadView();
|
||||
|
||||
gpu_helper_.BindFramebuffer(small_mask_texture);
|
||||
glActiveTexture(GL_TEXTURE1);
|
||||
glBindTexture(GL_TEXTURE_2D, read_view.name());
|
||||
glUseProgram(mask_program_20_);
|
||||
GlRender();
|
||||
glBindTexture(GL_TEXTURE_2D, 0);
|
||||
glFlush();
|
||||
}
|
||||
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
|
||||
// Upsample small mask into output.
|
||||
mediapipe::GlTexture output_texture = gpu_helper_.CreateDestinationTexture(
|
||||
output_width, output_height,
|
||||
mediapipe::GpuBufferFormat::kBGRA32); // actually GL_RGBA8
|
||||
|
||||
// Run shader, upsample result.
|
||||
{
|
||||
gpu_helper_.BindFramebuffer(output_texture);
|
||||
glActiveTexture(GL_TEXTURE1);
|
||||
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
glBindTexture(GL_TEXTURE_2D, small_mask_texture.id());
|
||||
#else
|
||||
glBindTexture(GL_TEXTURE_2D, small_mask_texture.name());
|
||||
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
glUseProgram(upsample_program_);
|
||||
GlRender();
|
||||
glBindTexture(GL_TEXTURE_2D, 0);
|
||||
glFlush();
|
||||
}
|
||||
|
||||
// Send out image as GPU packet.
|
||||
auto output_image = output_texture.GetFrame<Image>();
|
||||
cc->Outputs().Tag(kMaskTag).Add(output_image.release(), cc->InputTimestamp());
|
||||
|
||||
// Cleanup
|
||||
output_texture.Release();
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
void TensorsToSegmentationCalculator::GlRender() {
|
||||
#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
|
||||
};
|
||||
|
||||
// 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 TensorsToSegmentationCalculator::LoadOptions(
|
||||
CalculatorContext* cc) {
|
||||
// Get calculator options specified in the graph.
|
||||
options_ = cc->Options<::mediapipe::TensorsToSegmentationCalculatorOptions>();
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status TensorsToSegmentationCalculator::InitGpu(CalculatorContext* cc) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> absl::Status {
|
||||
// A shader to process a segmentation tensor into an output mask.
|
||||
// Currently uses 4 channels for output, and sets R+A channels as mask value.
|
||||
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
// GLES 3.1
|
||||
const tflite::gpu::uint3 workgroup_size = {kWorkgroupSize, kWorkgroupSize,
|
||||
1};
|
||||
const std::string shader_header =
|
||||
absl::StrCat(tflite::gpu::gl::GetShaderHeader(workgroup_size), R"(
|
||||
precision highp float;
|
||||
|
||||
layout(rgba8, binding = 0) writeonly uniform highp image2D output_texture;
|
||||
|
||||
uniform ivec2 out_size;
|
||||
)");
|
||||
/* Shader defines will be inserted here. */
|
||||
|
||||
const std::string shader_src_main = R"(
|
||||
layout(std430, binding = 2) readonly buffer B0 {
|
||||
#ifdef TWO_CHANNEL_INPUT
|
||||
vec2 elements[];
|
||||
#else
|
||||
float elements[];
|
||||
#endif // TWO_CHANNEL_INPUT
|
||||
} input_data; // data tensor
|
||||
|
||||
void main() {
|
||||
int out_width = out_size.x;
|
||||
int out_height = out_size.y;
|
||||
|
||||
ivec2 gid = ivec2(gl_GlobalInvocationID.xy);
|
||||
if (gid.x >= out_width || gid.y >= out_height) { return; }
|
||||
int linear_index = gid.y * out_width + gid.x;
|
||||
|
||||
#ifdef TWO_CHANNEL_INPUT
|
||||
vec2 input_value = input_data.elements[linear_index];
|
||||
#else
|
||||
vec2 input_value = vec2(input_data.elements[linear_index], 0.0);
|
||||
#endif // TWO_CHANNEL_INPUT
|
||||
|
||||
// Run activation function.
|
||||
// One and only one of FN_SOFTMAX,FN_SIGMOID,FN_NONE will be defined.
|
||||
#ifdef FN_SOFTMAX
|
||||
// Only two channel input tensor is supported.
|
||||
vec2 input_px = input_value.rg;
|
||||
float shift = max(input_px.r, input_px.g);
|
||||
float softmax_denom = exp(input_px.r - shift) + exp(input_px.g - shift);
|
||||
float new_mask_value =
|
||||
exp(input_px[OUTPUT_LAYER_INDEX] - shift) / softmax_denom;
|
||||
#endif // FN_SOFTMAX
|
||||
|
||||
#ifdef FN_SIGMOID
|
||||
float new_mask_value = 1.0 / (exp(-input_value.r) + 1.0);
|
||||
#endif // FN_SIGMOID
|
||||
|
||||
#ifdef FN_NONE
|
||||
float new_mask_value = input_value.r;
|
||||
#endif // FN_NONE
|
||||
|
||||
#ifdef FLIP_Y_COORD
|
||||
int y_coord = out_height - gid.y - 1;
|
||||
#else
|
||||
int y_coord = gid.y;
|
||||
#endif // defined(FLIP_Y_COORD)
|
||||
ivec2 output_coordinate = ivec2(gid.x, y_coord);
|
||||
|
||||
vec4 out_value = vec4(new_mask_value, 0.0, 0.0, new_mask_value);
|
||||
imageStore(output_texture, output_coordinate, out_value);
|
||||
})";
|
||||
|
||||
#elif MEDIAPIPE_METAL_ENABLED
|
||||
// METAL
|
||||
const std::string shader_header = R"(
|
||||
#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
)";
|
||||
/* Shader defines will be inserted here. */
|
||||
|
||||
const std::string shader_src_main = R"(
|
||||
kernel void segmentationKernel(
|
||||
#ifdef TWO_CHANNEL_INPUT
|
||||
device float2* elements [[ buffer(0) ]],
|
||||
#else
|
||||
device float* elements [[ buffer(0) ]],
|
||||
#endif // TWO_CHANNEL_INPUT
|
||||
texture2d<float, access::write> output_texture [[ texture(1) ]],
|
||||
constant uint* out_size [[ buffer(2) ]],
|
||||
uint2 gid [[ thread_position_in_grid ]])
|
||||
{
|
||||
uint out_width = out_size[0];
|
||||
uint out_height = out_size[1];
|
||||
|
||||
if (gid.x >= out_width || gid.y >= out_height) { return; }
|
||||
uint linear_index = gid.y * out_width + gid.x;
|
||||
|
||||
#ifdef TWO_CHANNEL_INPUT
|
||||
float2 input_value = elements[linear_index];
|
||||
#else
|
||||
float2 input_value = float2(elements[linear_index], 0.0);
|
||||
#endif // TWO_CHANNEL_INPUT
|
||||
|
||||
// Run activation function.
|
||||
// One and only one of FN_SOFTMAX,FN_SIGMOID,FN_NONE will be defined.
|
||||
#ifdef FN_SOFTMAX
|
||||
// Only two channel input tensor is supported.
|
||||
float2 input_px = input_value.xy;
|
||||
float shift = max(input_px.x, input_px.y);
|
||||
float softmax_denom = exp(input_px.r - shift) + exp(input_px.g - shift);
|
||||
float new_mask_value =
|
||||
exp(input_px[OUTPUT_LAYER_INDEX] - shift) / softmax_denom;
|
||||
#endif // FN_SOFTMAX
|
||||
|
||||
#ifdef FN_SIGMOID
|
||||
float new_mask_value = 1.0 / (exp(-input_value.x) + 1.0);
|
||||
#endif // FN_SIGMOID
|
||||
|
||||
#ifdef FN_NONE
|
||||
float new_mask_value = input_value.x;
|
||||
#endif // FN_NONE
|
||||
|
||||
#ifdef FLIP_Y_COORD
|
||||
int y_coord = out_height - gid.y - 1;
|
||||
#else
|
||||
int y_coord = gid.y;
|
||||
#endif // defined(FLIP_Y_COORD)
|
||||
uint2 output_coordinate = uint2(gid.x, y_coord);
|
||||
|
||||
float4 out_value = float4(new_mask_value, 0.0, 0.0, new_mask_value);
|
||||
output_texture.write(out_value, output_coordinate);
|
||||
}
|
||||
)";
|
||||
|
||||
#else
|
||||
// GLES 2.0
|
||||
const std::string shader_header = absl::StrCat(
|
||||
std::string(mediapipe::kMediaPipeFragmentShaderPreamble), R"(
|
||||
DEFAULT_PRECISION(mediump, float)
|
||||
)");
|
||||
/* Shader defines will be inserted here. */
|
||||
|
||||
const std::string shader_src_main = R"(
|
||||
in vec2 sample_coordinate;
|
||||
|
||||
uniform sampler2D input_texture;
|
||||
|
||||
#ifdef GL_ES
|
||||
#define fragColor gl_FragColor
|
||||
#else
|
||||
out vec4 fragColor;
|
||||
#endif // defined(GL_ES);
|
||||
|
||||
void main() {
|
||||
|
||||
vec4 input_value = texture2D(input_texture, sample_coordinate);
|
||||
vec2 gid = sample_coordinate;
|
||||
|
||||
// Run activation function.
|
||||
// One and only one of FN_SOFTMAX,FN_SIGMOID,FN_NONE will be defined.
|
||||
|
||||
#ifdef FN_SOFTMAX
|
||||
// Only two channel input tensor is supported.
|
||||
vec2 input_px = input_value.rg;
|
||||
float shift = max(input_px.r, input_px.g);
|
||||
float softmax_denom = exp(input_px.r - shift) + exp(input_px.g - shift);
|
||||
float new_mask_value =
|
||||
exp(mix(input_px.r, input_px.g, float(OUTPUT_LAYER_INDEX)) - shift) / softmax_denom;
|
||||
#endif // FN_SOFTMAX
|
||||
|
||||
#ifdef FN_SIGMOID
|
||||
float new_mask_value = 1.0 / (exp(-input_value.r) + 1.0);
|
||||
#endif // FN_SIGMOID
|
||||
|
||||
#ifdef FN_NONE
|
||||
float new_mask_value = input_value.r;
|
||||
#endif // FN_NONE
|
||||
|
||||
#ifdef FLIP_Y_COORD
|
||||
float y_coord = 1.0 - gid.y;
|
||||
#else
|
||||
float y_coord = gid.y;
|
||||
#endif // defined(FLIP_Y_COORD)
|
||||
vec2 output_coordinate = vec2(gid.x, y_coord);
|
||||
|
||||
vec4 out_value = vec4(new_mask_value, 0.0, 0.0, new_mask_value);
|
||||
fragColor = out_value;
|
||||
})";
|
||||
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
|
||||
// Shader defines.
|
||||
typedef mediapipe::TensorsToSegmentationCalculatorOptions Options;
|
||||
const std::string output_layer_index =
|
||||
"\n#define OUTPUT_LAYER_INDEX int(" +
|
||||
std::to_string(options_.output_layer_index()) + ")";
|
||||
const std::string flip_y_coord =
|
||||
DoesGpuTextureStartAtBottom() ? "\n#define FLIP_Y_COORD" : "";
|
||||
const std::string fn_none =
|
||||
options_.activation() == Options::NONE ? "\n#define FN_NONE" : "";
|
||||
const std::string fn_sigmoid =
|
||||
options_.activation() == Options::SIGMOID ? "\n#define FN_SIGMOID" : "";
|
||||
const std::string fn_softmax =
|
||||
options_.activation() == Options::SOFTMAX ? "\n#define FN_SOFTMAX" : "";
|
||||
const std::string two_channel = options_.activation() == Options::SOFTMAX
|
||||
? "\n#define TWO_CHANNEL_INPUT"
|
||||
: "";
|
||||
const std::string shader_defines =
|
||||
absl::StrCat(output_layer_index, flip_y_coord, fn_softmax, fn_sigmoid,
|
||||
fn_none, two_channel);
|
||||
|
||||
// Build full shader.
|
||||
const std::string shader_src_no_previous =
|
||||
absl::StrCat(shader_header, shader_defines, shader_src_main);
|
||||
|
||||
// Vertex shader attributes.
|
||||
const GLint attr_location[NUM_ATTRIBUTES] = {
|
||||
ATTRIB_VERTEX,
|
||||
ATTRIB_TEXTURE_POSITION,
|
||||
};
|
||||
const GLchar* attr_name[NUM_ATTRIBUTES] = {
|
||||
"position",
|
||||
"texture_coordinate",
|
||||
};
|
||||
|
||||
// Main shader program & parameters
|
||||
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
GlShader shader_without_previous;
|
||||
MP_RETURN_IF_ERROR(GlShader::CompileShader(
|
||||
GL_COMPUTE_SHADER, shader_src_no_previous, &shader_without_previous));
|
||||
mask_program_31_ = absl::make_unique<GlProgram>();
|
||||
MP_RETURN_IF_ERROR(GlProgram::CreateWithShader(shader_without_previous,
|
||||
mask_program_31_.get()));
|
||||
#elif MEDIAPIPE_METAL_ENABLED
|
||||
id<MTLDevice> device = metal_helper_.mtlDevice;
|
||||
NSString* library_source =
|
||||
[NSString stringWithUTF8String:shader_src_no_previous.c_str()];
|
||||
NSError* error = nil;
|
||||
id<MTLLibrary> library = [device newLibraryWithSource:library_source
|
||||
options:nullptr
|
||||
error:&error];
|
||||
RET_CHECK(library != nil) << "Couldn't create shader library "
|
||||
<< [[error localizedDescription] UTF8String];
|
||||
id<MTLFunction> kernel_func = nil;
|
||||
kernel_func = [library newFunctionWithName:@"segmentationKernel"];
|
||||
RET_CHECK(kernel_func != nil) << "Couldn't create kernel function.";
|
||||
mask_program_ =
|
||||
[device newComputePipelineStateWithFunction:kernel_func error:&error];
|
||||
RET_CHECK(mask_program_ != nil) << "Couldn't create pipeline state " <<
|
||||
[[error localizedDescription] UTF8String];
|
||||
#else
|
||||
mediapipe::GlhCreateProgram(
|
||||
mediapipe::kBasicVertexShader, shader_src_no_previous.c_str(),
|
||||
NUM_ATTRIBUTES, &attr_name[0], attr_location, &mask_program_20_);
|
||||
RET_CHECK(mask_program_20_) << "Problem initializing the program.";
|
||||
glUseProgram(mask_program_20_);
|
||||
glUniform1i(glGetUniformLocation(mask_program_20_, "input_texture"), 1);
|
||||
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
|
||||
|
||||
// Simple pass-through program, used for hardware upsampling.
|
||||
mediapipe::GlhCreateProgram(
|
||||
mediapipe::kBasicVertexShader, mediapipe::kBasicTexturedFragmentShader,
|
||||
NUM_ATTRIBUTES, &attr_name[0], attr_location, &upsample_program_);
|
||||
RET_CHECK(upsample_program_) << "Problem initializing the program.";
|
||||
glUseProgram(upsample_program_);
|
||||
glUniform1i(glGetUniformLocation(upsample_program_, "video_frame"), 1);
|
||||
|
||||
return absl::OkStatus();
|
||||
}));
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,46 @@
|
||||
// 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";
|
||||
import "mediapipe/gpu/gpu_origin.proto";
|
||||
|
||||
message TensorsToSegmentationCalculatorOptions {
|
||||
extend mediapipe.CalculatorOptions {
|
||||
optional TensorsToSegmentationCalculatorOptions ext = 374311106;
|
||||
}
|
||||
|
||||
// For CONVENTIONAL mode in OpenGL, textures start at bottom and needs
|
||||
// to be flipped vertically as tensors are expected to start at top.
|
||||
// (DEFAULT or unset is interpreted as CONVENTIONAL.)
|
||||
optional GpuOrigin.Mode gpu_origin = 1;
|
||||
|
||||
// Supported activation functions for filtering.
|
||||
enum Activation {
|
||||
NONE = 0; // Assumes 1-channel input tensor.
|
||||
SIGMOID = 1; // Assumes 1-channel input tensor.
|
||||
SOFTMAX = 2; // Assumes 2-channel input tensor.
|
||||
}
|
||||
// Activation function to apply to input tensor.
|
||||
// Softmax requires a 2-channel tensor, see output_layer_index below.
|
||||
optional Activation activation = 2 [default = NONE];
|
||||
|
||||
// Channel to use for processing tensor.
|
||||
// Only applies when using activation=SOFTMAX.
|
||||
// Works on two channel input tensor only.
|
||||
optional int32 output_layer_index = 3 [default = 1];
|
||||
}
|
||||
@@ -859,6 +859,7 @@ cc_library(
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/util/filtering:one_euro_filter",
|
||||
"//mediapipe/util/filtering:relative_velocity_filter",
|
||||
|
||||
@@ -323,7 +323,7 @@ absl::Status DetectionsToRectsCalculator::ComputeRotation(
|
||||
DetectionSpec DetectionsToRectsCalculator::GetDetectionSpec(
|
||||
const CalculatorContext* cc) {
|
||||
absl::optional<std::pair<int, int>> image_size;
|
||||
if (cc->Inputs().HasTag(kImageSizeTag)) {
|
||||
if (HasTagValue(cc->Inputs(), kImageSizeTag)) {
|
||||
image_size = cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
|
||||
}
|
||||
|
||||
|
||||
@@ -157,6 +157,12 @@ TEST(DetectionsToRectsCalculatorTest, DetectionKeyPointsToRect) {
|
||||
/*image_size=*/{640, 480});
|
||||
MP_ASSERT_OK(status_or_value);
|
||||
EXPECT_THAT(status_or_value.value(), RectEq(480, 360, 320, 240));
|
||||
|
||||
status_or_value = RunDetectionKeyPointsToRectCalculation(
|
||||
/*detection=*/DetectionWithKeyPoints({{0.25f, 0.25f}, {0.75f, 0.75f}}),
|
||||
/*image_size=*/{0, 0});
|
||||
MP_ASSERT_OK(status_or_value);
|
||||
EXPECT_THAT(status_or_value.value(), RectEq(0, 0, 0, 0));
|
||||
}
|
||||
|
||||
TEST(DetectionsToRectsCalculatorTest, DetectionToNormalizedRect) {
|
||||
|
||||
@@ -18,6 +18,7 @@
|
||||
#include "mediapipe/calculators/util/landmarks_smoothing_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
#include "mediapipe/util/filtering/one_euro_filter.h"
|
||||
@@ -30,6 +31,7 @@ namespace {
|
||||
constexpr char kNormalizedLandmarksTag[] = "NORM_LANDMARKS";
|
||||
constexpr char kLandmarksTag[] = "LANDMARKS";
|
||||
constexpr char kImageSizeTag[] = "IMAGE_SIZE";
|
||||
constexpr char kObjectScaleRoiTag[] = "OBJECT_SCALE_ROI";
|
||||
constexpr char kNormalizedFilteredLandmarksTag[] = "NORM_FILTERED_LANDMARKS";
|
||||
constexpr char kFilteredLandmarksTag[] = "FILTERED_LANDMARKS";
|
||||
|
||||
@@ -94,6 +96,18 @@ float GetObjectScale(const LandmarkList& landmarks) {
|
||||
return (object_width + object_height) / 2.0f;
|
||||
}
|
||||
|
||||
float GetObjectScale(const NormalizedRect& roi, const int image_width,
|
||||
const int image_height) {
|
||||
const float object_width = roi.width() * image_width;
|
||||
const float object_height = roi.height() * image_height;
|
||||
|
||||
return (object_width + object_height) / 2.0f;
|
||||
}
|
||||
|
||||
float GetObjectScale(const Rect& roi) {
|
||||
return (roi.width() + roi.height()) / 2.0f;
|
||||
}
|
||||
|
||||
// Abstract class for various landmarks filters.
|
||||
class LandmarksFilter {
|
||||
public:
|
||||
@@ -103,6 +117,7 @@ class LandmarksFilter {
|
||||
|
||||
virtual absl::Status Apply(const LandmarkList& in_landmarks,
|
||||
const absl::Duration& timestamp,
|
||||
const absl::optional<float> object_scale_opt,
|
||||
LandmarkList* out_landmarks) = 0;
|
||||
};
|
||||
|
||||
@@ -111,6 +126,7 @@ class NoFilter : public LandmarksFilter {
|
||||
public:
|
||||
absl::Status Apply(const LandmarkList& in_landmarks,
|
||||
const absl::Duration& timestamp,
|
||||
const absl::optional<float> object_scale_opt,
|
||||
LandmarkList* out_landmarks) override {
|
||||
*out_landmarks = in_landmarks;
|
||||
return absl::OkStatus();
|
||||
@@ -136,13 +152,15 @@ class VelocityFilter : public LandmarksFilter {
|
||||
|
||||
absl::Status Apply(const LandmarkList& in_landmarks,
|
||||
const absl::Duration& timestamp,
|
||||
const absl::optional<float> object_scale_opt,
|
||||
LandmarkList* out_landmarks) override {
|
||||
// Get value scale as inverse value of the object scale.
|
||||
// If value is too small smoothing will be disabled and landmarks will be
|
||||
// returned as is.
|
||||
float value_scale = 1.0f;
|
||||
if (!disable_value_scaling_) {
|
||||
const float object_scale = GetObjectScale(in_landmarks);
|
||||
const float object_scale =
|
||||
object_scale_opt ? *object_scale_opt : GetObjectScale(in_landmarks);
|
||||
if (object_scale < min_allowed_object_scale_) {
|
||||
*out_landmarks = in_landmarks;
|
||||
return absl::OkStatus();
|
||||
@@ -205,12 +223,14 @@ class VelocityFilter : public LandmarksFilter {
|
||||
class OneEuroFilterImpl : public LandmarksFilter {
|
||||
public:
|
||||
OneEuroFilterImpl(double frequency, double min_cutoff, double beta,
|
||||
double derivate_cutoff, float min_allowed_object_scale)
|
||||
double derivate_cutoff, float min_allowed_object_scale,
|
||||
bool disable_value_scaling)
|
||||
: frequency_(frequency),
|
||||
min_cutoff_(min_cutoff),
|
||||
beta_(beta),
|
||||
derivate_cutoff_(derivate_cutoff),
|
||||
min_allowed_object_scale_(min_allowed_object_scale) {}
|
||||
min_allowed_object_scale_(min_allowed_object_scale),
|
||||
disable_value_scaling_(disable_value_scaling) {}
|
||||
|
||||
absl::Status Reset() override {
|
||||
x_filters_.clear();
|
||||
@@ -221,16 +241,24 @@ class OneEuroFilterImpl : public LandmarksFilter {
|
||||
|
||||
absl::Status Apply(const LandmarkList& in_landmarks,
|
||||
const absl::Duration& timestamp,
|
||||
const absl::optional<float> object_scale_opt,
|
||||
LandmarkList* out_landmarks) override {
|
||||
// Initialize filters once.
|
||||
MP_RETURN_IF_ERROR(InitializeFiltersIfEmpty(in_landmarks.landmark_size()));
|
||||
|
||||
const float object_scale = GetObjectScale(in_landmarks);
|
||||
if (object_scale < min_allowed_object_scale_) {
|
||||
*out_landmarks = in_landmarks;
|
||||
return absl::OkStatus();
|
||||
// Get value scale as inverse value of the object scale.
|
||||
// If value is too small smoothing will be disabled and landmarks will be
|
||||
// returned as is.
|
||||
float value_scale = 1.0f;
|
||||
if (!disable_value_scaling_) {
|
||||
const float object_scale =
|
||||
object_scale_opt ? *object_scale_opt : GetObjectScale(in_landmarks);
|
||||
if (object_scale < min_allowed_object_scale_) {
|
||||
*out_landmarks = in_landmarks;
|
||||
return absl::OkStatus();
|
||||
}
|
||||
value_scale = 1.0f / object_scale;
|
||||
}
|
||||
const float value_scale = 1.0f / object_scale;
|
||||
|
||||
// Filter landmarks. Every axis of every landmark is filtered separately.
|
||||
for (int i = 0; i < in_landmarks.landmark_size(); ++i) {
|
||||
@@ -277,6 +305,7 @@ class OneEuroFilterImpl : public LandmarksFilter {
|
||||
double beta_;
|
||||
double derivate_cutoff_;
|
||||
double min_allowed_object_scale_;
|
||||
bool disable_value_scaling_;
|
||||
|
||||
std::vector<OneEuroFilter> x_filters_;
|
||||
std::vector<OneEuroFilter> y_filters_;
|
||||
@@ -292,6 +321,10 @@ class OneEuroFilterImpl : public LandmarksFilter {
|
||||
// IMAGE_SIZE: A std::pair<int, int> represention of image width and height.
|
||||
// Required to perform all computations in absolute coordinates to avoid any
|
||||
// influence of normalized values.
|
||||
// OBJECT_SCALE_ROI (optional): A NormRect or Rect (depending on the format of
|
||||
// input landmarks) used to determine the object scale for some of the
|
||||
// filters. If not provided - object scale will be calculated from
|
||||
// landmarks.
|
||||
//
|
||||
// Outputs:
|
||||
// NORM_FILTERED_LANDMARKS: A NormalizedLandmarkList of smoothed landmarks.
|
||||
@@ -301,6 +334,7 @@ class OneEuroFilterImpl : public LandmarksFilter {
|
||||
// calculator: "LandmarksSmoothingCalculator"
|
||||
// input_stream: "NORM_LANDMARKS:pose_landmarks"
|
||||
// input_stream: "IMAGE_SIZE:image_size"
|
||||
// input_stream: "OBJECT_SCALE_ROI:roi"
|
||||
// output_stream: "NORM_FILTERED_LANDMARKS:pose_landmarks_filtered"
|
||||
// options: {
|
||||
// [mediapipe.LandmarksSmoothingCalculatorOptions.ext] {
|
||||
@@ -330,9 +364,17 @@ absl::Status LandmarksSmoothingCalculator::GetContract(CalculatorContract* cc) {
|
||||
cc->Outputs()
|
||||
.Tag(kNormalizedFilteredLandmarksTag)
|
||||
.Set<NormalizedLandmarkList>();
|
||||
|
||||
if (cc->Inputs().HasTag(kObjectScaleRoiTag)) {
|
||||
cc->Inputs().Tag(kObjectScaleRoiTag).Set<NormalizedRect>();
|
||||
}
|
||||
} else {
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<LandmarkList>();
|
||||
cc->Outputs().Tag(kFilteredLandmarksTag).Set<LandmarkList>();
|
||||
|
||||
if (cc->Inputs().HasTag(kObjectScaleRoiTag)) {
|
||||
cc->Inputs().Tag(kObjectScaleRoiTag).Set<Rect>();
|
||||
}
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
@@ -357,7 +399,8 @@ absl::Status LandmarksSmoothingCalculator::Open(CalculatorContext* cc) {
|
||||
options.one_euro_filter().min_cutoff(),
|
||||
options.one_euro_filter().beta(),
|
||||
options.one_euro_filter().derivate_cutoff(),
|
||||
options.one_euro_filter().min_allowed_object_scale());
|
||||
options.one_euro_filter().min_allowed_object_scale(),
|
||||
options.one_euro_filter().disable_value_scaling());
|
||||
} else {
|
||||
RET_CHECK_FAIL()
|
||||
<< "Landmarks filter is either not specified or not supported";
|
||||
@@ -389,13 +432,20 @@ absl::Status LandmarksSmoothingCalculator::Process(CalculatorContext* cc) {
|
||||
std::tie(image_width, image_height) =
|
||||
cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
|
||||
|
||||
absl::optional<float> object_scale;
|
||||
if (cc->Inputs().HasTag(kObjectScaleRoiTag) &&
|
||||
!cc->Inputs().Tag(kObjectScaleRoiTag).IsEmpty()) {
|
||||
auto& roi = cc->Inputs().Tag(kObjectScaleRoiTag).Get<NormalizedRect>();
|
||||
object_scale = GetObjectScale(roi, image_width, image_height);
|
||||
}
|
||||
|
||||
auto in_landmarks = absl::make_unique<LandmarkList>();
|
||||
NormalizedLandmarksToLandmarks(in_norm_landmarks, image_width, image_height,
|
||||
in_landmarks.get());
|
||||
|
||||
auto out_landmarks = absl::make_unique<LandmarkList>();
|
||||
MP_RETURN_IF_ERROR(landmarks_filter_->Apply(*in_landmarks, timestamp,
|
||||
out_landmarks.get()));
|
||||
MP_RETURN_IF_ERROR(landmarks_filter_->Apply(
|
||||
*in_landmarks, timestamp, object_scale, out_landmarks.get()));
|
||||
|
||||
auto out_norm_landmarks = absl::make_unique<NormalizedLandmarkList>();
|
||||
LandmarksToNormalizedLandmarks(*out_landmarks, image_width, image_height,
|
||||
@@ -408,9 +458,16 @@ absl::Status LandmarksSmoothingCalculator::Process(CalculatorContext* cc) {
|
||||
const auto& in_landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<LandmarkList>();
|
||||
|
||||
absl::optional<float> object_scale;
|
||||
if (cc->Inputs().HasTag(kObjectScaleRoiTag) &&
|
||||
!cc->Inputs().Tag(kObjectScaleRoiTag).IsEmpty()) {
|
||||
auto& roi = cc->Inputs().Tag(kObjectScaleRoiTag).Get<Rect>();
|
||||
object_scale = GetObjectScale(roi);
|
||||
}
|
||||
|
||||
auto out_landmarks = absl::make_unique<LandmarkList>();
|
||||
MP_RETURN_IF_ERROR(
|
||||
landmarks_filter_->Apply(in_landmarks, timestamp, out_landmarks.get()));
|
||||
MP_RETURN_IF_ERROR(landmarks_filter_->Apply(
|
||||
in_landmarks, timestamp, object_scale, out_landmarks.get()));
|
||||
|
||||
cc->Outputs()
|
||||
.Tag(kFilteredLandmarksTag)
|
||||
|
||||
@@ -41,9 +41,9 @@ message LandmarksSmoothingCalculatorOptions {
|
||||
optional float min_allowed_object_scale = 3 [default = 1e-6];
|
||||
|
||||
// Disable value scaling based on object size and use `1.0` instead.
|
||||
// Value scale is calculated as inverse value of object size. Object size is
|
||||
// calculated as maximum side of rectangular bounding box of the object in
|
||||
// XY plane.
|
||||
// If not disabled, value scale is calculated as inverse value of object
|
||||
// size. Object size is calculated as maximum side of rectangular bounding
|
||||
// box of the object in XY plane.
|
||||
optional bool disable_value_scaling = 4 [default = false];
|
||||
}
|
||||
|
||||
@@ -72,6 +72,12 @@ message LandmarksSmoothingCalculatorOptions {
|
||||
// If calculated object scale is less than given value smoothing will be
|
||||
// disabled and landmarks will be returned as is.
|
||||
optional float min_allowed_object_scale = 5 [default = 1e-6];
|
||||
|
||||
// Disable value scaling based on object size and use `1.0` instead.
|
||||
// If not disabled, value scale is calculated as inverse value of object
|
||||
// size. Object size is calculated as maximum side of rectangular bounding
|
||||
// box of the object in XY plane.
|
||||
optional bool disable_value_scaling = 6 [default = false];
|
||||
}
|
||||
|
||||
oneof filter_options {
|
||||
|
||||
@@ -40,7 +40,7 @@ constexpr char kRectTag[] = "NORM_RECT";
|
||||
// Input:
|
||||
// LANDMARKS: A LandmarkList representing world landmarks in the rectangle.
|
||||
// NORM_RECT: An NormalizedRect representing a normalized rectangle in image
|
||||
// coordinates.
|
||||
// coordinates. (Optional)
|
||||
//
|
||||
// Output:
|
||||
// LANDMARKS: A LandmarkList representing world landmarks projected (rotated
|
||||
@@ -59,7 +59,9 @@ class WorldLandmarkProjectionCalculator : public CalculatorBase {
|
||||
public:
|
||||
static absl::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<LandmarkList>();
|
||||
cc->Inputs().Tag(kRectTag).Set<NormalizedRect>();
|
||||
if (cc->Inputs().HasTag(kRectTag)) {
|
||||
cc->Inputs().Tag(kRectTag).Set<NormalizedRect>();
|
||||
}
|
||||
cc->Outputs().Tag(kLandmarksTag).Set<LandmarkList>();
|
||||
|
||||
return absl::OkStatus();
|
||||
@@ -74,13 +76,24 @@ class WorldLandmarkProjectionCalculator : public CalculatorBase {
|
||||
absl::Status Process(CalculatorContext* cc) override {
|
||||
// Check that landmarks and rect are not empty.
|
||||
if (cc->Inputs().Tag(kLandmarksTag).IsEmpty() ||
|
||||
cc->Inputs().Tag(kRectTag).IsEmpty()) {
|
||||
(cc->Inputs().HasTag(kRectTag) &&
|
||||
cc->Inputs().Tag(kRectTag).IsEmpty())) {
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
const auto& in_landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<LandmarkList>();
|
||||
const auto& in_rect = cc->Inputs().Tag(kRectTag).Get<NormalizedRect>();
|
||||
std::function<void(const Landmark&, Landmark*)> rotate_fn;
|
||||
if (cc->Inputs().HasTag(kRectTag)) {
|
||||
const auto& in_rect = cc->Inputs().Tag(kRectTag).Get<NormalizedRect>();
|
||||
const float cosa = std::cos(in_rect.rotation());
|
||||
const float sina = std::sin(in_rect.rotation());
|
||||
rotate_fn = [cosa, sina](const Landmark& in_landmark,
|
||||
Landmark* out_landmark) {
|
||||
out_landmark->set_x(cosa * in_landmark.x() - sina * in_landmark.y());
|
||||
out_landmark->set_y(sina * in_landmark.x() + cosa * in_landmark.y());
|
||||
};
|
||||
}
|
||||
|
||||
auto out_landmarks = absl::make_unique<LandmarkList>();
|
||||
for (int i = 0; i < in_landmarks.landmark_size(); ++i) {
|
||||
@@ -89,11 +102,9 @@ class WorldLandmarkProjectionCalculator : public CalculatorBase {
|
||||
Landmark* out_landmark = out_landmarks->add_landmark();
|
||||
*out_landmark = in_landmark;
|
||||
|
||||
const float angle = in_rect.rotation();
|
||||
out_landmark->set_x(std::cos(angle) * in_landmark.x() -
|
||||
std::sin(angle) * in_landmark.y());
|
||||
out_landmark->set_y(std::sin(angle) * in_landmark.x() +
|
||||
std::cos(angle) * in_landmark.y());
|
||||
if (rotate_fn) {
|
||||
rotate_fn(in_landmark, out_landmark);
|
||||
}
|
||||
}
|
||||
|
||||
cc->Outputs()
|
||||
|
||||
Reference in New Issue
Block a user