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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: ff83882955f1a1e2a043ff4e71278be9d7217bbe
This commit is contained in:
MediaPipe Team
2021-05-05 14:56:16 -04:00
committed by chuoling
parent ecb5b5f44a
commit a9b643e0f5
210 changed files with 5312 additions and 3838 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/core/BUILD
+29 -2
View File
@@ -451,8 +451,8 @@ cc_library(
)
cc_library(
name = "nonzero_calculator",
srcs = ["nonzero_calculator.cc"],
name = "non_zero_calculator",
srcs = ["non_zero_calculator.cc"],
visibility = [
"//visibility:public",
],
@@ -464,6 +464,21 @@ cc_library(
alwayslink = 1,
)
cc_test(
name = "non_zero_calculator_test",
size = "small",
srcs = ["non_zero_calculator_test.cc"],
deps = [
":non_zero_calculator",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework:timestamp",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:status",
"//mediapipe/framework/tool:validate_type",
],
)
cc_test(
name = "mux_calculator_test",
srcs = ["mux_calculator_test.cc"],
@@ -665,6 +680,18 @@ cc_library(
alwayslink = 1,
)
cc_library(
name = "default_side_packet_calculator",
srcs = ["default_side_packet_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
],
alwayslink = 1,
)
cc_library(
name = "side_packet_to_stream_calculator",
srcs = ["side_packet_to_stream_calculator.cc"],
mediapipe/calculators/core/default_side_packet_calculator.cc
+103
View File
@@ -0,0 +1,103 @@
// Copyright 2019 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 "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
namespace mediapipe {
namespace {
constexpr char kOptionalValueTag[] = "OPTIONAL_VALUE";
constexpr char kDefaultValueTag[] = "DEFAULT_VALUE";
constexpr char kValueTag[] = "VALUE";
} // namespace
// Outputs side packet default value if optional value is not provided.
//
// This calculator utilizes the fact that MediaPipe automatically removes
// optional side packets of the calculator configuration (i.e. OPTIONAL_VALUE).
// And if it happens - returns default value, otherwise - returns optional
// value.
//
// Input:
// OPTIONAL_VALUE (optional) - AnyType (but same type as DEFAULT_VALUE)
// Optional side packet value that is outputted by the calculator as is if
// provided.
//
// DEFAULT_VALUE - AnyType
// Default side pack value that is outputted by the calculator if
// OPTIONAL_VALUE is not provided.
//
// Output:
// VALUE - AnyType (but same type as DEFAULT_VALUE)
// Either OPTIONAL_VALUE (if provided) or DEFAULT_VALUE (otherwise).
//
// Usage example:
// node {
// calculator: "DefaultSidePacketCalculator"
// input_side_packet: "OPTIONAL_VALUE:segmentation_mask_enabled_optional"
// input_side_packet: "DEFAULT_VALUE:segmentation_mask_enabled_default"
// output_side_packet: "VALUE:segmentation_mask_enabled"
// }
class DefaultSidePacketCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
};
REGISTER_CALCULATOR(DefaultSidePacketCalculator);
absl::Status DefaultSidePacketCalculator::GetContract(CalculatorContract* cc) {
RET_CHECK(cc->InputSidePackets().HasTag(kDefaultValueTag))
<< "Default value must be provided";
cc->InputSidePackets().Tag(kDefaultValueTag).SetAny();
// Optional input side packet can be unspecified. In this case MediaPipe will
// remove it from the calculator config.
if (cc->InputSidePackets().HasTag(kOptionalValueTag)) {
cc->InputSidePackets()
.Tag(kOptionalValueTag)
.SetSameAs(&cc->InputSidePackets().Tag(kDefaultValueTag));
}
RET_CHECK(cc->OutputSidePackets().HasTag(kValueTag));
cc->OutputSidePackets().Tag(kValueTag).SetSameAs(
&cc->InputSidePackets().Tag(kDefaultValueTag));
return absl::OkStatus();
}
absl::Status DefaultSidePacketCalculator::Open(CalculatorContext* cc) {
// If optional value is provided it is returned as the calculator output.
if (cc->InputSidePackets().HasTag(kOptionalValueTag)) {
auto& packet = cc->InputSidePackets().Tag(kOptionalValueTag);
cc->OutputSidePackets().Tag(kValueTag).Set(packet);
return absl::OkStatus();
}
// If no optional value
auto& packet = cc->InputSidePackets().Tag(kDefaultValueTag);
cc->OutputSidePackets().Tag(kValueTag).Set(packet);
return absl::OkStatus();
}
absl::Status DefaultSidePacketCalculator::Process(CalculatorContext* cc) {
return absl::OkStatus();
}
} // namespace mediapipe
mediapipe/calculators/core/nonzero_calculator.cc → mediapipe/calculators/core/non_zero_calculator.cc
+16 -4
View File
@@ -23,14 +23,26 @@ namespace api2 {
class NonZeroCalculator : public Node {
public:
static constexpr Input<int>::SideFallback kIn{"INPUT"};
static constexpr Output<int> kOut{"OUTPUT"};
static constexpr Output<int>::Optional kOut{"OUTPUT"};
static constexpr Output<bool>::Optional kBooleanOut{"OUTPUT_BOOL"};
MEDIAPIPE_NODE_CONTRACT(kIn, kOut);
MEDIAPIPE_NODE_CONTRACT(kIn, kOut, kBooleanOut);
absl::Status UpdateContract(CalculatorContract* cc) {
RET_CHECK(kOut(cc).IsConnected() || kBooleanOut(cc).IsConnected())
<< "At least one output stream is expected.";
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) final {
if (!kIn(cc).IsEmpty()) {
auto output = std::make_unique<int>((*kIn(cc) != 0) ? 1 : 0);
kOut(cc).Send(std::move(output));
bool isNonZero = *kIn(cc) != 0;
if (kOut(cc).IsConnected()) {
kOut(cc).Send(std::make_unique<int>(isNonZero ? 1 : 0));
}
if (kBooleanOut(cc).IsConnected()) {
kBooleanOut(cc).Send(std::make_unique<bool>(isNonZero));
}
}
return absl::OkStatus();
}
mediapipe/calculators/core/non_zero_calculator_test.cc
+93
View File
@@ -0,0 +1,93 @@
// 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 "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_runner.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/framework/timestamp.h"
#include "mediapipe/framework/tool/validate_type.h"
namespace mediapipe {
class NonZeroCalculatorTest : public ::testing::Test {
protected:
NonZeroCalculatorTest()
: runner_(
R"pb(
calculator: "NonZeroCalculator"
input_stream: "INPUT:input"
output_stream: "OUTPUT:output"
output_stream: "OUTPUT_BOOL:output_bool"
)pb") {}
void SetInput(const std::vector<int>& inputs) {
int timestamp = 0;
for (const auto input : inputs) {
runner_.MutableInputs()
->Get("INPUT", 0)
.packets.push_back(MakePacket<int>(input).At(Timestamp(timestamp++)));
}
}
std::vector<int> GetOutput() {
std::vector<int> result;
for (const auto output : runner_.Outputs().Get("OUTPUT", 0).packets) {
result.push_back(output.Get<int>());
}
return result;
}
std::vector<bool> GetOutputBool() {
std::vector<bool> result;
for (const auto output : runner_.Outputs().Get("OUTPUT_BOOL", 0).packets) {
result.push_back(output.Get<bool>());
}
return result;
}
CalculatorRunner runner_;
};
TEST_F(NonZeroCalculatorTest, ProducesZeroOutputForZeroInput) {
SetInput({0});
MP_ASSERT_OK(runner_.Run());
EXPECT_THAT(GetOutput(), ::testing::ElementsAre(0));
EXPECT_THAT(GetOutputBool(), ::testing::ElementsAre(false));
}
TEST_F(NonZeroCalculatorTest, ProducesNonZeroOutputForNonZeroInput) {
SetInput({1, 2, 3, -4, 5});
MP_ASSERT_OK(runner_.Run());
EXPECT_THAT(GetOutput(), ::testing::ElementsAre(1, 1, 1, 1, 1));
EXPECT_THAT(GetOutputBool(),
::testing::ElementsAre(true, true, true, true, true));
}
TEST_F(NonZeroCalculatorTest, SwitchesBetweenNonZeroAndZeroOutput) {
SetInput({1, 0, 3, 0, 5});
MP_ASSERT_OK(runner_.Run());
EXPECT_THAT(GetOutput(), ::testing::ElementsAre(1, 0, 1, 0, 1));
EXPECT_THAT(GetOutputBool(),
::testing::ElementsAre(true, false, true, false, true));
}
} // namespace mediapipe
mediapipe/calculators/image/image_cropping_calculator.cc
+1 -1
View File
@@ -285,7 +285,7 @@ absl::Status ImageCroppingCalculator::RenderGpu(CalculatorContext* cc) {
// Run cropping shader on GPU.
{
gpu_helper_.BindFramebuffer(dst_tex); // GL_TEXTURE0
gpu_helper_.BindFramebuffer(dst_tex);
glActiveTexture(GL_TEXTURE1);
glBindTexture(src_tex.target(), src_tex.name());
mediapipe/calculators/image/image_transformation_calculator.cc
+1 -1
View File
@@ -546,7 +546,7 @@ absl::Status ImageTransformationCalculator::RenderGpu(CalculatorContext* cc) {
auto dst = gpu_helper_.CreateDestinationTexture(output_width, output_height,
input.format());
gpu_helper_.BindFramebuffer(dst); // GL_TEXTURE0
gpu_helper_.BindFramebuffer(dst);
glActiveTexture(GL_TEXTURE1);
glBindTexture(src1.target(), src1.name());
mediapipe/calculators/image/recolor_calculator.cc
+7 -1
View File
@@ -209,6 +209,9 @@ absl::Status RecolorCalculator::Close(CalculatorContext* cc) {
absl::Status RecolorCalculator::RenderCpu(CalculatorContext* cc) {
if (cc->Inputs().Tag(kMaskCpuTag).IsEmpty()) {
cc->Outputs()
.Tag(kImageFrameTag)
.AddPacket(cc->Inputs().Tag(kImageFrameTag).Value());
return absl::OkStatus();
}
// Get inputs and setup output.
@@ -270,6 +273,9 @@ absl::Status RecolorCalculator::RenderCpu(CalculatorContext* cc) {
absl::Status RecolorCalculator::RenderGpu(CalculatorContext* cc) {
if (cc->Inputs().Tag(kMaskGpuTag).IsEmpty()) {
cc->Outputs()
.Tag(kGpuBufferTag)
.AddPacket(cc->Inputs().Tag(kGpuBufferTag).Value());
return absl::OkStatus();
}
#if !MEDIAPIPE_DISABLE_GPU
@@ -287,7 +293,7 @@ absl::Status RecolorCalculator::RenderGpu(CalculatorContext* cc) {
// Run recolor shader on GPU.
{
gpu_helper_.BindFramebuffer(dst_tex); // GL_TEXTURE0
gpu_helper_.BindFramebuffer(dst_tex);
glActiveTexture(GL_TEXTURE1);
glBindTexture(img_tex.target(), img_tex.name());
mediapipe/calculators/image/set_alpha_calculator.cc
+2 -2
View File
@@ -323,7 +323,7 @@ absl::Status SetAlphaCalculator::RenderGpu(CalculatorContext* cc) {
const auto& alpha_mask =
cc->Inputs().Tag(kInputAlphaTagGpu).Get<mediapipe::GpuBuffer>();
auto alpha_texture = gpu_helper_.CreateSourceTexture(alpha_mask);
gpu_helper_.BindFramebuffer(output_texture); // GL_TEXTURE0
gpu_helper_.BindFramebuffer(output_texture);
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, input_texture.name());
glActiveTexture(GL_TEXTURE2);
@@ -335,7 +335,7 @@ absl::Status SetAlphaCalculator::RenderGpu(CalculatorContext* cc) {
glBindTexture(GL_TEXTURE_2D, 0);
alpha_texture.Release();
} else {
gpu_helper_.BindFramebuffer(output_texture); // GL_TEXTURE0
gpu_helper_.BindFramebuffer(output_texture);
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, input_texture.name());
GlRender(cc); // use value from options
mediapipe/calculators/tensor/BUILD
+2
View File
@@ -490,6 +490,7 @@ cc_library(
"//mediapipe/framework/port:statusor",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:port",
"//mediapipe/gpu:gpu_origin_cc_proto",
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [":image_to_tensor_calculator_gpu_deps"],
@@ -526,6 +527,7 @@ mediapipe_proto_library(
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
"//mediapipe/gpu:gpu_origin_proto",
],
)
mediapipe/calculators/tensor/image_to_tensor_calculator.cc
+4 -3
View File
@@ -31,6 +31,7 @@
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/gpu/gpu_origin.pb.h"
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/gpu/gpu_buffer.h"
@@ -236,7 +237,7 @@ class ImageToTensorCalculator : public Node {
}
private:
bool DoesInputStartAtBottom() {
bool DoesGpuInputStartAtBottom() {
return options_.gpu_origin() != mediapipe::GpuOrigin_Mode_TOP_LEFT;
}
@@ -290,11 +291,11 @@ class ImageToTensorCalculator : public Node {
#elif MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
ASSIGN_OR_RETURN(gpu_converter_,
CreateImageToGlBufferTensorConverter(
cc, DoesInputStartAtBottom(), GetBorderMode()));
cc, DoesGpuInputStartAtBottom(), GetBorderMode()));
#else
ASSIGN_OR_RETURN(gpu_converter_,
CreateImageToGlTextureTensorConverter(
cc, DoesInputStartAtBottom(), GetBorderMode()));
cc, DoesGpuInputStartAtBottom(), GetBorderMode()));
#endif // MEDIAPIPE_METAL_ENABLED
#endif // !MEDIAPIPE_DISABLE_GPU
}
mediapipe/calculators/tensor/image_to_tensor_calculator.proto
+1 -14
View File
@@ -17,20 +17,7 @@ syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
message GpuOrigin {
enum Mode {
DEFAULT = 0;
// OpenGL: bottom-left origin
// Metal : top-left origin
CONVENTIONAL = 1;
// OpenGL: top-left origin
// Metal : top-left origin
TOP_LEFT = 2;
}
}
import "mediapipe/gpu/gpu_origin.proto";
message ImageToTensorCalculatorOptions {
extend mediapipe.CalculatorOptions {
mediapipe/calculators/tensor/inference_calculator_gl.cc
+2 -1
View File
@@ -317,7 +317,8 @@ absl::Status InferenceCalculatorGlImpl::LoadModel(CalculatorContext* cc) {
absl::Status InferenceCalculatorGlImpl::LoadDelegate(CalculatorContext* cc) {
// Configure and create the delegate.
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
options.compile_options.precision_loss_allowed = 1;
options.compile_options.precision_loss_allowed =
allow_precision_loss_ ? 1 : 0;
options.compile_options.preferred_gl_object_type =
TFLITE_GL_OBJECT_TYPE_FASTEST;
options.compile_options.dynamic_batch_enabled = 0;
mediapipe/calculators/tensor/inference_calculator_metal.cc
+19 -9
View File
@@ -97,6 +97,7 @@ class InferenceCalculatorMetalImpl
Packet<TfLiteModelPtr> model_packet_;
std::unique_ptr<tflite::Interpreter> interpreter_;
TfLiteDelegatePtr delegate_;
bool allow_precision_loss_ = false;
#if MEDIAPIPE_TFLITE_METAL_INFERENCE
MPPMetalHelper* gpu_helper_ = nullptr;
@@ -122,6 +123,9 @@ absl::Status InferenceCalculatorMetalImpl::UpdateContract(
}
absl::Status InferenceCalculatorMetalImpl::Open(CalculatorContext* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
allow_precision_loss_ = options.delegate().gpu().allow_precision_loss();
MP_RETURN_IF_ERROR(LoadModel(cc));
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
@@ -222,7 +226,7 @@ absl::Status InferenceCalculatorMetalImpl::LoadDelegate(CalculatorContext* cc) {
// Configure and create the delegate.
TFLGpuDelegateOptions options;
options.allow_precision_loss = true;
options.allow_precision_loss = allow_precision_loss_;
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypeDoNotWait;
delegate_ =
TfLiteDelegatePtr(TFLGpuDelegateCreate(&options), &TFLGpuDelegateDelete);
@@ -239,7 +243,9 @@ absl::Status InferenceCalculatorMetalImpl::LoadDelegate(CalculatorContext* cc) {
tensor->dims->data + tensor->dims->size};
dims.back() = RoundUp(dims.back(), 4);
gpu_buffers_in_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat16, Tensor::Shape{dims}));
allow_precision_loss_ ? Tensor::ElementType::kFloat16
: Tensor::ElementType::kFloat32,
Tensor::Shape{dims}));
auto buffer_view =
gpu_buffers_in_[i]->GetMtlBufferWriteView(gpu_helper_.mtlDevice);
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
@@ -261,7 +267,9 @@ absl::Status InferenceCalculatorMetalImpl::LoadDelegate(CalculatorContext* cc) {
output_shapes_[i] = {dims};
dims.back() = RoundUp(dims.back(), 4);
gpu_buffers_out_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat16, Tensor::Shape{dims}));
allow_precision_loss_ ? Tensor::ElementType::kFloat16
: Tensor::ElementType::kFloat32,
Tensor::Shape{dims}));
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
delegate_.get(), output_indices[i],
gpu_buffers_out_[i]
@@ -271,17 +279,19 @@ absl::Status InferenceCalculatorMetalImpl::LoadDelegate(CalculatorContext* cc) {
}
// Create converter for GPU input.
converter_to_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
isFloat16:true
convertToPBHWC4:true];
converter_to_BPHWC4_ =
[[TFLBufferConvert alloc] initWithDevice:device
isFloat16:allow_precision_loss_
convertToPBHWC4:true];
if (converter_to_BPHWC4_ == nil) {
return mediapipe::InternalError(
"Error initializating input buffer converter");
}
// Create converter for GPU output.
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
isFloat16:true
convertToPBHWC4:false];
converter_from_BPHWC4_ =
[[TFLBufferConvert alloc] initWithDevice:device
isFloat16:allow_precision_loss_
convertToPBHWC4:false];
if (converter_from_BPHWC4_ == nil) {
return absl::InternalError("Error initializating output buffer converter");
}
mediapipe/calculators/tensor/tensors_to_classification_calculator.cc
+10 -1
View File
@@ -89,7 +89,8 @@ absl::Status TensorsToClassificationCalculator::Open(CalculatorContext* cc) {
ASSIGN_OR_RETURN(string_path,
PathToResourceAsFile(options_.label_map_path()));
std::string label_map_string;
MP_RETURN_IF_ERROR(file::GetContents(string_path, &label_map_string));
MP_RETURN_IF_ERROR(
mediapipe::GetResourceContents(string_path, &label_map_string));
std::istringstream stream(label_map_string);
std::string line;
@@ -98,6 +99,14 @@ absl::Status TensorsToClassificationCalculator::Open(CalculatorContext* cc) {
label_map_[i++] = line;
}
label_map_loaded_ = true;
} else if (options_.has_label_map()) {
for (int i = 0; i < options_.label_map().entries_size(); ++i) {
const auto& entry = options_.label_map().entries(i);
RET_CHECK(!label_map_.contains(entry.id()))
<< "Duplicate id found: " << entry.id();
label_map_[entry.id()] = entry.label();
}
label_map_loaded_ = true;
}
return absl::OkStatus();
mediapipe/calculators/tensor/tensors_to_classification_calculator.proto
+12
View File
@@ -25,6 +25,14 @@ message TensorsToClassificationCalculatorOptions {
optional TensorsToClassificationCalculatorOptions ext = 335742638;
}
message LabelMap {
message Entry {
optional int32 id = 1;
optional string label = 2;
}
repeated Entry entries = 1;
}
// Score threshold for perserving the class.
optional float min_score_threshold = 1;
// Number of highest scoring labels to output. If top_k is not positive then
@@ -32,6 +40,10 @@ message TensorsToClassificationCalculatorOptions {
optional int32 top_k = 2;
// Path to a label map file for getting the actual name of class ids.
optional string label_map_path = 3;
// Label map. (Can be used instead of label_map_path.)
// NOTE: "label_map_path", if specified, takes precedence over "label_map".
optional LabelMap label_map = 5;
// Whether the input is a single float for binary classification.
// When true, only a single float is expected in the input tensor and the
// label map, if provided, is expected to have exactly two labels.
mediapipe/calculators/tensor/tensors_to_classification_calculator_test.cc
+35
View File
@@ -115,6 +115,41 @@ TEST_F(TensorsToClassificationCalculatorTest, CorrectOutputWithLabelMapPath) {
}
}
TEST_F(TensorsToClassificationCalculatorTest, CorrectOutputWithLabelMap) {
mediapipe::CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToClassificationCalculator"
input_stream: "TENSORS:tensors"
output_stream: "CLASSIFICATIONS:classifications"
options {
[mediapipe.TensorsToClassificationCalculatorOptions.ext] {
label_map {
entries { id: 0, label: "ClassA" }
entries { id: 1, label: "ClassB" }
entries { id: 2, label: "ClassC" }
}
}
}
)pb"));
BuildGraph(&runner, {0, 0.5, 1});
MP_ASSERT_OK(runner.Run());
const auto& output_packets_ = runner.Outputs().Tag("CLASSIFICATIONS").packets;
EXPECT_EQ(1, output_packets_.size());
const auto& classification_list =
output_packets_[0].Get<ClassificationList>();
EXPECT_EQ(3, classification_list.classification_size());
// Verify that the label field is set.
for (int i = 0; i < classification_list.classification_size(); ++i) {
EXPECT_EQ(i, classification_list.classification(i).index());
EXPECT_EQ(i * 0.5, classification_list.classification(i).score());
ASSERT_TRUE(classification_list.classification(i).has_label());
}
}
TEST_F(TensorsToClassificationCalculatorTest,
CorrectOutputWithLabelMinScoreThreshold) {
mediapipe::CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
mediapipe/calculators/tensorflow/tensor_to_image_frame_calculator.cc
+54 -26
View File
@@ -34,15 +34,28 @@ constexpr char kTensor[] = "TENSOR";
} // namespace
// Input:
// Tensor of type DT_FLOAT, with values between 0-255 (SRGB or GRAY8). The
// shape can be HxWx{3,1} or simply HxW.
// Tensor of type DT_FLOAT or DT_UINT8, with values between 0-255
// (SRGB or GRAY8). The shape can be HxWx{3,1} or simply HxW.
//
// Optionally supports a scale factor that can scale 0-1 value ranges to 0-255.
// For DT_FLOAT tensors, optionally supports a scale factor that can scale 0-1
// value ranges to 0-255.
//
// Output:
// ImageFrame containing the values of the tensor cast as uint8 (SRGB or GRAY8)
//
// Possible extensions: support other input ranges, maybe 4D tensors.
//
// Example:
// node {
// calculator: "TensorToImageFrameCalculator"
// input_stream: "TENSOR:3d_float_tensor"
// output_stream: "IMAGE:image_frame"
// options {
// [mediapipe.TensorToImageFrameCalculatorOptions.ext] {
// scale_factor: 1.0 # set to 255.0 for [0,1] -> [0,255] scaling
// }
// }
// }
class TensorToImageFrameCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc);
@@ -57,8 +70,8 @@ class TensorToImageFrameCalculator : public CalculatorBase {
REGISTER_CALCULATOR(TensorToImageFrameCalculator);
absl::Status TensorToImageFrameCalculator::GetContract(CalculatorContract* cc) {
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
<< "Only one input stream is supported.";
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1)
<< "Only one output stream is supported.";
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
<< "One input stream must be provided.";
RET_CHECK(cc->Inputs().HasTag(kTensor))
@@ -91,29 +104,44 @@ absl::Status TensorToImageFrameCalculator::Process(CalculatorContext* cc) {
RET_CHECK_EQ(depth, 3) << "Output tensor depth must be 3 or 1.";
}
}
const int32 total_size =
input_tensor.dim_size(0) * input_tensor.dim_size(1) * depth;
std::unique_ptr<uint8[]> buffer(new uint8[total_size]);
auto data = input_tensor.flat<float>().data();
for (int i = 0; i < total_size; ++i) {
float d = scale_factor_ * data[i];
if (d < 0) d = 0;
if (d > 255) d = 255;
buffer[i] = d;
int32 height = input_tensor.dim_size(0);
int32 width = input_tensor.dim_size(1);
auto format = (depth == 3 ? ImageFormat::SRGB : ImageFormat::GRAY8);
const int32 total_size = height * width * depth;
::std::unique_ptr<const ImageFrame> output;
if (input_tensor.dtype() == tensorflow::DT_FLOAT) {
// Allocate buffer with alignments.
std::unique_ptr<uint8_t[]> buffer(
new (std::align_val_t(EIGEN_MAX_ALIGN_BYTES)) uint8_t[total_size]);
auto data = input_tensor.flat<float>().data();
for (int i = 0; i < total_size; ++i) {
float d = scale_factor_ * data[i];
if (d < 0) d = 0;
if (d > 255) d = 255;
buffer[i] = d;
}
output = ::absl::make_unique<ImageFrame>(format, width, height,
width * depth, buffer.release());
} else if (input_tensor.dtype() == tensorflow::DT_UINT8) {
if (scale_factor_ != 1.0) {
return absl::InvalidArgumentError("scale_factor_ given for uint8 tensor");
}
// tf::Tensor has internally ref-counted buffer. The following code make the
// ImageFrame own the copied Tensor through the deleter, which increases
// the refcount of the buffer and allow us to use the shared buffer as the
// image. This allows us to create an ImageFrame object without copying
// buffer. const ImageFrame prevents the buffer from being modified later.
auto copy = new tf::Tensor(input_tensor);
output = ::absl::make_unique<const ImageFrame>(
format, width, height, width * depth, copy->flat<uint8_t>().data(),
[copy](uint8*) { delete copy; });
} else {
return absl::InvalidArgumentError(
absl::StrCat("Expected float or uint8 tensor, received ",
DataTypeString(input_tensor.dtype())));
}
::std::unique_ptr<ImageFrame> output;
if (depth == 3) {
output = ::absl::make_unique<ImageFrame>(
ImageFormat::SRGB, input_tensor.dim_size(1), input_tensor.dim_size(0),
input_tensor.dim_size(1) * 3, buffer.release());
} else if (depth == 1) {
output = ::absl::make_unique<ImageFrame>(
ImageFormat::GRAY8, input_tensor.dim_size(1), input_tensor.dim_size(0),
input_tensor.dim_size(1), buffer.release());
} else {
return absl::InvalidArgumentError("Unrecognized image depth.");
}
cc->Outputs().Tag(kImage).Add(output.release(), cc->InputTimestamp());
return absl::OkStatus();
mediapipe/calculators/tensorflow/tensor_to_image_frame_calculator_test.cc
+40 -26
View File
@@ -29,6 +29,7 @@ constexpr char kImage[] = "IMAGE";
} // namespace
template <class TypeParam>
class TensorToImageFrameCalculatorTest : public ::testing::Test {
protected:
void SetUpRunner() {
@@ -42,14 +43,20 @@ class TensorToImageFrameCalculatorTest : public ::testing::Test {
std::unique_ptr<CalculatorRunner> runner_;
};
TEST_F(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrame) {
SetUpRunner();
using TensorToImageFrameCalculatorTestTypes = ::testing::Types<float, uint8_t>;
TYPED_TEST_CASE(TensorToImageFrameCalculatorTest,
TensorToImageFrameCalculatorTestTypes);
TYPED_TEST(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrame) {
// TYPED_TEST requires explicit "this->"
this->SetUpRunner();
auto& runner = this->runner_;
constexpr int kWidth = 16;
constexpr int kHeight = 8;
const tf::TensorShape tensor_shape(
std::vector<tf::int64>{kHeight, kWidth, 3});
auto tensor = absl::make_unique<tf::Tensor>(tf::DT_FLOAT, tensor_shape);
auto tensor_vec = tensor->flat<float>().data();
const tf::TensorShape tensor_shape{kHeight, kWidth, 3};
auto tensor = absl::make_unique<tf::Tensor>(
tf::DataTypeToEnum<TypeParam>::v(), tensor_shape);
auto tensor_vec = tensor->template flat<TypeParam>().data();
// Writing sequence of integers as floats which we want back (as they were
// written).
@@ -58,15 +65,16 @@ TEST_F(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrame) {
}
const int64 time = 1234;
runner_->MutableInputs()->Tag(kTensor).packets.push_back(
runner->MutableInputs()->Tag(kTensor).packets.push_back(
Adopt(tensor.release()).At(Timestamp(time)));
EXPECT_TRUE(runner_->Run().ok());
EXPECT_TRUE(runner->Run().ok());
const std::vector<Packet>& output_packets =
runner_->Outputs().Tag(kImage).packets;
runner->Outputs().Tag(kImage).packets;
EXPECT_EQ(1, output_packets.size());
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
const ImageFrame& output_image = output_packets[0].Get<ImageFrame>();
EXPECT_EQ(ImageFormat::SRGB, output_image.Format());
EXPECT_EQ(kWidth, output_image.Width());
EXPECT_EQ(kHeight, output_image.Height());
@@ -76,14 +84,15 @@ TEST_F(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrame) {
}
}
TEST_F(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrameGray) {
SetUpRunner();
TYPED_TEST(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrameGray) {
this->SetUpRunner();
auto& runner = this->runner_;
constexpr int kWidth = 16;
constexpr int kHeight = 8;
const tf::TensorShape tensor_shape(
std::vector<tf::int64>{kHeight, kWidth, 1});
auto tensor = absl::make_unique<tf::Tensor>(tf::DT_FLOAT, tensor_shape);
auto tensor_vec = tensor->flat<float>().data();
const tf::TensorShape tensor_shape{kHeight, kWidth, 1};
auto tensor = absl::make_unique<tf::Tensor>(
tf::DataTypeToEnum<TypeParam>::v(), tensor_shape);
auto tensor_vec = tensor->template flat<TypeParam>().data();
// Writing sequence of integers as floats which we want back (as they were
// written).
@@ -92,15 +101,16 @@ TEST_F(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrameGray) {
}
const int64 time = 1234;
runner_->MutableInputs()->Tag(kTensor).packets.push_back(
runner->MutableInputs()->Tag(kTensor).packets.push_back(
Adopt(tensor.release()).At(Timestamp(time)));
EXPECT_TRUE(runner_->Run().ok());
EXPECT_TRUE(runner->Run().ok());
const std::vector<Packet>& output_packets =
runner_->Outputs().Tag(kImage).packets;
runner->Outputs().Tag(kImage).packets;
EXPECT_EQ(1, output_packets.size());
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
const ImageFrame& output_image = output_packets[0].Get<ImageFrame>();
EXPECT_EQ(ImageFormat::GRAY8, output_image.Format());
EXPECT_EQ(kWidth, output_image.Width());
EXPECT_EQ(kHeight, output_image.Height());
@@ -110,13 +120,16 @@ TEST_F(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrameGray) {
}
}
TEST_F(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrame2DGray) {
SetUpRunner();
TYPED_TEST(TensorToImageFrameCalculatorTest,
Converts3DTensorToImageFrame2DGray) {
this->SetUpRunner();
auto& runner = this->runner_;
constexpr int kWidth = 16;
constexpr int kHeight = 8;
const tf::TensorShape tensor_shape(std::vector<tf::int64>{kHeight, kWidth});
auto tensor = absl::make_unique<tf::Tensor>(tf::DT_FLOAT, tensor_shape);
auto tensor_vec = tensor->flat<float>().data();
const tf::TensorShape tensor_shape{kHeight, kWidth};
auto tensor = absl::make_unique<tf::Tensor>(
tf::DataTypeToEnum<TypeParam>::v(), tensor_shape);
auto tensor_vec = tensor->template flat<TypeParam>().data();
// Writing sequence of integers as floats which we want back (as they were
// written).
@@ -125,15 +138,16 @@ TEST_F(TensorToImageFrameCalculatorTest, Converts3DTensorToImageFrame2DGray) {
}
const int64 time = 1234;
runner_->MutableInputs()->Tag(kTensor).packets.push_back(
runner->MutableInputs()->Tag(kTensor).packets.push_back(
Adopt(tensor.release()).At(Timestamp(time)));
EXPECT_TRUE(runner_->Run().ok());
EXPECT_TRUE(runner->Run().ok());
const std::vector<Packet>& output_packets =
runner_->Outputs().Tag(kImage).packets;
runner->Outputs().Tag(kImage).packets;
EXPECT_EQ(1, output_packets.size());
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
const ImageFrame& output_image = output_packets[0].Get<ImageFrame>();
EXPECT_EQ(ImageFormat::GRAY8, output_image.Format());
EXPECT_EQ(kWidth, output_image.Width());
EXPECT_EQ(kHeight, output_image.Height());
mediapipe/calculators/tensorflow/tensor_to_matrix_calculator.cc
+11 -23
View File
@@ -91,8 +91,6 @@ absl::Status FillTimeSeriesHeaderIfValid(const Packet& header_packet,
// the input data when it arrives in Process(). In particular, if the header
// states that we produce a 1xD column vector, the input tensor must also be 1xD
//
// This designed was discussed in http://g/speakeranalysis/4uyx7cNRwJY and
// http://g/daredevil-project/VB26tcseUy8.
// Example Config
// node: {
// calculator: "TensorToMatrixCalculator"
@@ -158,22 +156,17 @@ absl::Status TensorToMatrixCalculator::Open(CalculatorContext* cc) {
if (header_status.ok()) {
if (cc->Options<TensorToMatrixCalculatorOptions>()
.has_time_series_header_overrides()) {
// From design discussions with Daredevil, we only want to support single
// sample per packet for now, so we hardcode the sample_rate based on the
// packet_rate of the REFERENCE and fail noisily if we cannot. An
// alternative would be to calculate the sample_rate from the reference
// sample_rate and the change in num_samples between the reference and
// override headers:
// sample_rate_output = sample_rate_reference /
// (num_samples_override / num_samples_reference)
// This only supports a single sample per packet for now, so we hardcode
// the sample_rate based on the packet_rate of the REFERENCE and fail
// if we cannot.
const TimeSeriesHeader& override_header =
cc->Options<TensorToMatrixCalculatorOptions>()
.time_series_header_overrides();
input_header->MergeFrom(override_header);
CHECK(input_header->has_packet_rate())
RET_CHECK(input_header->has_packet_rate())
<< "The TimeSeriesHeader.packet_rate must be set.";
if (!override_header.has_sample_rate()) {
CHECK_EQ(input_header->num_samples(), 1)
RET_CHECK_EQ(input_header->num_samples(), 1)
<< "Currently the time series can only output single samples.";
input_header->set_sample_rate(input_header->packet_rate());
}
@@ -186,20 +179,16 @@ absl::Status TensorToMatrixCalculator::Open(CalculatorContext* cc) {
}
absl::Status TensorToMatrixCalculator::Process(CalculatorContext* cc) {
// Daredevil requested CHECK for noisy failures rather than quieter RET_CHECK
// failures. These are absolute conditions of the graph for the graph to be
// valid, and if it is violated by any input anywhere, the graph will be
// invalid for all inputs. A hard CHECK will enable faster debugging by
// immediately exiting and more prominently displaying error messages.
// Do not replace with RET_CHECKs.
// Verify that each reference stream packet corresponds to a tensor packet
// otherwise the header information is invalid. If we don't have a reference
// stream, Process() is only called when we have an input tensor and this is
// always True.
CHECK(cc->Inputs().HasTag(kTensor))
RET_CHECK(cc->Inputs().HasTag(kTensor))
<< "Tensor stream not available at same timestamp as the reference "
"stream.";
RET_CHECK(!cc->Inputs().Tag(kTensor).IsEmpty()) << "Tensor stream is empty.";
RET_CHECK_OK(cc->Inputs().Tag(kTensor).Value().ValidateAsType<tf::Tensor>())
<< "Tensor stream packet does not contain a Tensor.";
const tf::Tensor& input_tensor = cc->Inputs().Tag(kTensor).Get<tf::Tensor>();
CHECK(1 == input_tensor.dims() || 2 == input_tensor.dims())
@@ -207,13 +196,12 @@ absl::Status TensorToMatrixCalculator::Process(CalculatorContext* cc) {
const int32 length = input_tensor.dim_size(input_tensor.dims() - 1);
const int32 width = (1 == input_tensor.dims()) ? 1 : input_tensor.dim_size(0);
if (header_.has_num_channels()) {
CHECK_EQ(length, header_.num_channels())
RET_CHECK_EQ(length, header_.num_channels())
<< "The number of channels at runtime does not match the header.";
}
if (header_.has_num_samples()) {
CHECK_EQ(width, header_.num_samples())
RET_CHECK_EQ(width, header_.num_samples())
<< "The number of samples at runtime does not match the header.";
;
}
auto output = absl::make_unique<Matrix>(width, length);
*output =
mediapipe/calculators/tensorflow/tensorflow_inference_calculator.cc
+511 -356
View File
@@ -98,388 +98,543 @@ class InferenceState {
// This calculator performs inference on a trained TensorFlow model.
//
// A mediapipe::TensorFlowSession with a model loaded and ready for use.
// For this calculator it must include a tag_to_tensor_map.
cc->InputSidePackets().Tag("SESSION").Set<TensorFlowSession>();
if (cc->InputSidePackets().HasTag("RECURRENT_INIT_TENSORS")) {
cc->InputSidePackets()
.Tag("RECURRENT_INIT_TENSORS")
.Set<std::unique_ptr<std::map<std::string, tf::Tensor>>>();
}
return absl::OkStatus();
}
// TensorFlow Sessions can be created from checkpoint paths, frozen models, or
// the SavedModel system. See the TensorFlowSessionFrom* packet generators for
// details. Each of these methods defines a mapping between MediaPipe streams
// and TensorFlow tensors. All of this information is passed in as an
// input_side_packet.
//
// The input and output streams are TensorFlow tensors labeled by tags. The tags
// for the streams are matched to feeds and fetchs in a TensorFlow session using
// a named_signature.generic_signature in the ModelManifest. The
// generic_signature is used as key-value pairs between the MediaPipe tag and
// the TensorFlow tensor. The signature_name in the options proto determines
// which named_signature is used. The keys in the generic_signature must be
// valid MediaPipe tags ([A-Z0-9_]*, no lowercase or special characters). All of
// the tensors corresponding to tags in the signature for input_streams are fed
// to the model and for output_streams the tensors are fetched from the model.
//
// Other calculators are used to convert data to and from tensors, this op only
// handles the TensorFlow session and batching. Batching occurs by concatenating
// input tensors along the 0th dimension across timestamps. If the 0th dimension
// is not a batch dimension, this calculator will add a 0th dimension by
// default. Setting add_batch_dim_to_tensors to false disables the dimension
// addition. Once batch_size inputs have been provided, the batch will be run
// and the output tensors sent out on the output streams with timestamps
// corresponding to the input stream packets. Setting the batch_size to 1
// completely disables batching, but is indepdent of add_batch_dim_to_tensors.
//
// The TensorFlowInferenceCalculator also support feeding states recurrently for
// RNNs and LSTMs. Simply set the recurrent_tag_pair options to define the
// recurrent tensors. Initializing the recurrent state can be handled by the
// GraphTensorsPacketGenerator.
//
// The calculator updates two Counters to report timing information:
// --<name>-TotalTimeUsecs = Total time spent running inference (in usecs),
// --<name>-TotalProcessedTimestamps = # of instances processed
// (approximately batches processed * batch_size),
// where <name> is replaced with CalculatorGraphConfig::Node::name() if it
// exists, or with TensorFlowInferenceCalculator if the name is not set. The
// name must be set for timing information to be instance-specific in graphs
// with multiple TensorFlowInferenceCalculators.
//
// Example config:
// packet_generator {
// packet_generator: "TensorFlowSessionFromSavedModelGenerator"
// output_side_packet: "tensorflow_session"
// options {
// [mediapipe.TensorFlowSessionFromSavedModelGeneratorOptions.ext]: {
// saved_model_path: "/path/to/saved/model"
// signature_name: "mediapipe"
// }
// }
// }
// node {
// calculator: "TensorFlowInferenceCalculator"
// input_stream: "IMAGES:image_tensors_keyed_in_signature_by_tag"
// input_stream: "AUDIO:audio_tensors_keyed_in_signature_by_tag"
// output_stream: "LABELS:softmax_tensor_keyed_in_signature_by_tag"
// input_side_packet: "SESSION:tensorflow_session"
// }
//
// Where the input and output streams are treated as Packet<tf::Tensor> and
// the mediapipe_signature has tensor bindings between "IMAGES", "AUDIO", and
// "LABELS" and their respective tensors exported to /path/to/bundle. For an
// example of how this model was exported, see
// tensorflow_inference_test_graph_generator.py
//
// It is possible to use a GraphDef proto that was not exported by exporter (i.e
// without MetaGraph with bindings). Such GraphDef could contain all of its
// parameters in-lined (for example, it can be the output of freeze_graph.py).
// To instantiate a TensorFlow model from a GraphDef file, replace the
// packet_factory above with TensorFlowSessionFromFrozenGraphGenerator:
//
// packet_generator {
// packet_generator: "TensorFlowSessionFromFrozenGraphGenerator"
// output_side_packet: "SESSION:tensorflow_session"
// options {
// [mediapipe.TensorFlowSessionFromFrozenGraphGeneratorOptions.ext]: {
// graph_proto_path: "[PATH]"
// tag_to_tensor_names {
// key: "JPG_STRING"
// value: "input:0"
// }
// tag_to_tensor_names {
// key: "SOFTMAX"
// value: "softmax:0"
// }
// }
// }
// }
//
// It is also possible to use a GraphDef proto and checkpoint file that have not
// been frozen. This can be used to load graphs directly as they have been
// written from training. However, it is more brittle and you are encouraged to
// use a one of the more perminent formats described above. To instantiate a
// TensorFlow model from a GraphDef file and checkpoint, replace the
// packet_factory above with TensorFlowSessionFromModelCheckpointGenerator:
//
// packet_generator {
// packet_generator: "TensorFlowSessionFromModelCheckpointGenerator"
// output_side_packet: "SESSION:tensorflow_session"
// options {
// [mediapipe.TensorFlowSessionFromModelCheckpointGeneratorOptions.ext]: {
// graph_proto_path: "[PATH]"
// model_options {
// checkpoint_path: "[PATH2]"
// }
// tag_to_tensor_names {
// key: "JPG_STRING"
// value: "input:0"
// }
// tag_to_tensor_names {
// key: "SOFTMAX"
// value: "softmax:0"
// }
// }
// }
// }
class TensorFlowInferenceCalculator : public CalculatorBase {
public:
// Counters for recording timing information. The actual names have the value
// of CalculatorGraphConfig::Node::name() prepended.
static constexpr char kTotalUsecsCounterSuffix[] = "TotalTimeUsecs";
static constexpr char kTotalProcessedTimestampsCounterSuffix[] =
"TotalProcessedTimestamps";
static constexpr char kTotalSessionRunsTimeUsecsCounterSuffix[] =
"TotalSessionRunsTimeUsecs";
static constexpr char kTotalNumSessionRunsCounterSuffix[] =
"TotalNumSessionRuns";
std::unique_ptr<InferenceState> CreateInferenceState(CalculatorContext* cc)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
std::unique_ptr<InferenceState> inference_state =
absl::make_unique<InferenceState>();
if (cc->InputSidePackets().HasTag("RECURRENT_INIT_TENSORS") &&
!cc->InputSidePackets().Tag("RECURRENT_INIT_TENSORS").IsEmpty()) {
std::map<std::string, tf::Tensor>* init_tensor_map;
init_tensor_map = GetFromUniquePtr<std::map<std::string, tf::Tensor>>(
cc->InputSidePackets().Tag("RECURRENT_INIT_TENSORS"));
for (const auto& p : *init_tensor_map) {
inference_state->input_tensor_batches_[p.first].emplace_back(p.second);
TensorFlowInferenceCalculator() : session_(nullptr) {
clock_ = std::unique_ptr<mediapipe::Clock>(
mediapipe::MonotonicClock::CreateSynchronizedMonotonicClock());
}
static absl::Status GetContract(CalculatorContract* cc) {
const auto& options = cc->Options<TensorFlowInferenceCalculatorOptions>();
RET_CHECK(!cc->Inputs().GetTags().empty());
for (const std::string& tag : cc->Inputs().GetTags()) {
// The tensorflow::Tensor with the tag equal to the graph node. May
// have a TimeSeriesHeader if all present TimeSeriesHeaders match.
if (!options.batched_input()) {
cc->Inputs().Tag(tag).Set<tf::Tensor>();
} else {
cc->Inputs().Tag(tag).Set<std::vector<mediapipe::Packet>>();
}
}
}
return inference_state;
}
absl::Status Open(CalculatorContext* cc) override {
options_ = cc->Options<TensorFlowInferenceCalculatorOptions>();
RET_CHECK(cc->InputSidePackets().HasTag("SESSION"));
session_ = cc->InputSidePackets()
.Tag("SESSION")
.Get<TensorFlowSession>()
.session.get();
tag_to_tensor_map_ = cc->InputSidePackets()
.Tag("SESSION")
.Get<TensorFlowSession>()
.tag_to_tensor_map;
// Validate and store the recurrent tags
RET_CHECK(options_.has_batch_size());
RET_CHECK(options_.batch_size() == 1 || options_.recurrent_tag_pair().empty())
<< "To use recurrent_tag_pairs, batch_size must be 1.";
for (const auto& tag_pair : options_.recurrent_tag_pair()) {
const std::vector<std::string> tags = absl::StrSplit(tag_pair, ':');
RET_CHECK_EQ(tags.size(), 2)
<< "recurrent_tag_pair must be a colon "
"separated std::string with two components: "
<< tag_pair;
RET_CHECK(mediapipe::ContainsKey(tag_to_tensor_map_, tags[0]))
<< "Can't find tag '" << tags[0] << "' in signature "
<< options_.signature_name();
RET_CHECK(mediapipe::ContainsKey(tag_to_tensor_map_, tags[1]))
<< "Can't find tag '" << tags[1] << "' in signature "
<< options_.signature_name();
recurrent_feed_tags_.insert(tags[0]);
recurrent_fetch_tags_to_feed_tags_[tags[1]] = tags[0];
}
// Check that all tags are present in this signature bound to tensors.
for (const std::string& tag : cc->Inputs().GetTags()) {
RET_CHECK(mediapipe::ContainsKey(tag_to_tensor_map_, tag))
<< "Can't find tag '" << tag << "' in signature "
<< options_.signature_name();
}
for (const std::string& tag : cc->Outputs().GetTags()) {
RET_CHECK(mediapipe::ContainsKey(tag_to_tensor_map_, tag))
<< "Can't find tag '" << tag << "' in signature "
<< options_.signature_name();
}
{
absl::WriterMutexLock l(&mutex_);
inference_state_ = std::unique_ptr<InferenceState>();
}
if (options_.batch_size() == 1 || options_.batched_input()) {
cc->SetOffset(0);
}
return absl::OkStatus();
}
// Adds a batch dimension to the input tensor if specified in the calculator
// options.
absl::Status AddBatchDimension(tf::Tensor* input_tensor) {
if (options_.add_batch_dim_to_tensors()) {
tf::TensorShape new_shape(input_tensor->shape());
new_shape.InsertDim(0, 1);
RET_CHECK(input_tensor->CopyFrom(*input_tensor, new_shape))
<< "Could not add 0th dimension to tensor without changing its shape."
<< " Current shape: " << input_tensor->shape().DebugString();
}
return absl::OkStatus();
}
absl::Status AggregateTensorPacket(
const std::string& tag_name, const Packet& packet,
std::map<Timestamp, std::map<std::string, tf::Tensor>>*
input_tensors_by_tag_by_timestamp,
InferenceState* inference_state) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
tf::Tensor input_tensor(packet.Get<tf::Tensor>());
RET_CHECK_OK(AddBatchDimension(&input_tensor));
if (mediapipe::ContainsKey(recurrent_feed_tags_, tag_name)) {
// If we receive an input on a recurrent tag, override the state.
// It's OK to override the global state because there is just one
// input stream allowed for recurrent tensors.
inference_state_->input_tensor_batches_[tag_name].clear();
}
(*input_tensors_by_tag_by_timestamp)[packet.Timestamp()].insert(
std::make_pair(tag_name, input_tensor));
return absl::OkStatus();
}
// Removes the batch dimension of the output tensor if specified in the
// calculator options.
absl::Status RemoveBatchDimension(tf::Tensor* output_tensor) {
if (options_.add_batch_dim_to_tensors()) {
tf::TensorShape new_shape(output_tensor->shape());
new_shape.RemoveDim(0);
RET_CHECK(output_tensor->CopyFrom(*output_tensor, new_shape))
<< "Could not remove 0th dimension from tensor without changing its "
<< "shape. Current shape: " << output_tensor->shape().DebugString()
<< " (The expected first dimension is 1 for a batch element.)";
}
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) override {
std::unique_ptr<InferenceState> inference_state_to_process;
{
absl::WriterMutexLock l(&mutex_);
if (inference_state_ == nullptr) {
inference_state_ = CreateInferenceState(cc);
RET_CHECK(!cc->Outputs().GetTags().empty());
for (const std::string& tag : cc->Outputs().GetTags()) {
// The tensorflow::Tensor with tag equal to the graph node to
// output. Any TimeSeriesHeader from the inputs will be forwarded
// with channels set to 0.
cc->Outputs().Tag(tag).Set<tf::Tensor>();
}
std::map<Timestamp, std::map<std::string, tf::Tensor>>
input_tensors_by_tag_by_timestamp;
for (const std::string& tag_as_node_name : cc->Inputs().GetTags()) {
if (cc->Inputs().Tag(tag_as_node_name).IsEmpty()) {
// Recurrent tensors can be empty.
if (!mediapipe::ContainsKey(recurrent_feed_tags_, tag_as_node_name)) {
if (options_.skip_on_missing_features()) {
return absl::OkStatus();
} else {
return absl::InvalidArgumentError(absl::StrCat(
"Tag ", tag_as_node_name,
" not present at timestamp: ", cc->InputTimestamp().Value()));
// A mediapipe::TensorFlowSession with a model loaded and ready for use.
// For this calculator it must include a tag_to_tensor_map.
cc->InputSidePackets().Tag("SESSION").Set<TensorFlowSession>();
if (cc->InputSidePackets().HasTag("RECURRENT_INIT_TENSORS")) {
cc->InputSidePackets()
.Tag("RECURRENT_INIT_TENSORS")
.Set<std::unique_ptr<std::map<std::string, tf::Tensor>>>();
}
return absl::OkStatus();
}
std::unique_ptr<InferenceState> CreateInferenceState(CalculatorContext* cc)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
std::unique_ptr<InferenceState> inference_state =
absl::make_unique<InferenceState>();
if (cc->InputSidePackets().HasTag("RECURRENT_INIT_TENSORS") &&
!cc->InputSidePackets().Tag("RECURRENT_INIT_TENSORS").IsEmpty()) {
std::map<std::string, tf::Tensor>* init_tensor_map;
init_tensor_map = GetFromUniquePtr<std::map<std::string, tf::Tensor>>(
cc->InputSidePackets().Tag("RECURRENT_INIT_TENSORS"));
for (const auto& p : *init_tensor_map) {
inference_state->input_tensor_batches_[p.first].emplace_back(p.second);
}
}
return inference_state;
}
absl::Status Open(CalculatorContext* cc) override {
options_ = cc->Options<TensorFlowInferenceCalculatorOptions>();
RET_CHECK(cc->InputSidePackets().HasTag("SESSION"));
session_ = cc->InputSidePackets()
.Tag("SESSION")
.Get<TensorFlowSession>()
.session.get();
tag_to_tensor_map_ = cc->InputSidePackets()
.Tag("SESSION")
.Get<TensorFlowSession>()
.tag_to_tensor_map;
// Validate and store the recurrent tags
RET_CHECK(options_.has_batch_size());
RET_CHECK(options_.batch_size() == 1 ||
options_.recurrent_tag_pair().empty())
<< "To use recurrent_tag_pairs, batch_size must be 1.";
for (const auto& tag_pair : options_.recurrent_tag_pair()) {
const std::vector<std::string> tags = absl::StrSplit(tag_pair, ':');
RET_CHECK_EQ(tags.size(), 2)
<< "recurrent_tag_pair must be a colon "
"separated std::string with two components: "
<< tag_pair;
RET_CHECK(mediapipe::ContainsKey(tag_to_tensor_map_, tags[0]))
<< "Can't find tag '" << tags[0] << "' in signature "
<< options_.signature_name();
RET_CHECK(mediapipe::ContainsKey(tag_to_tensor_map_, tags[1]))
<< "Can't find tag '" << tags[1] << "' in signature "
<< options_.signature_name();
recurrent_feed_tags_.insert(tags[0]);
recurrent_fetch_tags_to_feed_tags_[tags[1]] = tags[0];
}
// Check that all tags are present in this signature bound to tensors.
for (const std::string& tag : cc->Inputs().GetTags()) {
RET_CHECK(mediapipe::ContainsKey(tag_to_tensor_map_, tag))
<< "Can't find tag '" << tag << "' in signature "
<< options_.signature_name();
}
for (const std::string& tag : cc->Outputs().GetTags()) {
RET_CHECK(mediapipe::ContainsKey(tag_to_tensor_map_, tag))
<< "Can't find tag '" << tag << "' in signature "
<< options_.signature_name();
}
{
absl::WriterMutexLock l(&mutex_);
inference_state_ = std::unique_ptr<InferenceState>();
}
if (options_.batch_size() == 1 || options_.batched_input()) {
cc->SetOffset(0);
}
return absl::OkStatus();
}
// Adds a batch dimension to the input tensor if specified in the calculator
// options.
absl::Status AddBatchDimension(tf::Tensor* input_tensor) {
if (options_.add_batch_dim_to_tensors()) {
tf::TensorShape new_shape(input_tensor->shape());
new_shape.InsertDim(0, 1);
RET_CHECK(input_tensor->CopyFrom(*input_tensor, new_shape))
<< "Could not add 0th dimension to tensor without changing its shape."
<< " Current shape: " << input_tensor->shape().DebugString();
}
return absl::OkStatus();
}
absl::Status AggregateTensorPacket(
const std::string& tag_name, const Packet& packet,
std::map<Timestamp, std::map<std::string, tf::Tensor>>*
input_tensors_by_tag_by_timestamp,
InferenceState* inference_state) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
tf::Tensor input_tensor(packet.Get<tf::Tensor>());
RET_CHECK_OK(AddBatchDimension(&input_tensor));
if (mediapipe::ContainsKey(recurrent_feed_tags_, tag_name)) {
// If we receive an input on a recurrent tag, override the state.
// It's OK to override the global state because there is just one
// input stream allowed for recurrent tensors.
inference_state_->input_tensor_batches_[tag_name].clear();
}
(*input_tensors_by_tag_by_timestamp)[packet.Timestamp()].insert(
std::make_pair(tag_name, input_tensor));
return absl::OkStatus();
}
// Removes the batch dimension of the output tensor if specified in the
// calculator options.
absl::Status RemoveBatchDimension(tf::Tensor* output_tensor) {
if (options_.add_batch_dim_to_tensors()) {
tf::TensorShape new_shape(output_tensor->shape());
new_shape.RemoveDim(0);
RET_CHECK(output_tensor->CopyFrom(*output_tensor, new_shape))
<< "Could not remove 0th dimension from tensor without changing its "
<< "shape. Current shape: " << output_tensor->shape().DebugString()
<< " (The expected first dimension is 1 for a batch element.)";
}
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) override {
std::unique_ptr<InferenceState> inference_state_to_process;
{
absl::WriterMutexLock l(&mutex_);
if (inference_state_ == nullptr) {
inference_state_ = CreateInferenceState(cc);
}
std::map<Timestamp, std::map<std::string, tf::Tensor>>
input_tensors_by_tag_by_timestamp;
for (const std::string& tag_as_node_name : cc->Inputs().GetTags()) {
if (cc->Inputs().Tag(tag_as_node_name).IsEmpty()) {
// Recurrent tensors can be empty.
if (!mediapipe::ContainsKey(recurrent_feed_tags_, tag_as_node_name)) {
if (options_.skip_on_missing_features()) {
return absl::OkStatus();
} else {
return absl::InvalidArgumentError(absl::StrCat(
"Tag ", tag_as_node_name,
" not present at timestamp: ", cc->InputTimestamp().Value()));
}
}
} else if (options_.batched_input()) {
const auto& tensor_packets =
cc->Inputs().Tag(tag_as_node_name).Get<std::vector<Packet>>();
if (tensor_packets.size() > options_.batch_size()) {
return absl::InvalidArgumentError(absl::StrCat(
"Batch for tag ", tag_as_node_name,
" has more packets than batch capacity. batch_size: ",
options_.batch_size(), " packets: ", tensor_packets.size()));
}
for (const auto& packet : tensor_packets) {
RET_CHECK_OK(AggregateTensorPacket(
tag_as_node_name, packet, &input_tensors_by_tag_by_timestamp,
inference_state_.get()));
}
} else {
RET_CHECK_OK(AggregateTensorPacket(
tag_as_node_name, cc->Inputs().Tag(tag_as_node_name).Value(),
&input_tensors_by_tag_by_timestamp, inference_state_.get()));
}
} else if (options_.batched_input()) {
const auto& tensor_packets =
cc->Inputs().Tag(tag_as_node_name).Get<std::vector<Packet>>();
if (tensor_packets.size() > options_.batch_size()) {
return absl::InvalidArgumentError(absl::StrCat(
"Batch for tag ", tag_as_node_name,
" has more packets than batch capacity. batch_size: ",
options_.batch_size(), " packets: ", tensor_packets.size()));
}
for (const auto& timestamp_and_input_tensors_by_tag :
input_tensors_by_tag_by_timestamp) {
inference_state_->batch_timestamps_.emplace_back(
timestamp_and_input_tensors_by_tag.first);
for (const auto& input_tensor_and_tag :
timestamp_and_input_tensors_by_tag.second) {
inference_state_->input_tensor_batches_[input_tensor_and_tag.first]
.emplace_back(input_tensor_and_tag.second);
}
for (const auto& packet : tensor_packets) {
RET_CHECK_OK(AggregateTensorPacket(tag_as_node_name, packet,
&input_tensors_by_tag_by_timestamp,
inference_state_.get()));
}
if (inference_state_->batch_timestamps_.size() == options_.batch_size() ||
options_.batched_input()) {
inference_state_to_process = std::move(inference_state_);
inference_state_ = std::unique_ptr<InferenceState>();
}
}
if (inference_state_to_process) {
MP_RETURN_IF_ERROR(
OutputBatch(cc, std::move(inference_state_to_process)));
}
return absl::OkStatus();
}
absl::Status Close(CalculatorContext* cc) override {
std::unique_ptr<InferenceState> inference_state_to_process = nullptr;
{
absl::WriterMutexLock l(&mutex_);
if (cc->GraphStatus().ok() && inference_state_ != nullptr &&
!inference_state_->batch_timestamps_.empty()) {
inference_state_to_process = std::move(inference_state_);
inference_state_ = std::unique_ptr<InferenceState>();
}
}
if (inference_state_to_process) {
MP_RETURN_IF_ERROR(
OutputBatch(cc, std::move(inference_state_to_process)));
}
return absl::OkStatus();
}
// When a batch of input tensors is ready to be run, runs TensorFlow and
// outputs the output tensors. The output tensors have timestamps matching
// the input tensor that formed that batch element. Any requested
// batch_dimension is added and removed. This code takes advantage of the fact
// that copying a tensor shares the same reference-counted, heap allocated
// memory buffer. Therefore, copies are cheap and should not cause the memory
// buffer to fall out of scope. In contrast, concat is only used where
// necessary.
absl::Status OutputBatch(CalculatorContext* cc,
std::unique_ptr<InferenceState> inference_state) {
const int64 start_time = absl::ToUnixMicros(clock_->TimeNow());
std::vector<std::pair<mediapipe::ProtoString, tf::Tensor>> input_tensors;
for (auto& keyed_tensors : inference_state->input_tensor_batches_) {
if (options_.batch_size() == 1) {
// Short circuit to avoid the cost of deep copying tensors in concat.
if (!keyed_tensors.second.empty()) {
input_tensors.emplace_back(tag_to_tensor_map_[keyed_tensors.first],
keyed_tensors.second[0]);
} else {
// The input buffer can be empty for recurrent tensors.
RET_CHECK(
mediapipe::ContainsKey(recurrent_feed_tags_, keyed_tensors.first))
<< "A non-recurrent tensor does not have an input: "
<< keyed_tensors.first;
}
} else {
RET_CHECK_OK(AggregateTensorPacket(
tag_as_node_name, cc->Inputs().Tag(tag_as_node_name).Value(),
&input_tensors_by_tag_by_timestamp, inference_state_.get()));
}
}
for (const auto& timestamp_and_input_tensors_by_tag :
input_tensors_by_tag_by_timestamp) {
inference_state_->batch_timestamps_.emplace_back(
timestamp_and_input_tensors_by_tag.first);
for (const auto& input_tensor_and_tag :
timestamp_and_input_tensors_by_tag.second) {
inference_state_->input_tensor_batches_[input_tensor_and_tag.first]
.emplace_back(input_tensor_and_tag.second);
}
}
if (inference_state_->batch_timestamps_.size() == options_.batch_size() ||
options_.batched_input()) {
inference_state_to_process = std::move(inference_state_);
inference_state_ = std::unique_ptr<InferenceState>();
}
}
if (inference_state_to_process) {
MP_RETURN_IF_ERROR(OutputBatch(cc, std::move(inference_state_to_process)));
}
return absl::OkStatus();
}
absl::Status Close(CalculatorContext* cc) override {
std::unique_ptr<InferenceState> inference_state_to_process = nullptr;
{
absl::WriterMutexLock l(&mutex_);
if (cc->GraphStatus().ok() && inference_state_ != nullptr &&
!inference_state_->batch_timestamps_.empty()) {
inference_state_to_process = std::move(inference_state_);
inference_state_ = std::unique_ptr<InferenceState>();
}
}
if (inference_state_to_process) {
MP_RETURN_IF_ERROR(OutputBatch(cc, std::move(inference_state_to_process)));
}
return absl::OkStatus();
}
// When a batch of input tensors is ready to be run, runs TensorFlow and
// outputs the output tensors. The output tensors have timestamps matching
// the input tensor that formed that batch element. Any requested
// batch_dimension is added and removed. This code takes advantage of the fact
// that copying a tensor shares the same reference-counted, heap allocated
// memory buffer. Therefore, copies are cheap and should not cause the memory
// buffer to fall out of scope. In contrast, concat is only used where
// necessary.
absl::Status OutputBatch(CalculatorContext* cc,
std::unique_ptr<InferenceState> inference_state) {
const int64 start_time = absl::ToUnixMicros(clock_->TimeNow());
std::vector<std::pair<mediapipe::ProtoString, tf::Tensor>> input_tensors;
for (auto& keyed_tensors : inference_state->input_tensor_batches_) {
if (options_.batch_size() == 1) {
// Short circuit to avoid the cost of deep copying tensors in concat.
if (!keyed_tensors.second.empty()) {
// Pad by replicating the first tens or, then ignore the values.
keyed_tensors.second.resize(options_.batch_size());
std::fill(keyed_tensors.second.begin() +
inference_state->batch_timestamps_.size(),
keyed_tensors.second.end(), keyed_tensors.second[0]);
tf::Tensor concated;
const tf::Status concat_status =
tf::tensor::Concat(keyed_tensors.second, &concated);
CHECK(concat_status.ok()) << concat_status.ToString();
input_tensors.emplace_back(tag_to_tensor_map_[keyed_tensors.first],
keyed_tensors.second[0]);
} else {
// The input buffer can be empty for recurrent tensors.
RET_CHECK(
mediapipe::ContainsKey(recurrent_feed_tags_, keyed_tensors.first))
<< "A non-recurrent tensor does not have an input: "
<< keyed_tensors.first;
concated);
}
} else {
// Pad by replicating the first tens or, then ignore the values.
keyed_tensors.second.resize(options_.batch_size());
std::fill(keyed_tensors.second.begin() +
inference_state->batch_timestamps_.size(),
keyed_tensors.second.end(), keyed_tensors.second[0]);
tf::Tensor concated;
const tf::Status concat_status =
tf::tensor::Concat(keyed_tensors.second, &concated);
CHECK(concat_status.ok()) << concat_status.ToString();
input_tensors.emplace_back(tag_to_tensor_map_[keyed_tensors.first],
concated);
}
}
inference_state->input_tensor_batches_.clear();
std::vector<mediapipe::ProtoString> output_tensor_names;
std::vector<std::string> output_name_in_signature;
for (const std::string& tag : cc->Outputs().GetTags()) {
output_tensor_names.emplace_back(tag_to_tensor_map_[tag]);
output_name_in_signature.emplace_back(tag);
}
for (const auto& tag_pair : recurrent_fetch_tags_to_feed_tags_) {
// Ensure that we always fetch the recurrent state tensors.
if (std::find(output_name_in_signature.begin(),
output_name_in_signature.end(),
tag_pair.first) == output_name_in_signature.end()) {
output_tensor_names.emplace_back(tag_to_tensor_map_[tag_pair.first]);
output_name_in_signature.emplace_back(tag_pair.first);
inference_state->input_tensor_batches_.clear();
std::vector<mediapipe::ProtoString> output_tensor_names;
std::vector<std::string> output_name_in_signature;
for (const std::string& tag : cc->Outputs().GetTags()) {
output_tensor_names.emplace_back(tag_to_tensor_map_[tag]);
output_name_in_signature.emplace_back(tag);
}
}
std::vector<tf::Tensor> outputs;
for (const auto& tag_pair : recurrent_fetch_tags_to_feed_tags_) {
// Ensure that we always fetch the recurrent state tensors.
if (std::find(output_name_in_signature.begin(),
output_name_in_signature.end(),
tag_pair.first) == output_name_in_signature.end()) {
output_tensor_names.emplace_back(tag_to_tensor_map_[tag_pair.first]);
output_name_in_signature.emplace_back(tag_pair.first);
}
}
std::vector<tf::Tensor> outputs;
SimpleSemaphore* session_run_throttle = nullptr;
if (options_.max_concurrent_session_runs() > 0) {
session_run_throttle =
get_session_run_throttle(options_.max_concurrent_session_runs());
session_run_throttle->Acquire(1);
}
const int64 run_start_time = absl::ToUnixMicros(clock_->TimeNow());
tf::Status tf_status;
{
SimpleSemaphore* session_run_throttle = nullptr;
if (options_.max_concurrent_session_runs() > 0) {
session_run_throttle =
get_session_run_throttle(options_.max_concurrent_session_runs());
session_run_throttle->Acquire(1);
}
const int64 run_start_time = absl::ToUnixMicros(clock_->TimeNow());
tf::Status tf_status;
{
#if !defined(MEDIAPIPE_MOBILE) && !defined(__APPLE__)
tensorflow::profiler::TraceMe trace(absl::string_view(cc->NodeName()));
tensorflow::profiler::TraceMe trace(absl::string_view(cc->NodeName()));
#endif
tf_status = session_->Run(input_tensors, output_tensor_names,
{} /* target_node_names */, &outputs);
}
tf_status = session_->Run(input_tensors, output_tensor_names,
{} /* target_node_names */, &outputs);
}
if (session_run_throttle != nullptr) {
session_run_throttle->Release(1);
}
if (session_run_throttle != nullptr) {
session_run_throttle->Release(1);
}
// RET_CHECK on the tf::Status object itself in order to print an
// informative error message.
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
// RET_CHECK on the tf::Status object itself in order to print an
// informative error message.
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
const int64 run_end_time = absl::ToUnixMicros(clock_->TimeNow());
cc->GetCounter(kTotalSessionRunsTimeUsecsCounterSuffix)
->IncrementBy(run_end_time - run_start_time);
cc->GetCounter(kTotalNumSessionRunsCounterSuffix)->Increment();
const int64 run_end_time = absl::ToUnixMicros(clock_->TimeNow());
cc->GetCounter(kTotalSessionRunsTimeUsecsCounterSuffix)
->IncrementBy(run_end_time - run_start_time);
cc->GetCounter(kTotalNumSessionRunsCounterSuffix)->Increment();
// Feed back the recurrent state.
for (const auto& tag_pair : recurrent_fetch_tags_to_feed_tags_) {
int pos = std::find(output_name_in_signature.begin(),
output_name_in_signature.end(), tag_pair.first) -
output_name_in_signature.begin();
inference_state->input_tensor_batches_[tag_pair.second].emplace_back(
outputs[pos]);
}
// Feed back the recurrent state.
for (const auto& tag_pair : recurrent_fetch_tags_to_feed_tags_) {
int pos = std::find(output_name_in_signature.begin(),
output_name_in_signature.end(), tag_pair.first) -
output_name_in_signature.begin();
inference_state->input_tensor_batches_[tag_pair.second].emplace_back(
outputs[pos]);
}
absl::WriterMutexLock l(&mutex_);
// Set that we want to split on each index of the 0th dimension.
std::vector<tf::int64> split_vector(options_.batch_size(), 1);
for (int i = 0; i < output_tensor_names.size(); ++i) {
if (options_.batch_size() == 1) {
if (cc->Outputs().HasTag(output_name_in_signature[i])) {
tf::Tensor output_tensor(outputs[i]);
RET_CHECK_OK(RemoveBatchDimension(&output_tensor));
cc->Outputs()
.Tag(output_name_in_signature[i])
.Add(new tf::Tensor(output_tensor),
inference_state->batch_timestamps_[0]);
}
} else {
std::vector<tf::Tensor> split_tensors;
const tf::Status split_status =
tf::tensor::Split(outputs[i], split_vector, &split_tensors);
CHECK(split_status.ok()) << split_status.ToString();
// Loop over timestamps so that we don't copy the padding.
for (int j = 0; j < inference_state->batch_timestamps_.size(); ++j) {
tf::Tensor output_tensor(split_tensors[j]);
RET_CHECK_OK(RemoveBatchDimension(&output_tensor));
cc->Outputs()
.Tag(output_name_in_signature[i])
.Add(new tf::Tensor(output_tensor),
inference_state->batch_timestamps_[j]);
absl::WriterMutexLock l(&mutex_);
// Set that we want to split on each index of the 0th dimension.
std::vector<tf::int64> split_vector(options_.batch_size(), 1);
for (int i = 0; i < output_tensor_names.size(); ++i) {
if (options_.batch_size() == 1) {
if (cc->Outputs().HasTag(output_name_in_signature[i])) {
tf::Tensor output_tensor(outputs[i]);
RET_CHECK_OK(RemoveBatchDimension(&output_tensor));
cc->Outputs()
.Tag(output_name_in_signature[i])
.Add(new tf::Tensor(output_tensor),
inference_state->batch_timestamps_[0]);
}
} else {
std::vector<tf::Tensor> split_tensors;
const tf::Status split_status =
tf::tensor::Split(outputs[i], split_vector, &split_tensors);
CHECK(split_status.ok()) << split_status.ToString();
// Loop over timestamps so that we don't copy the padding.
for (int j = 0; j < inference_state->batch_timestamps_.size(); ++j) {
tf::Tensor output_tensor(split_tensors[j]);
RET_CHECK_OK(RemoveBatchDimension(&output_tensor));
cc->Outputs()
.Tag(output_name_in_signature[i])
.Add(new tf::Tensor(output_tensor),
inference_state->batch_timestamps_[j]);
}
}
}
// Get end time and report.
const int64 end_time = absl::ToUnixMicros(clock_->TimeNow());
cc->GetCounter(kTotalUsecsCounterSuffix)
->IncrementBy(end_time - start_time);
cc->GetCounter(kTotalProcessedTimestampsCounterSuffix)
->IncrementBy(inference_state->batch_timestamps_.size());
// Make sure we hold on to the recursive state.
if (!options_.recurrent_tag_pair().empty()) {
inference_state_ = std::move(inference_state);
inference_state_->batch_timestamps_.clear();
}
return absl::OkStatus();
}
// Get end time and report.
const int64 end_time = absl::ToUnixMicros(clock_->TimeNow());
cc->GetCounter(kTotalUsecsCounterSuffix)->IncrementBy(end_time - start_time);
cc->GetCounter(kTotalProcessedTimestampsCounterSuffix)
->IncrementBy(inference_state->batch_timestamps_.size());
private:
// The Session object is provided by a packet factory and is owned by the
// MediaPipe framework. Individual calls are thread-safe, but session state
// may be shared across threads.
tf::Session* session_;
// Make sure we hold on to the recursive state.
if (!options_.recurrent_tag_pair().empty()) {
inference_state_ = std::move(inference_state);
inference_state_->batch_timestamps_.clear();
// A mapping between stream tags and the tensor names they are bound to.
std::map<std::string, std::string> tag_to_tensor_map_;
absl::Mutex mutex_;
std::unique_ptr<InferenceState> inference_state_ ABSL_GUARDED_BY(mutex_);
// The options for the calculator.
TensorFlowInferenceCalculatorOptions options_;
// Store the feed and fetch tags for feed/fetch recurrent networks.
std::set<std::string> recurrent_feed_tags_;
std::map<std::string, std::string> recurrent_fetch_tags_to_feed_tags_;
// Clock used to measure the computation time in OutputBatch().
std::unique_ptr<mediapipe::Clock> clock_;
// The static singleton semaphore to throttle concurrent session runs.
static SimpleSemaphore* get_session_run_throttle(
int32 max_concurrent_session_runs) {
static SimpleSemaphore* session_run_throttle =
new SimpleSemaphore(max_concurrent_session_runs);
return session_run_throttle;
}
return absl::OkStatus();
}
private:
// The Session object is provided by a packet factory and is owned by the
// MediaPipe framework. Individual calls are thread-safe, but session state may
// be shared across threads.
tf::Session* session_;
// A mapping between stream tags and the tensor names they are bound to.
std::map<std::string, std::string> tag_to_tensor_map_;
absl::Mutex mutex_;
std::unique_ptr<InferenceState> inference_state_ ABSL_GUARDED_BY(mutex_);
// The options for the calculator.
TensorFlowInferenceCalculatorOptions options_;
// Store the feed and fetch tags for feed/fetch recurrent networks.
std::set<std::string> recurrent_feed_tags_;
std::map<std::string, std::string> recurrent_fetch_tags_to_feed_tags_;
// Clock used to measure the computation time in OutputBatch().
std::unique_ptr<mediapipe::Clock> clock_;
// The static singleton semaphore to throttle concurrent session runs.
static SimpleSemaphore* get_session_run_throttle(
int32 max_concurrent_session_runs) {
static SimpleSemaphore* session_run_throttle =
new SimpleSemaphore(max_concurrent_session_runs);
return session_run_throttle;
}
}
;
};
REGISTER_CALCULATOR(TensorFlowInferenceCalculator);
constexpr char TensorFlowInferenceCalculator::kTotalUsecsCounterSuffix[];
mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_calculator.cc
+1
View File
@@ -80,6 +80,7 @@ const std::string MaybeConvertSignatureToTag(
// which in turn contains a TensorFlow Session ready for execution and a map
// between tags and tensor names.
//
//
// Example usage:
// node {
// calculator: "TensorFlowSessionFromSavedModelCalculator"
mediapipe/calculators/tensorflow/unpack_media_sequence_calculator.cc
+52 -30
View File
@@ -217,38 +217,41 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
first_timestamp_seen_ = recent_timestamp;
}
}
if (recent_timestamp > last_timestamp_seen) {
if (recent_timestamp > last_timestamp_seen &&
recent_timestamp < Timestamp::PostStream().Value()) {
last_timestamp_key_ = map_kv.first;
last_timestamp_seen = recent_timestamp;
}
}
}
if (!timestamps_.empty()) {
RET_CHECK(!last_timestamp_key_.empty())
<< "Something went wrong because the timestamp key is unset. "
"Example: "
<< sequence_->DebugString();
RET_CHECK_GT(last_timestamp_seen, Timestamp::PreStream().Value())
<< "Something went wrong because the last timestamp is unset. "
"Example: "
<< sequence_->DebugString();
RET_CHECK_LT(first_timestamp_seen_,
Timestamp::OneOverPostStream().Value())
<< "Something went wrong because the first timestamp is unset. "
"Example: "
<< sequence_->DebugString();
for (const auto& kv : timestamps_) {
if (!kv.second.empty() &&
kv.second[0] < Timestamp::PostStream().Value()) {
// These checks only make sense if any values are not PostStream, but
// only need to be made once.
RET_CHECK(!last_timestamp_key_.empty())
<< "Something went wrong because the timestamp key is unset. "
<< "Example: " << sequence_->DebugString();
RET_CHECK_GT(last_timestamp_seen, Timestamp::PreStream().Value())
<< "Something went wrong because the last timestamp is unset. "
<< "Example: " << sequence_->DebugString();
RET_CHECK_LT(first_timestamp_seen_,
Timestamp::OneOverPostStream().Value())
<< "Something went wrong because the first timestamp is unset. "
<< "Example: " << sequence_->DebugString();
break;
}
}
}
current_timestamp_index_ = 0;
process_poststream_ = false;
// Determine the data path and output it.
const auto& options = cc->Options<UnpackMediaSequenceCalculatorOptions>();
const auto& sequence = cc->InputSidePackets()
.Tag(kSequenceExampleTag)
.Get<tensorflow::SequenceExample>();
if (cc->Outputs().HasTag(kKeypointsTag)) {
keypoint_names_ = absl::StrSplit(options.keypoint_names(), ',');
default_keypoint_location_ = options.default_keypoint_location();
}
if (cc->OutputSidePackets().HasTag(kDataPath)) {
std::string root_directory = "";
if (cc->InputSidePackets().HasTag(kDatasetRootDirTag)) {
@@ -349,19 +352,30 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
// all packets on all streams that have a timestamp between the current
// reference timestep and the previous reference timestep. This ensures that
// we emit all timestamps in order, but also only emit a limited number in
// any particular call to Process().
int64 start_timestamp =
timestamps_[last_timestamp_key_][current_timestamp_index_];
if (current_timestamp_index_ == 0) {
start_timestamp = first_timestamp_seen_;
// any particular call to Process(). At the every end, we output the
// poststream packets. If we only have poststream packets,
// last_timestamp_key_ will be empty.
int64 start_timestamp = 0;
int64 end_timestamp = 0;
if (last_timestamp_key_.empty() || process_poststream_) {
process_poststream_ = true;
start_timestamp = Timestamp::PostStream().Value();
end_timestamp = Timestamp::OneOverPostStream().Value();
} else {
start_timestamp =
timestamps_[last_timestamp_key_][current_timestamp_index_];
if (current_timestamp_index_ == 0) {
start_timestamp = first_timestamp_seen_;
}
end_timestamp = start_timestamp + 1; // Base case at end of sequence.
if (current_timestamp_index_ <
timestamps_[last_timestamp_key_].size() - 1) {
end_timestamp =
timestamps_[last_timestamp_key_][current_timestamp_index_ + 1];
}
}
int64 end_timestamp = start_timestamp + 1; // Base case at end of sequence.
if (current_timestamp_index_ <
timestamps_[last_timestamp_key_].size() - 1) {
end_timestamp =
timestamps_[last_timestamp_key_][current_timestamp_index_ + 1];
}
for (const auto& map_kv : timestamps_) {
for (int i = 0; i < map_kv.second.size(); ++i) {
if (map_kv.second[i] >= start_timestamp &&
@@ -438,7 +452,14 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
if (current_timestamp_index_ < timestamps_[last_timestamp_key_].size()) {
return absl::OkStatus();
} else {
return tool::StatusStop();
if (process_poststream_) {
// Once we've processed the PostStream timestamp we can stop.
return tool::StatusStop();
} else {
// Otherwise, we still need to do one more pass to process it.
process_poststream_ = true;
return absl::OkStatus();
}
}
}
@@ -462,6 +483,7 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
std::vector<std::string> keypoint_names_;
// Default keypoint location when missing.
float default_keypoint_location_;
bool process_poststream_;
};
REGISTER_CALCULATOR(UnpackMediaSequenceCalculator);
} // namespace mediapipe
mediapipe/calculators/tensorflow/unpack_media_sequence_calculator_test.cc
+66
View File
@@ -412,6 +412,72 @@ TEST_F(UnpackMediaSequenceCalculatorTest, UnpacksTwoPostStreamFloatLists) {
::testing::Eq(Timestamp::PostStream()));
}
TEST_F(UnpackMediaSequenceCalculatorTest, UnpacksImageWithPostStreamFloatList) {
SetUpCalculator({"IMAGE:images"}, {});
auto input_sequence = absl::make_unique<tf::SequenceExample>();
std::string test_video_id = "test_video_id";
mpms::SetClipMediaId(test_video_id, input_sequence.get());
std::string test_image_string = "test_image_string";
int num_images = 1;
for (int i = 0; i < num_images; ++i) {
mpms::AddImageTimestamp(i, input_sequence.get());
mpms::AddImageEncoded(test_image_string, input_sequence.get());
}
mpms::AddFeatureFloats("FDENSE_MAX", {3.0f, 4.0f}, input_sequence.get());
mpms::AddFeatureTimestamp("FDENSE_MAX", Timestamp::PostStream().Value(),
input_sequence.get());
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
Adopt(input_sequence.release());
MP_ASSERT_OK(runner_->Run());
const std::vector<Packet>& output_packets =
runner_->Outputs().Tag("IMAGE").packets;
ASSERT_EQ(num_images, output_packets.size());
for (int i = 0; i < num_images; ++i) {
const std::string& output_image = output_packets[i].Get<std::string>();
ASSERT_EQ(output_image, test_image_string);
}
}
TEST_F(UnpackMediaSequenceCalculatorTest, UnpacksPostStreamFloatListWithImage) {
SetUpCalculator({"FLOAT_FEATURE_FDENSE_MAX:max"}, {});
auto input_sequence = absl::make_unique<tf::SequenceExample>();
std::string test_video_id = "test_video_id";
mpms::SetClipMediaId(test_video_id, input_sequence.get());
std::string test_image_string = "test_image_string";
int num_images = 1;
for (int i = 0; i < num_images; ++i) {
mpms::AddImageTimestamp(i, input_sequence.get());
mpms::AddImageEncoded(test_image_string, input_sequence.get());
}
mpms::AddFeatureFloats("FDENSE_MAX", {3.0f, 4.0f}, input_sequence.get());
mpms::AddFeatureTimestamp("FDENSE_MAX", Timestamp::PostStream().Value(),
input_sequence.get());
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
Adopt(input_sequence.release());
MP_ASSERT_OK(runner_->Run());
const std::vector<Packet>& fdense_max_packets =
runner_->Outputs().Tag("FLOAT_FEATURE_FDENSE_MAX").packets;
ASSERT_EQ(fdense_max_packets.size(), 1);
const auto& fdense_max_vector =
fdense_max_packets[0].Get<std::vector<float>>();
ASSERT_THAT(fdense_max_vector, ::testing::ElementsAreArray({3.0f, 4.0f}));
ASSERT_THAT(fdense_max_packets[0].Timestamp(),
::testing::Eq(Timestamp::PostStream()));
}
TEST_F(UnpackMediaSequenceCalculatorTest, GetDatasetFromPacket) {
SetUpCalculator({}, {"DATA_PATH:data_path"}, {"DATASET_ROOT:root"});
mediapipe/calculators/tflite/tflite_inference_calculator.cc
+7 -5
View File
@@ -904,7 +904,8 @@ absl::Status TfLiteInferenceCalculator::LoadDelegate(CalculatorContext* cc) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
// Configure and create the delegate.
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
options.compile_options.precision_loss_allowed = 1;
options.compile_options.precision_loss_allowed =
allow_precision_loss_ ? 1 : 0;
options.compile_options.preferred_gl_object_type =
TFLITE_GL_OBJECT_TYPE_FASTEST;
options.compile_options.dynamic_batch_enabled = 0;
@@ -968,7 +969,7 @@ absl::Status TfLiteInferenceCalculator::LoadDelegate(CalculatorContext* cc) {
const int kHalfSize = 2; // sizeof(half)
// Configure and create the delegate.
TFLGpuDelegateOptions options;
options.allow_precision_loss = true;
options.allow_precision_loss = allow_precision_loss_;
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypeActive;
if (!delegate_)
delegate_ = TfLiteDelegatePtr(TFLGpuDelegateCreate(&options),
@@ -1080,9 +1081,10 @@ absl::Status TfLiteInferenceCalculator::LoadDelegate(CalculatorContext* cc) {
}
// Create converter for GPU output.
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
isFloat16:true
convertToPBHWC4:false];
converter_from_BPHWC4_ =
[[TFLBufferConvert alloc] initWithDevice:device
isFloat16:allow_precision_loss_
convertToPBHWC4:false];
if (converter_from_BPHWC4_ == nil) {
return absl::InternalError(
"Error initializating output buffer converter");
mediapipe/calculators/tflite/tflite_tensors_to_segmentation_calculator.cc
+1 -1
View File
@@ -439,7 +439,7 @@ absl::Status TfLiteTensorsToSegmentationCalculator::ProcessGpu(
// Run shader, upsample result.
{
gpu_helper_.BindFramebuffer(output_texture); // GL_TEXTURE0
gpu_helper_.BindFramebuffer(output_texture);
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, small_mask_texture.id());
GlRender();
mediapipe/calculators/util/BUILD
+61
View File
@@ -821,6 +821,25 @@ cc_library(
alwayslink = 1,
)
cc_test(
name = "landmark_projection_calculator_test",
srcs = ["landmark_projection_calculator_test.cc"],
deps = [
":landmark_projection_calculator",
"//mediapipe/calculators/tensor:image_to_tensor_utils",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework/deps:message_matchers",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
"@com_google_absl//absl/memory",
"@com_google_googletest//:gtest_main",
],
)
mediapipe_proto_library(
name = "landmarks_smoothing_calculator_proto",
srcs = ["landmarks_smoothing_calculator.proto"],
@@ -1252,3 +1271,45 @@ cc_test(
"//mediapipe/framework/port:parse_text_proto",
],
)
mediapipe_proto_library(
name = "refine_landmarks_from_heatmap_calculator_proto",
srcs = ["refine_landmarks_from_heatmap_calculator.proto"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
],
)
cc_library(
name = "refine_landmarks_from_heatmap_calculator",
srcs = ["refine_landmarks_from_heatmap_calculator.cc"],
hdrs = ["refine_landmarks_from_heatmap_calculator.h"],
copts = select({
"//mediapipe:apple": [
"-x objective-c++",
"-fobjc-arc", # enable reference-counting
],
"//conditions:default": [],
}),
visibility = ["//visibility:public"],
deps = [
":refine_landmarks_from_heatmap_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:statusor",
],
alwayslink = 1,
)
cc_test(
name = "refine_landmarks_from_heatmap_calculator_test",
srcs = ["refine_landmarks_from_heatmap_calculator_test.cc"],
deps = [
":refine_landmarks_from_heatmap_calculator",
"//mediapipe/framework/port:gtest_main",
],
)
mediapipe/calculators/util/annotation_overlay_calculator.cc
+1 -1
View File
@@ -402,7 +402,7 @@ absl::Status AnnotationOverlayCalculator::RenderToGpu(CalculatorContext* cc,
// Blend overlay image in GPU shader.
{
gpu_helper_.BindFramebuffer(output_texture); // GL_TEXTURE0
gpu_helper_.BindFramebuffer(output_texture);
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, input_texture.name());
mediapipe/calculators/util/detection_label_id_to_text_calculator.cc
+7 -6
View File
@@ -54,6 +54,7 @@ class DetectionLabelIdToTextCalculator : public CalculatorBase {
private:
absl::node_hash_map<int, std::string> label_map_;
::mediapipe::DetectionLabelIdToTextCalculatorOptions options_;
};
REGISTER_CALCULATOR(DetectionLabelIdToTextCalculator);
@@ -68,13 +69,13 @@ absl::Status DetectionLabelIdToTextCalculator::GetContract(
absl::Status DetectionLabelIdToTextCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
const auto& options =
options_ =
cc->Options<::mediapipe::DetectionLabelIdToTextCalculatorOptions>();
if (options.has_label_map_path()) {
if (options_.has_label_map_path()) {
std::string string_path;
ASSIGN_OR_RETURN(string_path,
PathToResourceAsFile(options.label_map_path()));
PathToResourceAsFile(options_.label_map_path()));
std::string label_map_string;
MP_RETURN_IF_ERROR(file::GetContents(string_path, &label_map_string));
@@ -85,8 +86,8 @@ absl::Status DetectionLabelIdToTextCalculator::Open(CalculatorContext* cc) {
label_map_[i++] = line;
}
} else {
for (int i = 0; i < options.label_size(); ++i) {
label_map_[i] = options.label(i);
for (int i = 0; i < options_.label_size(); ++i) {
label_map_[i] = options_.label(i);
}
}
return absl::OkStatus();
@@ -106,7 +107,7 @@ absl::Status DetectionLabelIdToTextCalculator::Process(CalculatorContext* cc) {
}
}
// Remove label_id field if text labels exist.
if (has_text_label) {
if (has_text_label && !options_.keep_label_id()) {
output_detection.clear_label_id();
}
}
mediapipe/calculators/util/detection_label_id_to_text_calculator.proto
+5
View File
@@ -31,4 +31,9 @@ message DetectionLabelIdToTextCalculatorOptions {
// label: "label for id 1"
// ...
repeated string label = 2;
// By default, the `label_id` field from the input is stripped if a text label
// could be found. By setting this field to true, it is always copied to the
// output detections.
optional bool keep_label_id = 3;
}
mediapipe/calculators/util/labels_to_render_data_calculator.cc
+5 -1
View File
@@ -120,7 +120,11 @@ absl::Status LabelsToRenderDataCalculator::Process(CalculatorContext* cc) {
labels.resize(classifications.classification_size());
scores.resize(classifications.classification_size());
for (int i = 0; i < classifications.classification_size(); ++i) {
labels[i] = classifications.classification(i).label();
if (options_.use_display_name()) {
labels[i] = classifications.classification(i).display_name();
} else {
labels[i] = classifications.classification(i).label();
}
scores[i] = classifications.classification(i).score();
}
} else {
mediapipe/calculators/util/labels_to_render_data_calculator.proto
+3
View File
@@ -59,4 +59,7 @@ message LabelsToRenderDataCalculatorOptions {
BOTTOM_LEFT = 1;
}
optional Location location = 6 [default = TOP_LEFT];
// Uses Classification.display_name field instead of Classification.label.
optional bool use_display_name = 9 [default = false];
}
mediapipe/calculators/util/landmark_projection_calculator.cc
+124 -36
View File
@@ -13,6 +13,7 @@
// limitations under the License.
#include <cmath>
#include <functional>
#include <vector>
#include "mediapipe/calculators/util/landmark_projection_calculator.pb.h"
@@ -27,20 +28,32 @@ namespace {
constexpr char kLandmarksTag[] = "NORM_LANDMARKS";
constexpr char kRectTag[] = "NORM_RECT";
constexpr char kProjectionMatrix[] = "PROJECTION_MATRIX";
} // namespace
// Projects normalized landmarks in a rectangle to its original coordinates. The
// rectangle must also be in normalized coordinates.
// Projects normalized landmarks to its original coordinates.
// Input:
// NORM_LANDMARKS: A NormalizedLandmarkList representing landmarks
// in a normalized rectangle.
// NORM_RECT: An NormalizedRect representing a normalized rectangle in image
// coordinates.
// NORM_LANDMARKS - NormalizedLandmarkList
// Represents landmarks in a normalized rectangle if NORM_RECT is specified
// or landmarks that should be projected using PROJECTION_MATRIX if
// specified. (Prefer using PROJECTION_MATRIX as it eliminates need of
// letterbox removal step.)
// NORM_RECT - NormalizedRect
// Represents a normalized rectangle in image coordinates and results in
// landmarks with their locations adjusted to the image.
// PROJECTION_MATRIX - std::array<float, 16>
// A 4x4 row-major-order matrix that maps landmarks' locations from one
// coordinate system to another. In this case from the coordinate system of
// the normalized region of interest to the coordinate system of the image.
//
// Note: either NORM_RECT or PROJECTION_MATRIX has to be specified.
// Note: landmark's Z is projected in a custom way - it's scaled by width of
// the normalized region of interest used during landmarks detection.
//
// Output:
// NORM_LANDMARKS: A NormalizedLandmarkList representing landmarks
// with their locations adjusted to the image.
// NORM_LANDMARKS - NormalizedLandmarkList
// Landmarks with their locations adjusted according to the inputs.
//
// Usage example:
// node {
@@ -58,12 +71,27 @@ constexpr char kRectTag[] = "NORM_RECT";
// output_stream: "NORM_LANDMARKS:0:projected_landmarks_0"
// output_stream: "NORM_LANDMARKS:1:projected_landmarks_1"
// }
//
// node {
// calculator: "LandmarkProjectionCalculator"
// input_stream: "NORM_LANDMARKS:landmarks"
// input_stream: "PROECTION_MATRIX:matrix"
// output_stream: "NORM_LANDMARKS:projected_landmarks"
// }
//
// node {
// calculator: "LandmarkProjectionCalculator"
// input_stream: "NORM_LANDMARKS:0:landmarks_0"
// input_stream: "NORM_LANDMARKS:1:landmarks_1"
// input_stream: "PROECTION_MATRIX:matrix"
// output_stream: "NORM_LANDMARKS:0:projected_landmarks_0"
// output_stream: "NORM_LANDMARKS:1:projected_landmarks_1"
// }
class LandmarkProjectionCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag(kLandmarksTag) &&
cc->Inputs().HasTag(kRectTag))
<< "Missing one or more input streams.";
RET_CHECK(cc->Inputs().HasTag(kLandmarksTag))
<< "Missing NORM_LANDMARKS input.";
RET_CHECK_EQ(cc->Inputs().NumEntries(kLandmarksTag),
cc->Outputs().NumEntries(kLandmarksTag))
@@ -73,7 +101,14 @@ class LandmarkProjectionCalculator : public CalculatorBase {
id != cc->Inputs().EndId(kLandmarksTag); ++id) {
cc->Inputs().Get(id).Set<NormalizedLandmarkList>();
}
cc->Inputs().Tag(kRectTag).Set<NormalizedRect>();
RET_CHECK(cc->Inputs().HasTag(kRectTag) ^
cc->Inputs().HasTag(kProjectionMatrix))
<< "Either NORM_RECT or PROJECTION_MATRIX must be specified.";
if (cc->Inputs().HasTag(kRectTag)) {
cc->Inputs().Tag(kRectTag).Set<NormalizedRect>();
} else {
cc->Inputs().Tag(kProjectionMatrix).Set<std::array<float, 16>>();
}
for (CollectionItemId id = cc->Outputs().BeginId(kLandmarksTag);
id != cc->Outputs().EndId(kLandmarksTag); ++id) {
@@ -89,31 +124,50 @@ class LandmarkProjectionCalculator : public CalculatorBase {
return absl::OkStatus();
}
static void ProjectXY(const NormalizedLandmark& lm,
const std::array<float, 16>& matrix,
NormalizedLandmark* out) {
out->set_x(lm.x() * matrix[0] + lm.y() * matrix[1] + lm.z() * matrix[2] +
matrix[3]);
out->set_y(lm.x() * matrix[4] + lm.y() * matrix[5] + lm.z() * matrix[6] +
matrix[7]);
}
/**
* Landmark's Z scale is equal to a relative (to image) width of region of
* interest used during detection. To calculate based on matrix:
* 1. Project (0,0) --- (1,0) segment using matrix.
* 2. Calculate length of the projected segment.
*/
static float CalculateZScale(const std::array<float, 16>& matrix) {
NormalizedLandmark a;
a.set_x(0.0f);
a.set_y(0.0f);
NormalizedLandmark b;
b.set_x(1.0f);
b.set_y(0.0f);
NormalizedLandmark a_projected;
ProjectXY(a, matrix, &a_projected);
NormalizedLandmark b_projected;
ProjectXY(b, matrix, &b_projected);
return std::sqrt(std::pow(b_projected.x() - a_projected.x(), 2) +
std::pow(b_projected.y() - a_projected.y(), 2));
}
absl::Status Process(CalculatorContext* cc) override {
if (cc->Inputs().Tag(kRectTag).IsEmpty()) {
return absl::OkStatus();
}
const auto& input_rect = cc->Inputs().Tag(kRectTag).Get<NormalizedRect>();
const auto& options =
cc->Options<::mediapipe::LandmarkProjectionCalculatorOptions>();
CollectionItemId input_id = cc->Inputs().BeginId(kLandmarksTag);
CollectionItemId output_id = cc->Outputs().BeginId(kLandmarksTag);
// Number of inputs and outpus is the same according to the contract.
for (; input_id != cc->Inputs().EndId(kLandmarksTag);
++input_id, ++output_id) {
const auto& input_packet = cc->Inputs().Get(input_id);
if (input_packet.IsEmpty()) {
continue;
std::function<void(const NormalizedLandmark&, NormalizedLandmark*)>
project_fn;
if (cc->Inputs().HasTag(kRectTag)) {
if (cc->Inputs().Tag(kRectTag).IsEmpty()) {
return absl::OkStatus();
}
const auto& input_landmarks = input_packet.Get<NormalizedLandmarkList>();
NormalizedLandmarkList output_landmarks;
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
const NormalizedLandmark& landmark = input_landmarks.landmark(i);
NormalizedLandmark* new_landmark = output_landmarks.add_landmark();
const auto& input_rect = cc->Inputs().Tag(kRectTag).Get<NormalizedRect>();
const auto& options =
cc->Options<mediapipe::LandmarkProjectionCalculatorOptions>();
project_fn = [&input_rect, &options](const NormalizedLandmark& landmark,
NormalizedLandmark* new_landmark) {
// TODO: fix projection or deprecate (current projection
// calculations are incorrect for general case).
const float x = landmark.x() - 0.5f;
const float y = landmark.y() - 0.5f;
const float angle =
@@ -130,10 +184,44 @@ class LandmarkProjectionCalculator : public CalculatorBase {
new_landmark->set_x(new_x);
new_landmark->set_y(new_y);
new_landmark->set_z(new_z);
};
} else if (cc->Inputs().HasTag(kProjectionMatrix)) {
if (cc->Inputs().Tag(kProjectionMatrix).IsEmpty()) {
return absl::OkStatus();
}
const auto& project_mat =
cc->Inputs().Tag(kProjectionMatrix).Get<std::array<float, 16>>();
const float z_scale = CalculateZScale(project_mat);
project_fn = [&project_mat, z_scale](const NormalizedLandmark& lm,
NormalizedLandmark* new_landmark) {
*new_landmark = lm;
ProjectXY(lm, project_mat, new_landmark);
new_landmark->set_z(z_scale * lm.z());
};
} else {
return absl::InternalError("Either rect or matrix must be specified.");
}
CollectionItemId input_id = cc->Inputs().BeginId(kLandmarksTag);
CollectionItemId output_id = cc->Outputs().BeginId(kLandmarksTag);
// Number of inputs and outpus is the same according to the contract.
for (; input_id != cc->Inputs().EndId(kLandmarksTag);
++input_id, ++output_id) {
const auto& input_packet = cc->Inputs().Get(input_id);
if (input_packet.IsEmpty()) {
continue;
}
const auto& input_landmarks = input_packet.Get<NormalizedLandmarkList>();
NormalizedLandmarkList output_landmarks;
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
const NormalizedLandmark& landmark = input_landmarks.landmark(i);
NormalizedLandmark* new_landmark = output_landmarks.add_landmark();
project_fn(landmark, new_landmark);
}
cc->Outputs().Get(output_id).AddPacket(
MakePacket<NormalizedLandmarkList>(output_landmarks)
MakePacket<NormalizedLandmarkList>(std::move(output_landmarks))
.At(cc->InputTimestamp()));
}
return absl::OkStatus();
mediapipe/calculators/util/landmark_projection_calculator_test.cc
+240
View File
@@ -0,0 +1,240 @@
#include <array>
#include <vector>
#include "absl/memory/memory.h"
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_runner.h"
#include "mediapipe/framework/deps/message_matchers.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
namespace {
absl::StatusOr<mediapipe::NormalizedLandmarkList> RunCalculator(
mediapipe::NormalizedLandmarkList input, mediapipe::NormalizedRect rect) {
mediapipe::CalculatorRunner runner(
ParseTextProtoOrDie<mediapipe::CalculatorGraphConfig::Node>(R"pb(
calculator: "LandmarkProjectionCalculator"
input_stream: "NORM_LANDMARKS:landmarks"
input_stream: "NORM_RECT:rect"
output_stream: "NORM_LANDMARKS:projected_landmarks"
)pb"));
runner.MutableInputs()
->Tag("NORM_LANDMARKS")
.packets.push_back(
MakePacket<mediapipe::NormalizedLandmarkList>(std::move(input))
.At(Timestamp(1)));
runner.MutableInputs()
->Tag("NORM_RECT")
.packets.push_back(MakePacket<mediapipe::NormalizedRect>(std::move(rect))
.At(Timestamp(1)));
MP_RETURN_IF_ERROR(runner.Run());
const auto& output_packets = runner.Outputs().Tag("NORM_LANDMARKS").packets;
RET_CHECK_EQ(output_packets.size(), 1);
return output_packets[0].Get<mediapipe::NormalizedLandmarkList>();
}
TEST(LandmarkProjectionCalculatorTest, ProjectingWithDefaultRect) {
mediapipe::NormalizedLandmarkList landmarks =
ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 10, y: 20, z: -0.5 }
)pb");
mediapipe::NormalizedRect rect =
ParseTextProtoOrDie<mediapipe::NormalizedRect>(
R"pb(
x_center: 0.5,
y_center: 0.5,
width: 1.0,
height: 1.0,
rotation: 0.0
)pb");
auto status_or_result = RunCalculator(std::move(landmarks), std::move(rect));
MP_ASSERT_OK(status_or_result);
EXPECT_THAT(
status_or_result.value(),
EqualsProto(ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 10, y: 20, z: -0.5 }
)pb")));
}
mediapipe::NormalizedRect GetCroppedRect() {
return ParseTextProtoOrDie<mediapipe::NormalizedRect>(
R"pb(
x_center: 0.5, y_center: 0.5, width: 0.5, height: 2, rotation: 0.0
)pb");
}
mediapipe::NormalizedLandmarkList GetCroppedRectTestInput() {
return ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 1.0, y: 1.0, z: -0.5 }
)pb");
}
mediapipe::NormalizedLandmarkList GetCroppedRectTestExpectedResult() {
return ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 0.75, y: 1.5, z: -0.25 }
)pb");
}
TEST(LandmarkProjectionCalculatorTest, ProjectingWithCroppedRect) {
auto status_or_result =
RunCalculator(GetCroppedRectTestInput(), GetCroppedRect());
MP_ASSERT_OK(status_or_result);
EXPECT_THAT(status_or_result.value(),
EqualsProto(GetCroppedRectTestExpectedResult()));
}
absl::StatusOr<mediapipe::NormalizedLandmarkList> RunCalculator(
mediapipe::NormalizedLandmarkList input, std::array<float, 16> matrix) {
mediapipe::CalculatorRunner runner(
ParseTextProtoOrDie<mediapipe::CalculatorGraphConfig::Node>(R"pb(
calculator: "LandmarkProjectionCalculator"
input_stream: "NORM_LANDMARKS:landmarks"
input_stream: "PROJECTION_MATRIX:matrix"
output_stream: "NORM_LANDMARKS:projected_landmarks"
)pb"));
runner.MutableInputs()
->Tag("NORM_LANDMARKS")
.packets.push_back(
MakePacket<mediapipe::NormalizedLandmarkList>(std::move(input))
.At(Timestamp(1)));
runner.MutableInputs()
->Tag("PROJECTION_MATRIX")
.packets.push_back(MakePacket<std::array<float, 16>>(std::move(matrix))
.At(Timestamp(1)));
MP_RETURN_IF_ERROR(runner.Run());
const auto& output_packets = runner.Outputs().Tag("NORM_LANDMARKS").packets;
RET_CHECK_EQ(output_packets.size(), 1);
return output_packets[0].Get<mediapipe::NormalizedLandmarkList>();
}
TEST(LandmarkProjectionCalculatorTest, ProjectingWithIdentityMatrix) {
mediapipe::NormalizedLandmarkList landmarks =
ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 10, y: 20, z: -0.5 }
)pb");
// clang-format off
std::array<float, 16> matrix = {
1.0f, 0.0f, 0.0f, 0.0f,
0.0f, 1.0f, 0.0f, 0.0f,
0.0f, 0.0f, 1.0f, 0.0f,
0.0f, 0.0f, 0.0f, 1.0f,
};
// clang-format on
auto status_or_result =
RunCalculator(std::move(landmarks), std::move(matrix));
MP_ASSERT_OK(status_or_result);
EXPECT_THAT(
status_or_result.value(),
EqualsProto(ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 10, y: 20, z: -0.5 }
)pb")));
}
TEST(LandmarkProjectionCalculatorTest, ProjectingWithCroppedRectMatrix) {
constexpr int kRectWidth = 1280;
constexpr int kRectHeight = 720;
auto roi = GetRoi(kRectWidth, kRectHeight, GetCroppedRect());
std::array<float, 16> matrix;
GetRotatedSubRectToRectTransformMatrix(roi, kRectWidth, kRectHeight,
/*flip_horizontaly=*/false, &matrix);
auto status_or_result = RunCalculator(GetCroppedRectTestInput(), matrix);
MP_ASSERT_OK(status_or_result);
EXPECT_THAT(status_or_result.value(),
EqualsProto(GetCroppedRectTestExpectedResult()));
}
TEST(LandmarkProjectionCalculatorTest, ProjectingWithScaleMatrix) {
mediapipe::NormalizedLandmarkList landmarks =
ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 10, y: 20, z: -0.5 }
landmark { x: 5, y: 6, z: 7 }
)pb");
// clang-format off
std::array<float, 16> matrix = {
10.0f, 0.0f, 0.0f, 0.0f,
0.0f, 100.0f, 0.0f, 0.0f,
0.0f, 0.0f, 1.0f, 0.0f,
0.0f, 0.0f, 0.0f, 1.0f,
};
// clang-format on
auto status_or_result =
RunCalculator(std::move(landmarks), std::move(matrix));
MP_ASSERT_OK(status_or_result);
EXPECT_THAT(
status_or_result.value(),
EqualsProto(ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 100, y: 2000, z: -5 }
landmark { x: 50, y: 600, z: 70 }
)pb")));
}
TEST(LandmarkProjectionCalculatorTest, ProjectingWithTranslateMatrix) {
mediapipe::NormalizedLandmarkList landmarks =
ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 10, y: 20, z: -0.5 }
)pb");
// clang-format off
std::array<float, 16> matrix = {
1.0f, 0.0f, 0.0f, 1.0f,
0.0f, 1.0f, 0.0f, 2.0f,
0.0f, 0.0f, 1.0f, 0.0f,
0.0f, 0.0f, 0.0f, 1.0f,
};
// clang-format on
auto status_or_result =
RunCalculator(std::move(landmarks), std::move(matrix));
MP_ASSERT_OK(status_or_result);
EXPECT_THAT(
status_or_result.value(),
EqualsProto(ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 11, y: 22, z: -0.5 }
)pb")));
}
TEST(LandmarkProjectionCalculatorTest, ProjectingWithRotationMatrix) {
mediapipe::NormalizedLandmarkList landmarks =
ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 4, y: 0, z: -0.5 }
)pb");
// clang-format off
// 90 degrees rotation matrix
std::array<float, 16> matrix = {
0.0f, -1.0f, 0.0f, 0.0f,
1.0f, 0.0f, 0.0f, 0.0f,
0.0f, 0.0f, 1.0f, 0.0f,
0.0f, 0.0f, 0.0f, 1.0f,
};
// clang-format on
auto status_or_result =
RunCalculator(std::move(landmarks), std::move(matrix));
MP_ASSERT_OK(status_or_result);
EXPECT_THAT(
status_or_result.value(),
EqualsProto(ParseTextProtoOrDie<mediapipe::NormalizedLandmarkList>(R"pb(
landmark { x: 0, y: 4, z: -0.5 }
)pb")));
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/util/rect_to_render_scale_calculator.cc
+14
View File
@@ -1,3 +1,17 @@
// Copyright 2020 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 "mediapipe/calculators/util/rect_to_render_scale_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/rect.pb.h"
mediapipe/calculators/util/rect_to_render_scale_calculator.proto
+14
View File
@@ -1,3 +1,17 @@
// Copyright 2020 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;
mediapipe/calculators/util/refine_landmarks_from_heatmap_calculator.cc
+166
View File
@@ -0,0 +1,166 @@
// 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 "mediapipe/calculators/util/refine_landmarks_from_heatmap_calculator.h"
#include "mediapipe/calculators/util/refine_landmarks_from_heatmap_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
namespace mediapipe {
namespace {
inline float Sigmoid(float value) { return 1.0f / (1.0f + std::exp(-value)); }
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 size for heatmap tensor" << dims.size();
}
}
} // namespace
namespace api2 {
// Refines landmarks using correspond heatmap area.
//
// Input:
// NORM_LANDMARKS - Required. Input normalized landmarks to update.
// TENSORS - Required. Vector of input tensors. 0th element should be heatmap.
// The rest is unused.
// Output:
// NORM_LANDMARKS - Required. Updated normalized landmarks.
class RefineLandmarksFromHeatmapCalculatorImpl
: public NodeImpl<RefineLandmarksFromHeatmapCalculator,
RefineLandmarksFromHeatmapCalculatorImpl> {
public:
absl::Status Open(CalculatorContext* cc) override { return absl::OkStatus(); }
absl::Status Process(CalculatorContext* cc) override {
// Make sure we bypass landmarks if there is no detection.
if (kInLandmarks(cc).IsEmpty()) {
return absl::OkStatus();
}
// If for some reason heatmap is missing, just return original landmarks.
if (kInTensors(cc).IsEmpty()) {
kOutLandmarks(cc).Send(*kInLandmarks(cc));
return absl::OkStatus();
}
// Check basic prerequisites.
const auto& input_tensors = *kInTensors(cc);
RET_CHECK(!input_tensors.empty()) << "Empty input tensors list. First "
"element is expeced to be a heatmap";
const auto& hm_tensor = input_tensors[0];
const auto& in_lms = *kInLandmarks(cc);
auto hm_view = hm_tensor.GetCpuReadView();
auto hm_raw = hm_view.buffer<float>();
const auto& options =
cc->Options<mediapipe::RefineLandmarksFromHeatmapCalculatorOptions>();
ASSIGN_OR_RETURN(auto out_lms, RefineLandmarksFromHeatMap(
in_lms, hm_raw, hm_tensor.shape().dims,
options.kernel_size(),
options.min_confidence_to_refine()));
kOutLandmarks(cc).Send(std::move(out_lms));
return absl::OkStatus();
}
};
} // namespace api2
// Runs actual refinement
// High level algorithm:
//
// Heatmap is accepted as tensor in HWC layout where i-th channel is a heatmap
// for the i-th landmark.
//
// For each landmark we replace original value with a value calculated from the
// area in heatmap close to original landmark position (in particular are
// covered with kernel of size options.kernel_size). To calculate new coordinate
// from heatmap we calculate an weighted average inside the kernel. We update
// the landmark iff heatmap is confident in it's prediction i.e. max(heatmap) in
// kernel is at least options.min_confidence_to_refine big.
absl::StatusOr<mediapipe::NormalizedLandmarkList> RefineLandmarksFromHeatMap(
const mediapipe::NormalizedLandmarkList& in_lms,
const float* heatmap_raw_data, const std::vector<int>& heatmap_dims,
int kernel_size, float min_confidence_to_refine) {
ASSIGN_OR_RETURN(auto hm_dims, GetHwcFromDims(heatmap_dims));
auto [hm_height, hm_width, hm_channels] = hm_dims;
RET_CHECK_EQ(in_lms.landmark_size(), hm_channels)
<< "Expected heatmap to have number of layers == to number of "
"landmarks";
int hm_row_size = hm_width * hm_channels;
int hm_pixel_size = hm_channels;
mediapipe::NormalizedLandmarkList out_lms = in_lms;
for (int lm_index = 0; lm_index < out_lms.landmark_size(); ++lm_index) {
int center_col = out_lms.landmark(lm_index).x() * hm_width;
int center_row = out_lms.landmark(lm_index).y() * hm_height;
// Point is outside of the image let's keep it intact.
if (center_col < 0 || center_col >= hm_width || center_row < 0 ||
center_col >= hm_height) {
continue;
}
int offset = (kernel_size - 1) / 2;
// Calculate area to iterate over. Note that we decrease the kernel on
// the edges of the heatmap. Equivalent to zero border.
int begin_col = std::max(0, center_col - offset);
int end_col = std::min(hm_width, center_col + offset + 1);
int begin_row = std::max(0, center_row - offset);
int end_row = std::min(hm_height, center_row + offset + 1);
float sum = 0;
float weighted_col = 0;
float weighted_row = 0;
float max_value = 0;
// Main loop. Go over kernel and calculate weighted sum of coordinates,
// sum of weights and max weights.
for (int row = begin_row; row < end_row; ++row) {
for (int col = begin_col; col < end_col; ++col) {
// We expect memory to be in HWC layout without padding.
int idx = hm_row_size * row + hm_pixel_size * col + lm_index;
// Right now we hardcode sigmoid activation as it will be wasteful to
// calculate sigmoid for each value of heatmap in the model itself. If
// we ever have other activations it should be trivial to expand via
// options.
float confidence = Sigmoid(heatmap_raw_data[idx]);
sum += confidence;
max_value = std::max(max_value, confidence);
weighted_col += col * confidence;
weighted_row += row * confidence;
}
}
if (max_value >= min_confidence_to_refine && sum > 0) {
out_lms.mutable_landmark(lm_index)->set_x(weighted_col / hm_width / sum);
out_lms.mutable_landmark(lm_index)->set_y(weighted_row / hm_height / sum);
}
}
return out_lms;
}
} // namespace mediapipe
mediapipe/calculators/util/refine_landmarks_from_heatmap_calculator.h
+50
View File
@@ -0,0 +1,50 @@
// 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.
#ifndef MEDIAPIPE_CALCULATORS_UTIL_REFINE_LANDMARKS_FROM_HEATMAP_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_UTIL_REFINE_LANDMARKS_FROM_HEATMAP_CALCULATOR_H_
#include <vector>
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/statusor.h"
namespace mediapipe {
namespace api2 {
class RefineLandmarksFromHeatmapCalculator : public NodeIntf {
public:
static constexpr Input<mediapipe::NormalizedLandmarkList> kInLandmarks{
"NORM_LANDMARKS"};
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
static constexpr Output<mediapipe::NormalizedLandmarkList> kOutLandmarks{
"NORM_LANDMARKS"};
MEDIAPIPE_NODE_INTERFACE(RefineLandmarksFromHeatmapCalculator, kInLandmarks,
kInTensors, kOutLandmarks);
};
} // namespace api2
// Exposed for testing.
absl::StatusOr<mediapipe::NormalizedLandmarkList> RefineLandmarksFromHeatMap(
const mediapipe::NormalizedLandmarkList& in_lms,
const float* heatmap_raw_data, const std::vector<int>& heatmap_dims,
int kernel_size, float min_confidence_to_refine);
} // namespace mediapipe
#endif // MEDIAPIPE_CALCULATORS_UTIL_REFINE_LANDMARKS_FROM_HEATMAP_CALCULATOR_H_
mediapipe/calculators/util/refine_landmarks_from_heatmap_calculator.proto
+27
View File
@@ -0,0 +1,27 @@
// 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 RefineLandmarksFromHeatmapCalculatorOptions {
extend mediapipe.CalculatorOptions {
optional RefineLandmarksFromHeatmapCalculatorOptions ext = 362281653;
}
optional int32 kernel_size = 1 [default = 9];
optional float min_confidence_to_refine = 2 [default = 0.5];
}
mediapipe/calculators/util/refine_landmarks_from_heatmap_calculator_test.cc
+152
View File
@@ -0,0 +1,152 @@
// 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 "mediapipe/calculators/util/refine_landmarks_from_heatmap_calculator.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
namespace {
mediapipe::NormalizedLandmarkList vec_to_lms(
const std::vector<std::pair<float, float>>& inp) {
mediapipe::NormalizedLandmarkList ret;
for (const auto& it : inp) {
auto new_lm = ret.add_landmark();
new_lm->set_x(it.first);
new_lm->set_y(it.second);
}
return ret;
}
std::vector<std::pair<float, float>> lms_to_vec(
const mediapipe::NormalizedLandmarkList& lst) {
std::vector<std::pair<float, float>> ret;
for (const auto& lm : lst.landmark()) {
ret.push_back({lm.x(), lm.y()});
}
return ret;
}
std::vector<float> CHW_to_HWC(std::vector<float> inp, int height, int width,
int depth) {
std::vector<float> ret(inp.size());
const float* inp_ptr = inp.data();
for (int c = 0; c < depth; ++c) {
for (int row = 0; row < height; ++row) {
for (int col = 0; col < width; ++col) {
int dest_idx = width * depth * row + depth * col + c;
ret[dest_idx] = *inp_ptr;
++inp_ptr;
}
}
}
return ret;
}
using testing::ElementsAre;
using testing::FloatEq;
using testing::Pair;
TEST(RefineLandmarksFromHeatmapTest, Smoke) {
float z = -10000000000000000;
// clang-format off
std::vector<float> hm = {
z, z, z,
1, z, z,
z, z, z};
// clang-format on
auto ret_or_error = RefineLandmarksFromHeatMap(vec_to_lms({{0.5, 0.5}}),
hm.data(), {3, 3, 1}, 3, 0.1);
MP_EXPECT_OK(ret_or_error);
EXPECT_THAT(lms_to_vec(*ret_or_error),
ElementsAre(Pair(FloatEq(0), FloatEq(1 / 3.))));
}
TEST(RefineLandmarksFromHeatmapTest, MultiLayer) {
float z = -10000000000000000;
// clang-format off
std::vector<float> hm = CHW_to_HWC({
z, z, z,
1, z, z,
z, z, z,
z, z, z,
1, z, z,
z, z, z,
z, z, z,
1, z, z,
z, z, z}, 3, 3, 3);
// clang-format on
auto ret_or_error = RefineLandmarksFromHeatMap(
vec_to_lms({{0.5, 0.5}, {0.5, 0.5}, {0.5, 0.5}}), hm.data(), {3, 3, 3}, 3,
0.1);
MP_EXPECT_OK(ret_or_error);
EXPECT_THAT(lms_to_vec(*ret_or_error),
ElementsAre(Pair(FloatEq(0), FloatEq(1 / 3.)),
Pair(FloatEq(0), FloatEq(1 / 3.)),
Pair(FloatEq(0), FloatEq(1 / 3.))));
}
TEST(RefineLandmarksFromHeatmapTest, KeepIfNotSure) {
float z = -10000000000000000;
// clang-format off
std::vector<float> hm = CHW_to_HWC({
z, z, z,
0, z, z,
z, z, z,
z, z, z,
0, z, z,
z, z, z,
z, z, z,
0, z, z,
z, z, z}, 3, 3, 3);
// clang-format on
auto ret_or_error = RefineLandmarksFromHeatMap(
vec_to_lms({{0.5, 0.5}, {0.5, 0.5}, {0.5, 0.5}}), hm.data(), {3, 3, 3}, 3,
0.6);
MP_EXPECT_OK(ret_or_error);
EXPECT_THAT(lms_to_vec(*ret_or_error),
ElementsAre(Pair(FloatEq(0.5), FloatEq(0.5)),
Pair(FloatEq(0.5), FloatEq(0.5)),
Pair(FloatEq(0.5), FloatEq(0.5))));
}
TEST(RefineLandmarksFromHeatmapTest, Border) {
float z = -10000000000000000;
// clang-format off
std::vector<float> hm = CHW_to_HWC({
z, z, z,
0, z, 0,
z, z, z,
z, z, z,
0, z, 0,
z, z, 0}, 3, 3, 2);
// clang-format on
auto ret_or_error = RefineLandmarksFromHeatMap(
vec_to_lms({{0.0, 0.0}, {0.9, 0.9}}), hm.data(), {3, 3, 2}, 3, 0.1);
MP_EXPECT_OK(ret_or_error);
EXPECT_THAT(lms_to_vec(*ret_or_error),
ElementsAre(Pair(FloatEq(0), FloatEq(1 / 3.)),
Pair(FloatEq(2 / 3.), FloatEq(1 / 6. + 2 / 6.))));
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/util/thresholding_calculator.cc
+1 -1
View File
@@ -101,7 +101,7 @@ absl::Status ThresholdingCalculator::Open(CalculatorContext* cc) {
}
if (cc->InputSidePackets().HasTag("THRESHOLD")) {
threshold_ = cc->InputSidePackets().Tag("THRESHOLD").Get<float>();
threshold_ = cc->InputSidePackets().Tag("THRESHOLD").Get<double>();
}
return absl::OkStatus();
}