Project import generated by Copybara.
GitOrigin-RevId: 1610e588e497817fae2d9a458093ab6a370e2972
This commit is contained in:
MediaPipe Team
@@ -64,8 +64,9 @@ std::string ToString(GateState state) {
|
||||
// ALLOW or DISALLOW can also be specified as an input side packet. The rules
|
||||
// for evaluation remain the same as above.
|
||||
//
|
||||
// ALLOW/DISALLOW inputs must be specified either using input stream or
|
||||
// via input side packet but not both.
|
||||
// ALLOW/DISALLOW inputs must be specified either using input stream or via
|
||||
// input side packet but not both. If neither is specified, the behavior is then
|
||||
// determined by the "allow" field in the calculator options.
|
||||
//
|
||||
// Intended to be used with the default input stream handler, which synchronizes
|
||||
// all data input streams with the ALLOW/DISALLOW control input stream.
|
||||
@@ -92,20 +93,22 @@ class GateCalculator : public CalculatorBase {
|
||||
cc->InputSidePackets().HasTag(kDisallowTag);
|
||||
bool input_via_stream =
|
||||
cc->Inputs().HasTag(kAllowTag) || cc->Inputs().HasTag(kDisallowTag);
|
||||
// Only one of input_side_packet or input_stream may specify ALLOW/DISALLOW
|
||||
// input.
|
||||
RET_CHECK(input_via_side_packet ^ input_via_stream);
|
||||
|
||||
// Only one of input_side_packet or input_stream may specify
|
||||
// ALLOW/DISALLOW input.
|
||||
if (input_via_side_packet) {
|
||||
RET_CHECK(!input_via_stream);
|
||||
RET_CHECK(cc->InputSidePackets().HasTag(kAllowTag) ^
|
||||
cc->InputSidePackets().HasTag(kDisallowTag));
|
||||
|
||||
if (cc->InputSidePackets().HasTag(kAllowTag)) {
|
||||
cc->InputSidePackets().Tag(kAllowTag).Set<bool>();
|
||||
cc->InputSidePackets().Tag(kAllowTag).Set<bool>().Optional();
|
||||
} else {
|
||||
cc->InputSidePackets().Tag(kDisallowTag).Set<bool>();
|
||||
cc->InputSidePackets().Tag(kDisallowTag).Set<bool>().Optional();
|
||||
}
|
||||
} else {
|
||||
}
|
||||
if (input_via_stream) {
|
||||
RET_CHECK(!input_via_side_packet);
|
||||
RET_CHECK(cc->Inputs().HasTag(kAllowTag) ^
|
||||
cc->Inputs().HasTag(kDisallowTag));
|
||||
|
||||
@@ -139,7 +142,6 @@ class GateCalculator : public CalculatorBase {
|
||||
}
|
||||
|
||||
absl::Status Open(CalculatorContext* cc) final {
|
||||
use_side_packet_for_allow_disallow_ = false;
|
||||
if (cc->InputSidePackets().HasTag(kAllowTag)) {
|
||||
use_side_packet_for_allow_disallow_ = true;
|
||||
allow_by_side_packet_decision_ =
|
||||
@@ -158,12 +160,20 @@ class GateCalculator : public CalculatorBase {
|
||||
const auto& options = cc->Options<::mediapipe::GateCalculatorOptions>();
|
||||
empty_packets_as_allow_ = options.empty_packets_as_allow();
|
||||
|
||||
if (!use_side_packet_for_allow_disallow_ &&
|
||||
!cc->Inputs().HasTag(kAllowTag) && !cc->Inputs().HasTag(kDisallowTag)) {
|
||||
use_option_for_allow_disallow_ = true;
|
||||
allow_by_option_decision_ = options.allow();
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status Process(CalculatorContext* cc) final {
|
||||
bool allow = empty_packets_as_allow_;
|
||||
if (use_side_packet_for_allow_disallow_) {
|
||||
if (use_option_for_allow_disallow_) {
|
||||
allow = allow_by_option_decision_;
|
||||
} else if (use_side_packet_for_allow_disallow_) {
|
||||
allow = allow_by_side_packet_decision_;
|
||||
} else {
|
||||
if (cc->Inputs().HasTag(kAllowTag) &&
|
||||
@@ -217,8 +227,10 @@ class GateCalculator : public CalculatorBase {
|
||||
GateState last_gate_state_ = GATE_UNINITIALIZED;
|
||||
int num_data_streams_;
|
||||
bool empty_packets_as_allow_;
|
||||
bool use_side_packet_for_allow_disallow_;
|
||||
bool use_side_packet_for_allow_disallow_ = false;
|
||||
bool allow_by_side_packet_decision_;
|
||||
bool use_option_for_allow_disallow_ = false;
|
||||
bool allow_by_option_decision_;
|
||||
};
|
||||
REGISTER_CALCULATOR(GateCalculator);
|
||||
|
||||
|
||||
@@ -29,4 +29,8 @@ message GateCalculatorOptions {
|
||||
// disallowing the corresponding packets in the data input streams. Setting
|
||||
// this option to true inverts that, allowing the data packets to go through.
|
||||
optional bool empty_packets_as_allow = 1;
|
||||
|
||||
// Whether to allow or disallow the input streams to pass when no
|
||||
// ALLOW/DISALLOW input or side input is specified.
|
||||
optional bool allow = 2 [default = false];
|
||||
}
|
||||
|
||||
@@ -113,6 +113,68 @@ TEST_F(GateCalculatorTest, InvalidInputs) {
|
||||
)")));
|
||||
}
|
||||
|
||||
TEST_F(GateCalculatorTest, AllowByALLOWOptionToTrue) {
|
||||
SetRunner(R"(
|
||||
calculator: "GateCalculator"
|
||||
input_stream: "test_input"
|
||||
output_stream: "test_output"
|
||||
options: {
|
||||
[mediapipe.GateCalculatorOptions.ext] {
|
||||
allow: true
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
constexpr int64 kTimestampValue0 = 42;
|
||||
RunTimeStep(kTimestampValue0, true);
|
||||
constexpr int64 kTimestampValue1 = 43;
|
||||
RunTimeStep(kTimestampValue1, false);
|
||||
|
||||
const std::vector<Packet>& output = runner()->Outputs().Get("", 0).packets;
|
||||
ASSERT_EQ(2, output.size());
|
||||
EXPECT_EQ(kTimestampValue0, output[0].Timestamp().Value());
|
||||
EXPECT_EQ(kTimestampValue1, output[1].Timestamp().Value());
|
||||
EXPECT_EQ(true, output[0].Get<bool>());
|
||||
EXPECT_EQ(false, output[1].Get<bool>());
|
||||
}
|
||||
|
||||
TEST_F(GateCalculatorTest, DisallowByALLOWOptionSetToFalse) {
|
||||
SetRunner(R"(
|
||||
calculator: "GateCalculator"
|
||||
input_stream: "test_input"
|
||||
output_stream: "test_output"
|
||||
options: {
|
||||
[mediapipe.GateCalculatorOptions.ext] {
|
||||
allow: false
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
constexpr int64 kTimestampValue0 = 42;
|
||||
RunTimeStep(kTimestampValue0, true);
|
||||
constexpr int64 kTimestampValue1 = 43;
|
||||
RunTimeStep(kTimestampValue1, false);
|
||||
|
||||
const std::vector<Packet>& output = runner()->Outputs().Get("", 0).packets;
|
||||
ASSERT_EQ(0, output.size());
|
||||
}
|
||||
|
||||
TEST_F(GateCalculatorTest, DisallowByALLOWOptionNotSet) {
|
||||
SetRunner(R"(
|
||||
calculator: "GateCalculator"
|
||||
input_stream: "test_input"
|
||||
output_stream: "test_output"
|
||||
)");
|
||||
|
||||
constexpr int64 kTimestampValue0 = 42;
|
||||
RunTimeStep(kTimestampValue0, true);
|
||||
constexpr int64 kTimestampValue1 = 43;
|
||||
RunTimeStep(kTimestampValue1, false);
|
||||
|
||||
const std::vector<Packet>& output = runner()->Outputs().Get("", 0).packets;
|
||||
ASSERT_EQ(0, output.size());
|
||||
}
|
||||
|
||||
TEST_F(GateCalculatorTest, AllowByALLOWSidePacketSetToTrue) {
|
||||
SetRunner(R"(
|
||||
calculator: "GateCalculator"
|
||||
|
||||
@@ -661,3 +661,138 @@ cc_test(
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "affine_transformation",
|
||||
hdrs = ["affine_transformation.h"],
|
||||
deps = ["@com_google_absl//absl/status:statusor"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "affine_transformation_runner_gl",
|
||||
srcs = ["affine_transformation_runner_gl.cc"],
|
||||
hdrs = ["affine_transformation_runner_gl.h"],
|
||||
deps = [
|
||||
":affine_transformation",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/gpu:gl_calculator_helper",
|
||||
"//mediapipe/gpu:gl_simple_shaders",
|
||||
"//mediapipe/gpu:gpu_buffer",
|
||||
"//mediapipe/gpu:gpu_origin_cc_proto",
|
||||
"//mediapipe/gpu:shader_util",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/status:statusor",
|
||||
"@eigen_archive//:eigen3",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "affine_transformation_runner_opencv",
|
||||
srcs = ["affine_transformation_runner_opencv.cc"],
|
||||
hdrs = ["affine_transformation_runner_opencv.h"],
|
||||
deps = [
|
||||
":affine_transformation",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/port:opencv_core",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/status:statusor",
|
||||
"@eigen_archive//:eigen3",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_proto_library(
|
||||
name = "warp_affine_calculator_proto",
|
||||
srcs = ["warp_affine_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_options_proto",
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/gpu:gpu_origin_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "warp_affine_calculator",
|
||||
srcs = ["warp_affine_calculator.cc"],
|
||||
hdrs = ["warp_affine_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":affine_transformation",
|
||||
":affine_transformation_runner_opencv",
|
||||
":warp_affine_calculator_cc_proto",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/status:statusor",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/api2:node",
|
||||
"//mediapipe/framework/api2:port",
|
||||
"//mediapipe/framework/formats:image",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
] + select({
|
||||
"//mediapipe/gpu:disable_gpu": [],
|
||||
"//conditions:default": [
|
||||
"//mediapipe/gpu:gl_calculator_helper",
|
||||
"//mediapipe/gpu:gpu_buffer",
|
||||
":affine_transformation_runner_gl",
|
||||
],
|
||||
}),
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "warp_affine_calculator_test",
|
||||
srcs = ["warp_affine_calculator_test.cc"],
|
||||
data = [
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/input.jpg",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/large_sub_rect.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/large_sub_rect_border_zero.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/large_sub_rect_keep_aspect.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/large_sub_rect_keep_aspect_border_zero.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/large_sub_rect_keep_aspect_with_rotation.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/large_sub_rect_keep_aspect_with_rotation_border_zero.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/medium_sub_rect_keep_aspect.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/medium_sub_rect_keep_aspect_border_zero.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/medium_sub_rect_keep_aspect_with_rotation.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/medium_sub_rect_keep_aspect_with_rotation_border_zero.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/medium_sub_rect_with_rotation.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/medium_sub_rect_with_rotation_border_zero.png",
|
||||
"//mediapipe/calculators/tensor:testdata/image_to_tensor/noop_except_range.png",
|
||||
],
|
||||
tags = ["desktop_only_test"],
|
||||
deps = [
|
||||
":affine_transformation",
|
||||
":warp_affine_calculator",
|
||||
"//mediapipe/calculators/image:image_transformation_calculator",
|
||||
"//mediapipe/calculators/tensor:image_to_tensor_converter",
|
||||
"//mediapipe/calculators/tensor:image_to_tensor_utils",
|
||||
"//mediapipe/calculators/util:from_image_calculator",
|
||||
"//mediapipe/calculators/util:to_image_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/deps:file_path",
|
||||
"//mediapipe/framework/formats:image",
|
||||
"//mediapipe/framework/formats:image_format_cc_proto",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/formats:tensor",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:opencv_core",
|
||||
"//mediapipe/framework/port:opencv_imgcodecs",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/gpu:gpu_buffer_to_image_frame_calculator",
|
||||
"//mediapipe/gpu:image_frame_to_gpu_buffer_calculator",
|
||||
"@com_google_absl//absl/flags:flag",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
// 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_IMAGE_AFFINE_TRANSFORMATION_H_
|
||||
#define MEDIAPIPE_CALCULATORS_IMAGE_AFFINE_TRANSFORMATION_H_
|
||||
|
||||
#include <array>
|
||||
|
||||
#include "absl/status/statusor.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
class AffineTransformation {
|
||||
public:
|
||||
// Pixel extrapolation method.
|
||||
// When converting image to tensor it may happen that tensor needs to read
|
||||
// pixels outside image boundaries. Border mode helps to specify how such
|
||||
// pixels will be calculated.
|
||||
enum class BorderMode { kZero, kReplicate };
|
||||
|
||||
struct Size {
|
||||
int width;
|
||||
int height;
|
||||
};
|
||||
|
||||
template <typename InputT, typename OutputT>
|
||||
class Runner {
|
||||
public:
|
||||
virtual ~Runner() = default;
|
||||
|
||||
// Transforms input into output using @matrix as following:
|
||||
// output(x, y) = input(matrix[0] * x + matrix[1] * y + matrix[3],
|
||||
// matrix[4] * x + matrix[5] * y + matrix[7])
|
||||
// where x and y ranges are defined by @output_size.
|
||||
virtual absl::StatusOr<OutputT> Run(const InputT& input,
|
||||
const std::array<float, 16>& matrix,
|
||||
const Size& output_size,
|
||||
BorderMode border_mode) = 0;
|
||||
};
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_IMAGE_AFFINE_TRANSFORMATION_H_
|
||||
@@ -0,0 +1,354 @@
|
||||
// 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/image/affine_transformation_runner_gl.h"
|
||||
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
|
||||
#include "Eigen/Core"
|
||||
#include "Eigen/Geometry"
|
||||
#include "Eigen/LU"
|
||||
#include "absl/memory/memory.h"
|
||||
#include "absl/status/status.h"
|
||||
#include "absl/status/statusor.h"
|
||||
#include "mediapipe/calculators/image/affine_transformation.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "mediapipe/gpu/gl_simple_shaders.h"
|
||||
#include "mediapipe/gpu/gpu_buffer.h"
|
||||
#include "mediapipe/gpu/gpu_origin.pb.h"
|
||||
#include "mediapipe/gpu/shader_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
using mediapipe::GlCalculatorHelper;
|
||||
using mediapipe::GlhCreateProgram;
|
||||
using mediapipe::GlTexture;
|
||||
using mediapipe::GpuBuffer;
|
||||
using mediapipe::GpuOrigin;
|
||||
|
||||
bool IsMatrixVerticalFlipNeeded(GpuOrigin::Mode gpu_origin) {
|
||||
switch (gpu_origin) {
|
||||
case GpuOrigin::DEFAULT:
|
||||
case GpuOrigin::CONVENTIONAL:
|
||||
#ifdef __APPLE__
|
||||
return false;
|
||||
#else
|
||||
return true;
|
||||
#endif // __APPLE__
|
||||
case GpuOrigin::TOP_LEFT:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef __APPLE__
|
||||
#define GL_CLAMP_TO_BORDER_MAY_BE_SUPPORTED 0
|
||||
#else
|
||||
#define GL_CLAMP_TO_BORDER_MAY_BE_SUPPORTED 1
|
||||
#endif // __APPLE__
|
||||
|
||||
bool IsGlClampToBorderSupported(const mediapipe::GlContext& gl_context) {
|
||||
return gl_context.gl_major_version() > 3 ||
|
||||
(gl_context.gl_major_version() == 3 &&
|
||||
gl_context.gl_minor_version() >= 2);
|
||||
}
|
||||
|
||||
constexpr int kAttribVertex = 0;
|
||||
constexpr int kAttribTexturePosition = 1;
|
||||
constexpr int kNumAttributes = 2;
|
||||
|
||||
class GlTextureWarpAffineRunner
|
||||
: public AffineTransformation::Runner<GpuBuffer,
|
||||
std::unique_ptr<GpuBuffer>> {
|
||||
public:
|
||||
GlTextureWarpAffineRunner(std::shared_ptr<GlCalculatorHelper> gl_helper,
|
||||
GpuOrigin::Mode gpu_origin)
|
||||
: gl_helper_(gl_helper), gpu_origin_(gpu_origin) {}
|
||||
absl::Status Init() {
|
||||
return gl_helper_->RunInGlContext([this]() -> absl::Status {
|
||||
const GLint attr_location[kNumAttributes] = {
|
||||
kAttribVertex,
|
||||
kAttribTexturePosition,
|
||||
};
|
||||
const GLchar* attr_name[kNumAttributes] = {
|
||||
"position",
|
||||
"texture_coordinate",
|
||||
};
|
||||
|
||||
constexpr GLchar kVertShader[] = R"(
|
||||
in vec4 position;
|
||||
in mediump vec4 texture_coordinate;
|
||||
out mediump vec2 sample_coordinate;
|
||||
uniform mat4 transform_matrix;
|
||||
|
||||
void main() {
|
||||
gl_Position = position;
|
||||
vec4 tc = transform_matrix * texture_coordinate;
|
||||
sample_coordinate = tc.xy;
|
||||
}
|
||||
)";
|
||||
|
||||
constexpr GLchar kFragShader[] = R"(
|
||||
DEFAULT_PRECISION(mediump, float)
|
||||
in vec2 sample_coordinate;
|
||||
uniform sampler2D input_texture;
|
||||
|
||||
#ifdef GL_ES
|
||||
#define fragColor gl_FragColor
|
||||
#else
|
||||
out vec4 fragColor;
|
||||
#endif // defined(GL_ES);
|
||||
|
||||
void main() {
|
||||
vec4 color = texture2D(input_texture, sample_coordinate);
|
||||
#ifdef CUSTOM_ZERO_BORDER_MODE
|
||||
float out_of_bounds =
|
||||
float(sample_coordinate.x < 0.0 || sample_coordinate.x > 1.0 ||
|
||||
sample_coordinate.y < 0.0 || sample_coordinate.y > 1.0);
|
||||
color = mix(color, vec4(0.0, 0.0, 0.0, 0.0), out_of_bounds);
|
||||
#endif // defined(CUSTOM_ZERO_BORDER_MODE)
|
||||
fragColor = color;
|
||||
}
|
||||
)";
|
||||
|
||||
// Create program and set parameters.
|
||||
auto create_fn = [&](const std::string& vs,
|
||||
const std::string& fs) -> absl::StatusOr<Program> {
|
||||
GLuint program = 0;
|
||||
GlhCreateProgram(vs.c_str(), fs.c_str(), kNumAttributes, &attr_name[0],
|
||||
attr_location, &program);
|
||||
|
||||
RET_CHECK(program) << "Problem initializing warp affine program.";
|
||||
glUseProgram(program);
|
||||
glUniform1i(glGetUniformLocation(program, "input_texture"), 1);
|
||||
GLint matrix_id = glGetUniformLocation(program, "transform_matrix");
|
||||
return Program{.id = program, .matrix_id = matrix_id};
|
||||
};
|
||||
|
||||
const std::string vert_src =
|
||||
absl::StrCat(mediapipe::kMediaPipeVertexShaderPreamble, kVertShader);
|
||||
|
||||
const std::string frag_src = absl::StrCat(
|
||||
mediapipe::kMediaPipeFragmentShaderPreamble, kFragShader);
|
||||
|
||||
ASSIGN_OR_RETURN(program_, create_fn(vert_src, frag_src));
|
||||
|
||||
auto create_custom_zero_fn = [&]() -> absl::StatusOr<Program> {
|
||||
std::string custom_zero_border_mode_def = R"(
|
||||
#define CUSTOM_ZERO_BORDER_MODE
|
||||
)";
|
||||
const std::string frag_custom_zero_src =
|
||||
absl::StrCat(mediapipe::kMediaPipeFragmentShaderPreamble,
|
||||
custom_zero_border_mode_def, kFragShader);
|
||||
return create_fn(vert_src, frag_custom_zero_src);
|
||||
};
|
||||
#if GL_CLAMP_TO_BORDER_MAY_BE_SUPPORTED
|
||||
if (!IsGlClampToBorderSupported(gl_helper_->GetGlContext())) {
|
||||
ASSIGN_OR_RETURN(program_custom_zero_, create_custom_zero_fn());
|
||||
}
|
||||
#else
|
||||
ASSIGN_OR_RETURN(program_custom_zero_, create_custom_zero_fn());
|
||||
#endif // GL_CLAMP_TO_BORDER_MAY_BE_SUPPORTED
|
||||
|
||||
glGenFramebuffers(1, &framebuffer_);
|
||||
|
||||
// vertex storage
|
||||
glGenBuffers(2, vbo_);
|
||||
glGenVertexArrays(1, &vao_);
|
||||
|
||||
// vbo 0
|
||||
glBindBuffer(GL_ARRAY_BUFFER, vbo_[0]);
|
||||
glBufferData(GL_ARRAY_BUFFER, sizeof(mediapipe::kBasicSquareVertices),
|
||||
mediapipe::kBasicSquareVertices, GL_STATIC_DRAW);
|
||||
|
||||
// vbo 1
|
||||
glBindBuffer(GL_ARRAY_BUFFER, vbo_[1]);
|
||||
glBufferData(GL_ARRAY_BUFFER, sizeof(mediapipe::kBasicTextureVertices),
|
||||
mediapipe::kBasicTextureVertices, GL_STATIC_DRAW);
|
||||
|
||||
glBindBuffer(GL_ARRAY_BUFFER, 0);
|
||||
|
||||
return absl::OkStatus();
|
||||
});
|
||||
}
|
||||
|
||||
absl::StatusOr<std::unique_ptr<GpuBuffer>> Run(
|
||||
const GpuBuffer& input, const std::array<float, 16>& matrix,
|
||||
const AffineTransformation::Size& size,
|
||||
AffineTransformation::BorderMode border_mode) override {
|
||||
std::unique_ptr<GpuBuffer> gpu_buffer;
|
||||
MP_RETURN_IF_ERROR(
|
||||
gl_helper_->RunInGlContext([this, &input, &matrix, &size, &border_mode,
|
||||
&gpu_buffer]() -> absl::Status {
|
||||
auto input_texture = gl_helper_->CreateSourceTexture(input);
|
||||
auto output_texture = gl_helper_->CreateDestinationTexture(
|
||||
size.width, size.height, input.format());
|
||||
|
||||
MP_RETURN_IF_ERROR(
|
||||
RunInternal(input_texture, matrix, border_mode, &output_texture));
|
||||
gpu_buffer = output_texture.GetFrame<GpuBuffer>();
|
||||
return absl::OkStatus();
|
||||
}));
|
||||
|
||||
return gpu_buffer;
|
||||
}
|
||||
|
||||
absl::Status RunInternal(const GlTexture& texture,
|
||||
const std::array<float, 16>& matrix,
|
||||
AffineTransformation::BorderMode border_mode,
|
||||
GlTexture* output) {
|
||||
glDisable(GL_DEPTH_TEST);
|
||||
glBindFramebuffer(GL_FRAMEBUFFER, framebuffer_);
|
||||
glViewport(0, 0, output->width(), output->height());
|
||||
|
||||
glActiveTexture(GL_TEXTURE0);
|
||||
glBindTexture(GL_TEXTURE_2D, output->name());
|
||||
glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D,
|
||||
output->name(), 0);
|
||||
|
||||
glActiveTexture(GL_TEXTURE1);
|
||||
glBindTexture(texture.target(), texture.name());
|
||||
|
||||
// a) Filtering.
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
|
||||
|
||||
// b) Clamping.
|
||||
std::optional<Program> program = program_;
|
||||
switch (border_mode) {
|
||||
case AffineTransformation::BorderMode::kReplicate: {
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
|
||||
break;
|
||||
}
|
||||
case AffineTransformation::BorderMode::kZero: {
|
||||
#if GL_CLAMP_TO_BORDER_MAY_BE_SUPPORTED
|
||||
if (program_custom_zero_) {
|
||||
program = program_custom_zero_;
|
||||
} else {
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_BORDER);
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_BORDER);
|
||||
glTexParameterfv(GL_TEXTURE_2D, GL_TEXTURE_BORDER_COLOR,
|
||||
std::array<float, 4>{0.0f, 0.0f, 0.0f, 0.0f}.data());
|
||||
}
|
||||
#else
|
||||
RET_CHECK(program_custom_zero_)
|
||||
<< "Program must have been initialized.";
|
||||
program = program_custom_zero_;
|
||||
#endif // GL_CLAMP_TO_BORDER_MAY_BE_SUPPORTED
|
||||
break;
|
||||
}
|
||||
}
|
||||
glUseProgram(program->id);
|
||||
|
||||
Eigen::Matrix<float, 4, 4, Eigen::RowMajor> eigen_mat(matrix.data());
|
||||
if (IsMatrixVerticalFlipNeeded(gpu_origin_)) {
|
||||
// @matrix describes affine transformation in terms of TOP LEFT origin, so
|
||||
// in some cases/on some platforms an extra flipping should be done before
|
||||
// and after.
|
||||
const Eigen::Matrix<float, 4, 4, Eigen::RowMajor> flip_y(
|
||||
{{1.0f, 0.0f, 0.0f, 0.0f},
|
||||
{0.0f, -1.0f, 0.0f, 1.0f},
|
||||
{0.0f, 0.0f, 1.0f, 0.0f},
|
||||
{0.0f, 0.0f, 0.0f, 1.0f}});
|
||||
eigen_mat = flip_y * eigen_mat * flip_y;
|
||||
}
|
||||
|
||||
// If GL context is ES2, then GL_FALSE must be used for 'transpose'
|
||||
// GLboolean in glUniformMatrix4fv, or else INVALID_VALUE error is reported.
|
||||
// Hence, transposing the matrix and always passing transposed.
|
||||
eigen_mat.transposeInPlace();
|
||||
glUniformMatrix4fv(program->matrix_id, 1, GL_FALSE, eigen_mat.data());
|
||||
|
||||
// vao
|
||||
glBindVertexArray(vao_);
|
||||
|
||||
// vbo 0
|
||||
glBindBuffer(GL_ARRAY_BUFFER, vbo_[0]);
|
||||
glEnableVertexAttribArray(kAttribVertex);
|
||||
glVertexAttribPointer(kAttribVertex, 2, GL_FLOAT, 0, 0, nullptr);
|
||||
|
||||
// vbo 1
|
||||
glBindBuffer(GL_ARRAY_BUFFER, vbo_[1]);
|
||||
glEnableVertexAttribArray(kAttribTexturePosition);
|
||||
glVertexAttribPointer(kAttribTexturePosition, 2, GL_FLOAT, 0, 0, nullptr);
|
||||
|
||||
// draw
|
||||
glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
|
||||
|
||||
// Resetting to MediaPipe texture param defaults.
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
|
||||
|
||||
glDisableVertexAttribArray(kAttribVertex);
|
||||
glDisableVertexAttribArray(kAttribTexturePosition);
|
||||
glBindBuffer(GL_ARRAY_BUFFER, 0);
|
||||
glBindVertexArray(0);
|
||||
|
||||
glActiveTexture(GL_TEXTURE1);
|
||||
glBindTexture(GL_TEXTURE_2D, 0);
|
||||
glActiveTexture(GL_TEXTURE0);
|
||||
glBindTexture(GL_TEXTURE_2D, 0);
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
~GlTextureWarpAffineRunner() override {
|
||||
gl_helper_->RunInGlContext([this]() {
|
||||
// Release OpenGL resources.
|
||||
if (framebuffer_ != 0) glDeleteFramebuffers(1, &framebuffer_);
|
||||
if (program_.id != 0) glDeleteProgram(program_.id);
|
||||
if (program_custom_zero_ && program_custom_zero_->id != 0) {
|
||||
glDeleteProgram(program_custom_zero_->id);
|
||||
}
|
||||
if (vao_ != 0) glDeleteVertexArrays(1, &vao_);
|
||||
glDeleteBuffers(2, vbo_);
|
||||
});
|
||||
}
|
||||
|
||||
private:
|
||||
struct Program {
|
||||
GLuint id;
|
||||
GLint matrix_id;
|
||||
};
|
||||
std::shared_ptr<GlCalculatorHelper> gl_helper_;
|
||||
GpuOrigin::Mode gpu_origin_;
|
||||
GLuint vao_ = 0;
|
||||
GLuint vbo_[2] = {0, 0};
|
||||
Program program_;
|
||||
std::optional<Program> program_custom_zero_;
|
||||
GLuint framebuffer_ = 0;
|
||||
};
|
||||
|
||||
#undef GL_CLAMP_TO_BORDER_MAY_BE_SUPPORTED
|
||||
|
||||
} // namespace
|
||||
|
||||
absl::StatusOr<std::unique_ptr<
|
||||
AffineTransformation::Runner<GpuBuffer, std::unique_ptr<GpuBuffer>>>>
|
||||
CreateAffineTransformationGlRunner(
|
||||
std::shared_ptr<GlCalculatorHelper> gl_helper, GpuOrigin::Mode gpu_origin) {
|
||||
auto runner =
|
||||
absl::make_unique<GlTextureWarpAffineRunner>(gl_helper, gpu_origin);
|
||||
MP_RETURN_IF_ERROR(runner->Init());
|
||||
return runner;
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,36 @@
|
||||
// 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_IMAGE_AFFINE_TRANSFORMATION_RUNNER_GL_H_
|
||||
#define MEDIAPIPE_CALCULATORS_IMAGE_AFFINE_TRANSFORMATION_RUNNER_GL_H_
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "absl/status/statusor.h"
|
||||
#include "mediapipe/calculators/image/affine_transformation.h"
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "mediapipe/gpu/gpu_buffer.h"
|
||||
#include "mediapipe/gpu/gpu_origin.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
absl::StatusOr<std::unique_ptr<AffineTransformation::Runner<
|
||||
mediapipe::GpuBuffer, std::unique_ptr<mediapipe::GpuBuffer>>>>
|
||||
CreateAffineTransformationGlRunner(
|
||||
std::shared_ptr<mediapipe::GlCalculatorHelper> gl_helper,
|
||||
mediapipe::GpuOrigin::Mode gpu_origin);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_IMAGE_AFFINE_TRANSFORMATION_RUNNER_GL_H_
|
||||
@@ -0,0 +1,160 @@
|
||||
// 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/image/affine_transformation_runner_opencv.h"
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "absl/status/statusor.h"
|
||||
#include "mediapipe/calculators/image/affine_transformation.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/port/opencv_core_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
cv::BorderTypes GetBorderModeForOpenCv(
|
||||
AffineTransformation::BorderMode border_mode) {
|
||||
switch (border_mode) {
|
||||
case AffineTransformation::BorderMode::kZero:
|
||||
return cv::BORDER_CONSTANT;
|
||||
case AffineTransformation::BorderMode::kReplicate:
|
||||
return cv::BORDER_REPLICATE;
|
||||
}
|
||||
}
|
||||
|
||||
class OpenCvRunner
|
||||
: public AffineTransformation::Runner<ImageFrame, ImageFrame> {
|
||||
public:
|
||||
absl::StatusOr<ImageFrame> Run(
|
||||
const ImageFrame& input, const std::array<float, 16>& matrix,
|
||||
const AffineTransformation::Size& size,
|
||||
AffineTransformation::BorderMode border_mode) override {
|
||||
// OpenCV warpAffine works in absolute coordinates, so the transfom (which
|
||||
// accepts and produces relative coordinates) should be adjusted to first
|
||||
// normalize coordinates and then scale them.
|
||||
// clang-format off
|
||||
cv::Matx44f normalize_dst_coordinate({
|
||||
1.0f / size.width, 0.0f, 0.0f, 0.0f,
|
||||
0.0f, 1.0f / size.height, 0.0f, 0.0f,
|
||||
0.0f, 0.0f, 1.0f, 0.0f,
|
||||
0.0f, 0.0f, 0.0f, 1.0f});
|
||||
cv::Matx44f scale_src_coordinate({
|
||||
1.0f * input.Width(), 0.0f, 0.0f, 0.0f,
|
||||
0.0f, 1.0f * input.Height(), 0.0f, 0.0f,
|
||||
0.0f, 0.0f, 1.0f, 0.0f,
|
||||
0.0f, 0.0f, 0.0f, 1.0f});
|
||||
// clang-format on
|
||||
cv::Matx44f adjust_dst_coordinate;
|
||||
cv::Matx44f adjust_src_coordinate;
|
||||
// TODO: update to always use accurate implementation.
|
||||
constexpr bool kOpenCvCompatibility = true;
|
||||
if (kOpenCvCompatibility) {
|
||||
adjust_dst_coordinate = normalize_dst_coordinate;
|
||||
adjust_src_coordinate = scale_src_coordinate;
|
||||
} else {
|
||||
// To do an accurate affine image transformation and make "on-cpu" and
|
||||
// "on-gpu" calculations aligned - extra offset is required to select
|
||||
// correct pixels.
|
||||
//
|
||||
// Each destination pixel corresponds to some pixels region from source
|
||||
// image.(In case of downscaling there can be more than one pixel.) The
|
||||
// offset for x and y is calculated in the way, so pixel in the middle of
|
||||
// the region is selected.
|
||||
//
|
||||
// For simplicity sake, let's consider downscaling from 100x50 to 10x10
|
||||
// without a rotation:
|
||||
// 1. Each destination pixel corresponds to 10x5 region
|
||||
// X range: [0, .. , 9]
|
||||
// Y range: [0, .. , 4]
|
||||
// 2. Considering we have __discrete__ pixels, the center of the region is
|
||||
// between (4, 2) and (5, 2) pixels, let's assume it's a "pixel"
|
||||
// (4.5, 2).
|
||||
// 3. When using the above as an offset for every pixel select while
|
||||
// downscaling, resulting pixels are:
|
||||
// (4.5, 2), (14.5, 2), .. , (94.5, 2)
|
||||
// (4.5, 7), (14.5, 7), .. , (94.5, 7)
|
||||
// ..
|
||||
// (4.5, 47), (14.5, 47), .., (94.5, 47)
|
||||
// instead of:
|
||||
// (0, 0), (10, 0), .. , (90, 0)
|
||||
// (0, 5), (10, 7), .. , (90, 5)
|
||||
// ..
|
||||
// (0, 45), (10, 45), .., (90, 45)
|
||||
// The latter looks shifted.
|
||||
//
|
||||
// Offsets are needed, so that __discrete__ pixel at (0, 0) corresponds to
|
||||
// the same pixel as would __non discrete__ pixel at (0.5, 0.5). Hence,
|
||||
// transformation matrix should shift coordinates by (0.5, 0.5) as the
|
||||
// very first step.
|
||||
//
|
||||
// Due to the above shift, transformed coordinates would be valid for
|
||||
// float coordinates where pixel (0, 0) spans [0.0, 1.0) x [0.0, 1.0).
|
||||
// T0 make it valid for __discrete__ pixels, transformation matrix should
|
||||
// shift coordinate by (-0.5f, -0.5f) as the very last step. (E.g. if we
|
||||
// get (0.5f, 0.5f), then it's (0, 0) __discrete__ pixel.)
|
||||
// clang-format off
|
||||
cv::Matx44f shift_dst({1.0f, 0.0f, 0.0f, 0.5f,
|
||||
0.0f, 1.0f, 0.0f, 0.5f,
|
||||
0.0f, 0.0f, 1.0f, 0.0f,
|
||||
0.0f, 0.0f, 0.0f, 1.0f});
|
||||
cv::Matx44f shift_src({1.0f, 0.0f, 0.0f, -0.5f,
|
||||
0.0f, 1.0f, 0.0f, -0.5f,
|
||||
0.0f, 0.0f, 1.0f, 0.0f,
|
||||
0.0f, 0.0f, 0.0f, 1.0f});
|
||||
// clang-format on
|
||||
adjust_dst_coordinate = normalize_dst_coordinate * shift_dst;
|
||||
adjust_src_coordinate = shift_src * scale_src_coordinate;
|
||||
}
|
||||
|
||||
cv::Matx44f transform(matrix.data());
|
||||
cv::Matx44f transform_absolute =
|
||||
adjust_src_coordinate * transform * adjust_dst_coordinate;
|
||||
|
||||
cv::Mat in_mat = formats::MatView(&input);
|
||||
|
||||
cv::Mat cv_affine_transform(2, 3, CV_32F);
|
||||
cv_affine_transform.at<float>(0, 0) = transform_absolute.val[0];
|
||||
cv_affine_transform.at<float>(0, 1) = transform_absolute.val[1];
|
||||
cv_affine_transform.at<float>(0, 2) = transform_absolute.val[3];
|
||||
cv_affine_transform.at<float>(1, 0) = transform_absolute.val[4];
|
||||
cv_affine_transform.at<float>(1, 1) = transform_absolute.val[5];
|
||||
cv_affine_transform.at<float>(1, 2) = transform_absolute.val[7];
|
||||
|
||||
ImageFrame out_image(input.Format(), size.width, size.height);
|
||||
cv::Mat out_mat = formats::MatView(&out_image);
|
||||
|
||||
cv::warpAffine(in_mat, out_mat, cv_affine_transform,
|
||||
cv::Size(out_mat.cols, out_mat.rows),
|
||||
/*flags=*/cv::INTER_LINEAR | cv::WARP_INVERSE_MAP,
|
||||
GetBorderModeForOpenCv(border_mode));
|
||||
|
||||
return out_image;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
absl::StatusOr<
|
||||
std::unique_ptr<AffineTransformation::Runner<ImageFrame, ImageFrame>>>
|
||||
CreateAffineTransformationOpenCvRunner() {
|
||||
return absl::make_unique<OpenCvRunner>();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,32 @@
|
||||
// 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_IMAGE_AFFINE_TRANSFORMATION_RUNNER_OPENCV_H_
|
||||
#define MEDIAPIPE_CALCULATORS_IMAGE_AFFINE_TRANSFORMATION_RUNNER_OPENCV_H_
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "absl/status/statusor.h"
|
||||
#include "mediapipe/calculators/image/affine_transformation.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
absl::StatusOr<
|
||||
std::unique_ptr<AffineTransformation::Runner<ImageFrame, ImageFrame>>>
|
||||
CreateAffineTransformationOpenCvRunner();
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_IMAGE_AFFINE_TRANSFORMATION_RUNNER_OPENCV_H_
|
||||
@@ -262,6 +262,7 @@ absl::Status ScaleImageCalculator::InitializeFrameInfo(CalculatorContext* cc) {
|
||||
scale_image::FindOutputDimensions(crop_width_, crop_height_, //
|
||||
options_.target_width(), //
|
||||
options_.target_height(), //
|
||||
options_.target_max_area(), //
|
||||
options_.preserve_aspect_ratio(), //
|
||||
options_.scale_to_multiple_of(), //
|
||||
&output_width_, &output_height_));
|
||||
|
||||
@@ -28,6 +28,11 @@ message ScaleImageCalculatorOptions {
|
||||
optional int32 target_width = 1;
|
||||
optional int32 target_height = 2;
|
||||
|
||||
// If set, then automatically calculates a target_width and target_height that
|
||||
// has an area below the target max area. Aspect ratio preservation cannot be
|
||||
// disabled.
|
||||
optional int32 target_max_area = 15;
|
||||
|
||||
// If true, the image is scaled up or down proportionally so that it
|
||||
// fits inside the box represented by target_width and target_height.
|
||||
// Otherwise it is scaled to fit target_width and target_height
|
||||
|
||||
@@ -92,12 +92,21 @@ absl::Status FindOutputDimensions(int input_width, //
|
||||
int input_height, //
|
||||
int target_width, //
|
||||
int target_height, //
|
||||
int target_max_area, //
|
||||
bool preserve_aspect_ratio, //
|
||||
int scale_to_multiple_of, //
|
||||
int* output_width, int* output_height) {
|
||||
CHECK(output_width);
|
||||
CHECK(output_height);
|
||||
|
||||
if (target_max_area > 0 && input_width * input_height > target_max_area) {
|
||||
preserve_aspect_ratio = true;
|
||||
target_height = static_cast<int>(sqrt(static_cast<double>(target_max_area) /
|
||||
(static_cast<double>(input_width) /
|
||||
static_cast<double>(input_height))));
|
||||
target_width = -1; // Resize width to preserve aspect ratio.
|
||||
}
|
||||
|
||||
if (preserve_aspect_ratio) {
|
||||
RET_CHECK(scale_to_multiple_of == 2)
|
||||
<< "FindOutputDimensions always outputs width and height that are "
|
||||
@@ -164,5 +173,17 @@ absl::Status FindOutputDimensions(int input_width, //
|
||||
<< "Unable to set output dimensions based on target dimensions.";
|
||||
}
|
||||
|
||||
absl::Status FindOutputDimensions(int input_width, //
|
||||
int input_height, //
|
||||
int target_width, //
|
||||
int target_height, //
|
||||
bool preserve_aspect_ratio, //
|
||||
int scale_to_multiple_of, //
|
||||
int* output_width, int* output_height) {
|
||||
return FindOutputDimensions(
|
||||
input_width, input_height, target_width, target_height, -1,
|
||||
preserve_aspect_ratio, scale_to_multiple_of, output_width, output_height);
|
||||
}
|
||||
|
||||
} // namespace scale_image
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -34,15 +34,25 @@ absl::Status FindCropDimensions(int input_width, int input_height, //
|
||||
int* crop_width, int* crop_height, //
|
||||
int* col_start, int* row_start);
|
||||
|
||||
// Given an input width and height, a target width and height, whether to
|
||||
// preserve the aspect ratio, and whether to round-down to the multiple of a
|
||||
// given number nearest to the targets, determine the output width and height.
|
||||
// If target_width or target_height is non-positive, then they will be set to
|
||||
// the input_width and input_height respectively. If scale_to_multiple_of is
|
||||
// less than 1, it will be treated like 1. The output_width and
|
||||
// output_height will be reduced as necessary to preserve_aspect_ratio if the
|
||||
// option is specified. If preserving the aspect ratio is desired, you must set
|
||||
// scale_to_multiple_of to 2.
|
||||
// Given an input width and height, a target width and height or max area,
|
||||
// whether to preserve the aspect ratio, and whether to round-down to the
|
||||
// multiple of a given number nearest to the targets, determine the output width
|
||||
// and height. If target_width or target_height is non-positive, then they will
|
||||
// be set to the input_width and input_height respectively. If target_area is
|
||||
// non-positive, then it will be ignored. If scale_to_multiple_of is less than
|
||||
// 1, it will be treated like 1. The output_width and output_height will be
|
||||
// reduced as necessary to preserve_aspect_ratio if the option is specified. If
|
||||
// preserving the aspect ratio is desired, you must set scale_to_multiple_of
|
||||
// to 2.
|
||||
absl::Status FindOutputDimensions(int input_width, int input_height, //
|
||||
int target_width,
|
||||
int target_height, //
|
||||
int target_max_area, //
|
||||
bool preserve_aspect_ratio, //
|
||||
int scale_to_multiple_of, //
|
||||
int* output_width, int* output_height);
|
||||
|
||||
// Backwards compatible helper.
|
||||
absl::Status FindOutputDimensions(int input_width, int input_height, //
|
||||
int target_width,
|
||||
int target_height, //
|
||||
|
||||
@@ -79,49 +79,49 @@ TEST(ScaleImageUtilsTest, FindOutputDimensionsPreserveRatio) {
|
||||
int output_width;
|
||||
int output_height;
|
||||
// Not scale.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, -1, true, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, -1, -1, true, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(200, output_width);
|
||||
EXPECT_EQ(100, output_height);
|
||||
// Not scale with odd input size.
|
||||
MP_ASSERT_OK(FindOutputDimensions(201, 101, -1, -1, false, 1, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(201, 101, -1, -1, -1, false, 1,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(201, output_width);
|
||||
EXPECT_EQ(101, output_height);
|
||||
// Scale down by 1/2.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, true, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, -1, true, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(100, output_width);
|
||||
EXPECT_EQ(50, output_height);
|
||||
// Scale up, doubling dimensions.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, 200, true, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, 200, -1, true, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(400, output_width);
|
||||
EXPECT_EQ(200, output_height);
|
||||
// Fits a 2:1 image into a 150 x 150 box. Output dimensions are always
|
||||
// visible by 2.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 150, true, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 150, -1, true, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(150, output_width);
|
||||
EXPECT_EQ(74, output_height);
|
||||
// Fits a 2:1 image into a 400 x 50 box.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 400, 50, true, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 400, 50, -1, true, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(100, output_width);
|
||||
EXPECT_EQ(50, output_height);
|
||||
// Scale to multiple number with odd targe size.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 101, -1, true, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 101, -1, -1, true, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(100, output_width);
|
||||
EXPECT_EQ(50, output_height);
|
||||
// Scale to multiple number with odd targe size.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 101, -1, true, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 101, -1, -1, true, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(100, output_width);
|
||||
EXPECT_EQ(50, output_height);
|
||||
// Scale to odd size.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 151, 101, false, 1, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 151, 101, -1, false, 1,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(151, output_width);
|
||||
EXPECT_EQ(101, output_height);
|
||||
}
|
||||
@@ -131,18 +131,18 @@ TEST(ScaleImageUtilsTest, FindOutputDimensionsNoAspectRatio) {
|
||||
int output_width;
|
||||
int output_height;
|
||||
// Scale width only.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, false, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, -1, false, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(100, output_width);
|
||||
EXPECT_EQ(100, output_height);
|
||||
// Scale height only.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, 200, false, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, 200, -1, false, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(200, output_width);
|
||||
EXPECT_EQ(200, output_height);
|
||||
// Scale both dimensions.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 200, false, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 200, -1, false, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(150, output_width);
|
||||
EXPECT_EQ(200, output_height);
|
||||
}
|
||||
@@ -152,41 +152,78 @@ TEST(ScaleImageUtilsTest, FindOutputDimensionsDownScaleToMultipleOf) {
|
||||
int output_width;
|
||||
int output_height;
|
||||
// Set no targets, downscale to a multiple of 8.
|
||||
MP_ASSERT_OK(FindOutputDimensions(100, 100, -1, -1, false, 8, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(100, 100, -1, -1, -1, false, 8,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(96, output_width);
|
||||
EXPECT_EQ(96, output_height);
|
||||
// Set width target, downscale to a multiple of 8.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, false, 8, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, -1, false, 8,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(96, output_width);
|
||||
EXPECT_EQ(96, output_height);
|
||||
// Set height target, downscale to a multiple of 8.
|
||||
MP_ASSERT_OK(FindOutputDimensions(201, 101, -1, 201, false, 8, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(201, 101, -1, 201, -1, false, 8,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(200, output_width);
|
||||
EXPECT_EQ(200, output_height);
|
||||
// Set both targets, downscale to a multiple of 8.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 200, false, 8, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 200, -1, false, 8,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(144, output_width);
|
||||
EXPECT_EQ(200, output_height);
|
||||
// Doesn't throw error if keep aspect is true and downscale multiple is 2.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 400, 200, true, 2, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 400, 200, -1, true, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(400, output_width);
|
||||
EXPECT_EQ(200, output_height);
|
||||
// Throws error if keep aspect is true, but downscale multiple is not 2.
|
||||
ASSERT_THAT(FindOutputDimensions(200, 100, 400, 200, true, 4, &output_width,
|
||||
&output_height),
|
||||
ASSERT_THAT(FindOutputDimensions(200, 100, 400, 200, -1, true, 4,
|
||||
&output_width, &output_height),
|
||||
testing::Not(testing::status::IsOk()));
|
||||
// Downscaling to multiple ignored if multiple is less than 2.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 401, 201, false, 1, &output_width,
|
||||
&output_height));
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 401, 201, -1, false, 1,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(401, output_width);
|
||||
EXPECT_EQ(201, output_height);
|
||||
}
|
||||
|
||||
// Tests scaling without keeping the aspect ratio fixed.
|
||||
TEST(ScaleImageUtilsTest, FindOutputDimensionsMaxArea) {
|
||||
int output_width;
|
||||
int output_height;
|
||||
// Smaller area.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, -1, 9000, false, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_NEAR(
|
||||
200 / 100,
|
||||
static_cast<double>(output_width) / static_cast<double>(output_height),
|
||||
0.1f);
|
||||
EXPECT_LE(output_width * output_height, 9000);
|
||||
// Close to original area.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, -1, 19999, false, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_NEAR(
|
||||
200.0 / 100.0,
|
||||
static_cast<double>(output_width) / static_cast<double>(output_height),
|
||||
0.1f);
|
||||
EXPECT_LE(output_width * output_height, 19999);
|
||||
// Don't scale with larger area.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, -1, 20001, false, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(200, output_width);
|
||||
EXPECT_EQ(100, output_height);
|
||||
// Don't scale with equal area.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, -1, 20000, false, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(200, output_width);
|
||||
EXPECT_EQ(100, output_height);
|
||||
// Don't scale at all.
|
||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, -1, -1, false, 2,
|
||||
&output_width, &output_height));
|
||||
EXPECT_EQ(200, output_width);
|
||||
EXPECT_EQ(100, output_height);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace scale_image
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
// 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/image/warp_affine_calculator.h"
|
||||
|
||||
#include <array>
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
|
||||
#include "mediapipe/calculators/image/affine_transformation.h"
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
#include "mediapipe/calculators/image/affine_transformation_runner_gl.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
#include "absl/status/status.h"
|
||||
#include "absl/status/statusor.h"
|
||||
#include "mediapipe/calculators/image/affine_transformation_runner_opencv.h"
|
||||
#include "mediapipe/calculators/image/warp_affine_calculator.pb.h"
|
||||
#include "mediapipe/framework/api2/node.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "mediapipe/gpu/gpu_buffer.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
AffineTransformation::BorderMode GetBorderMode(
|
||||
mediapipe::WarpAffineCalculatorOptions::BorderMode border_mode) {
|
||||
switch (border_mode) {
|
||||
case mediapipe::WarpAffineCalculatorOptions::BORDER_ZERO:
|
||||
return AffineTransformation::BorderMode::kZero;
|
||||
case mediapipe::WarpAffineCalculatorOptions::BORDER_UNSPECIFIED:
|
||||
case mediapipe::WarpAffineCalculatorOptions::BORDER_REPLICATE:
|
||||
return AffineTransformation::BorderMode::kReplicate;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename ImageT>
|
||||
class WarpAffineRunnerHolder {};
|
||||
|
||||
template <>
|
||||
class WarpAffineRunnerHolder<ImageFrame> {
|
||||
public:
|
||||
using RunnerType = AffineTransformation::Runner<ImageFrame, ImageFrame>;
|
||||
absl::Status Open(CalculatorContext* cc) { return absl::OkStatus(); }
|
||||
absl::StatusOr<RunnerType*> GetRunner() {
|
||||
if (!runner_) {
|
||||
ASSIGN_OR_RETURN(runner_, CreateAffineTransformationOpenCvRunner());
|
||||
}
|
||||
return runner_.get();
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<RunnerType> runner_;
|
||||
};
|
||||
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
template <>
|
||||
class WarpAffineRunnerHolder<mediapipe::GpuBuffer> {
|
||||
public:
|
||||
using RunnerType =
|
||||
AffineTransformation::Runner<mediapipe::GpuBuffer,
|
||||
std::unique_ptr<mediapipe::GpuBuffer>>;
|
||||
absl::Status Open(CalculatorContext* cc) {
|
||||
gpu_origin_ =
|
||||
cc->Options<mediapipe::WarpAffineCalculatorOptions>().gpu_origin();
|
||||
gl_helper_ = std::make_shared<mediapipe::GlCalculatorHelper>();
|
||||
return gl_helper_->Open(cc);
|
||||
}
|
||||
absl::StatusOr<RunnerType*> GetRunner() {
|
||||
if (!runner_) {
|
||||
ASSIGN_OR_RETURN(
|
||||
runner_, CreateAffineTransformationGlRunner(gl_helper_, gpu_origin_));
|
||||
}
|
||||
return runner_.get();
|
||||
}
|
||||
|
||||
private:
|
||||
mediapipe::GpuOrigin::Mode gpu_origin_;
|
||||
std::shared_ptr<mediapipe::GlCalculatorHelper> gl_helper_;
|
||||
std::unique_ptr<RunnerType> runner_;
|
||||
};
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
template <>
|
||||
class WarpAffineRunnerHolder<mediapipe::Image> {
|
||||
public:
|
||||
absl::Status Open(CalculatorContext* cc) { return runner_.Open(cc); }
|
||||
absl::StatusOr<
|
||||
AffineTransformation::Runner<mediapipe::Image, mediapipe::Image>*>
|
||||
GetRunner() {
|
||||
return &runner_;
|
||||
}
|
||||
|
||||
private:
|
||||
class Runner : public AffineTransformation::Runner<mediapipe::Image,
|
||||
mediapipe::Image> {
|
||||
public:
|
||||
absl::Status Open(CalculatorContext* cc) {
|
||||
MP_RETURN_IF_ERROR(cpu_holder_.Open(cc));
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
MP_RETURN_IF_ERROR(gpu_holder_.Open(cc));
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
return absl::OkStatus();
|
||||
}
|
||||
absl::StatusOr<mediapipe::Image> Run(
|
||||
const mediapipe::Image& input, const std::array<float, 16>& matrix,
|
||||
const AffineTransformation::Size& size,
|
||||
AffineTransformation::BorderMode border_mode) override {
|
||||
if (input.UsesGpu()) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
ASSIGN_OR_RETURN(auto* runner, gpu_holder_.GetRunner());
|
||||
ASSIGN_OR_RETURN(auto result, runner->Run(input.GetGpuBuffer(), matrix,
|
||||
size, border_mode));
|
||||
return mediapipe::Image(*result);
|
||||
#else
|
||||
return absl::UnavailableError("GPU support is disabled");
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
ASSIGN_OR_RETURN(auto* runner, cpu_holder_.GetRunner());
|
||||
const auto& frame_ptr = input.GetImageFrameSharedPtr();
|
||||
// Wrap image into image frame.
|
||||
const ImageFrame image_frame(frame_ptr->Format(), frame_ptr->Width(),
|
||||
frame_ptr->Height(), frame_ptr->WidthStep(),
|
||||
const_cast<uint8_t*>(frame_ptr->PixelData()),
|
||||
[](uint8* data) {});
|
||||
ASSIGN_OR_RETURN(auto result,
|
||||
runner->Run(image_frame, matrix, size, border_mode));
|
||||
return mediapipe::Image(std::make_shared<ImageFrame>(std::move(result)));
|
||||
}
|
||||
|
||||
private:
|
||||
WarpAffineRunnerHolder<ImageFrame> cpu_holder_;
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
WarpAffineRunnerHolder<mediapipe::GpuBuffer> gpu_holder_;
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
};
|
||||
|
||||
Runner runner_;
|
||||
};
|
||||
|
||||
template <typename InterfaceT>
|
||||
class WarpAffineCalculatorImpl : public mediapipe::api2::NodeImpl<InterfaceT> {
|
||||
public:
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
static absl::Status UpdateContract(CalculatorContract* cc) {
|
||||
if constexpr (std::is_same_v<InterfaceT, WarpAffineCalculatorGpu> ||
|
||||
std::is_same_v<InterfaceT, WarpAffineCalculator>) {
|
||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
absl::Status Open(CalculatorContext* cc) override { return holder_.Open(cc); }
|
||||
|
||||
absl::Status Process(CalculatorContext* cc) override {
|
||||
if (InterfaceT::kInImage(cc).IsEmpty() ||
|
||||
InterfaceT::kMatrix(cc).IsEmpty() ||
|
||||
InterfaceT::kOutputSize(cc).IsEmpty()) {
|
||||
return absl::OkStatus();
|
||||
}
|
||||
const std::array<float, 16>& transform = *InterfaceT::kMatrix(cc);
|
||||
auto [out_width, out_height] = *InterfaceT::kOutputSize(cc);
|
||||
AffineTransformation::Size output_size;
|
||||
output_size.width = out_width;
|
||||
output_size.height = out_height;
|
||||
ASSIGN_OR_RETURN(auto* runner, holder_.GetRunner());
|
||||
ASSIGN_OR_RETURN(
|
||||
auto result,
|
||||
runner->Run(
|
||||
*InterfaceT::kInImage(cc), transform, output_size,
|
||||
GetBorderMode(cc->Options<mediapipe::WarpAffineCalculatorOptions>()
|
||||
.border_mode())));
|
||||
InterfaceT::kOutImage(cc).Send(std::move(result));
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
WarpAffineRunnerHolder<typename decltype(InterfaceT::kInImage)::PayloadT>
|
||||
holder_;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
MEDIAPIPE_NODE_IMPLEMENTATION(
|
||||
WarpAffineCalculatorImpl<WarpAffineCalculatorCpu>);
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
MEDIAPIPE_NODE_IMPLEMENTATION(
|
||||
WarpAffineCalculatorImpl<WarpAffineCalculatorGpu>);
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
MEDIAPIPE_NODE_IMPLEMENTATION(WarpAffineCalculatorImpl<WarpAffineCalculator>);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,94 @@
|
||||
// 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_IMAGE_WARP_AFFINE_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_IMAGE_WARP_AFFINE_CALCULATOR_H_
|
||||
|
||||
#include "mediapipe/framework/api2/node.h"
|
||||
#include "mediapipe/framework/api2/port.h"
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
#include "mediapipe/gpu/gpu_buffer.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Runs affine transformation.
|
||||
//
|
||||
// Input:
|
||||
// IMAGE - Image/ImageFrame/GpuBuffer
|
||||
//
|
||||
// MATRIX - std::array<float, 16>
|
||||
// Used as following:
|
||||
// output(x, y) = input(matrix[0] * x + matrix[1] * y + matrix[3],
|
||||
// matrix[4] * x + matrix[5] * y + matrix[7])
|
||||
// where x and y ranges are defined by @OUTPUT_SIZE.
|
||||
//
|
||||
// OUTPUT_SIZE - std::pair<int, int>
|
||||
// Size of the output image.
|
||||
//
|
||||
// Output:
|
||||
// IMAGE - Image/ImageFrame/GpuBuffer
|
||||
//
|
||||
// Note:
|
||||
// - Output image type and format are the same as the input one.
|
||||
//
|
||||
// Usage example:
|
||||
// node {
|
||||
// calculator: "WarpAffineCalculator(Cpu|Gpu)"
|
||||
// input_stream: "IMAGE:image"
|
||||
// input_stream: "MATRIX:matrix"
|
||||
// input_stream: "OUTPUT_SIZE:size"
|
||||
// output_stream: "IMAGE:transformed_image"
|
||||
// options: {
|
||||
// [mediapipe.WarpAffineCalculatorOptions.ext] {
|
||||
// border_mode: BORDER_ZERO
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
template <typename ImageT>
|
||||
class WarpAffineCalculatorIntf : public mediapipe::api2::NodeIntf {
|
||||
public:
|
||||
static constexpr mediapipe::api2::Input<ImageT> kInImage{"IMAGE"};
|
||||
static constexpr mediapipe::api2::Input<std::array<float, 16>> kMatrix{
|
||||
"MATRIX"};
|
||||
static constexpr mediapipe::api2::Input<std::pair<int, int>> kOutputSize{
|
||||
"OUTPUT_SIZE"};
|
||||
static constexpr mediapipe::api2::Output<ImageT> kOutImage{"IMAGE"};
|
||||
};
|
||||
|
||||
class WarpAffineCalculatorCpu : public WarpAffineCalculatorIntf<ImageFrame> {
|
||||
public:
|
||||
MEDIAPIPE_NODE_INTERFACE(WarpAffineCalculatorCpu, kInImage, kMatrix,
|
||||
kOutputSize, kOutImage);
|
||||
};
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
class WarpAffineCalculatorGpu
|
||||
: public WarpAffineCalculatorIntf<mediapipe::GpuBuffer> {
|
||||
public:
|
||||
MEDIAPIPE_NODE_INTERFACE(WarpAffineCalculatorGpu, kInImage, kMatrix,
|
||||
kOutputSize, kOutImage);
|
||||
};
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
class WarpAffineCalculator : public WarpAffineCalculatorIntf<mediapipe::Image> {
|
||||
public:
|
||||
MEDIAPIPE_NODE_INTERFACE(WarpAffineCalculator, kInImage, kMatrix, kOutputSize,
|
||||
kOutImage);
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_IMAGE_WARP_AFFINE_CALCULATOR_H_
|
||||
@@ -0,0 +1,46 @@
|
||||
// Copyright 2021 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
import "mediapipe/gpu/gpu_origin.proto";
|
||||
|
||||
message WarpAffineCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional WarpAffineCalculatorOptions ext = 373693895;
|
||||
}
|
||||
|
||||
// Pixel extrapolation methods. See @border_mode.
|
||||
enum BorderMode {
|
||||
BORDER_UNSPECIFIED = 0;
|
||||
BORDER_ZERO = 1;
|
||||
BORDER_REPLICATE = 2;
|
||||
}
|
||||
|
||||
// Pixel extrapolation method.
|
||||
// When converting image to tensor it may happen that tensor needs to read
|
||||
// pixels outside image boundaries. Border mode helps to specify how such
|
||||
// pixels will be calculated.
|
||||
//
|
||||
// BORDER_REPLICATE is used by default.
|
||||
optional BorderMode border_mode = 1;
|
||||
|
||||
// For CONVENTIONAL mode for OpenGL, input image starts at bottom and needs
|
||||
// to be flipped vertically as tensors are expected to start at top.
|
||||
// (DEFAULT or unset interpreted as CONVENTIONAL.)
|
||||
optional GpuOrigin.Mode gpu_origin = 2;
|
||||
}
|
||||
@@ -0,0 +1,615 @@
|
||||
// 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 <cmath>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/flags/flag.h"
|
||||
#include "absl/memory/memory.h"
|
||||
#include "absl/strings/substitute.h"
|
||||
#include "mediapipe/calculators/image/affine_transformation.h"
|
||||
#include "mediapipe/calculators/tensor/image_to_tensor_converter.h"
|
||||
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/deps/file_path.h"
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
#include "mediapipe/framework/formats/image_format.pb.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/formats/tensor.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/opencv_core_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_imgcodecs_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
cv::Mat GetRgb(absl::string_view path) {
|
||||
cv::Mat bgr = cv::imread(file::JoinPath("./", path));
|
||||
cv::Mat rgb(bgr.rows, bgr.cols, CV_8UC3);
|
||||
int from_to[] = {0, 2, 1, 1, 2, 0};
|
||||
cv::mixChannels(&bgr, 1, &rgb, 1, from_to, 3);
|
||||
return rgb;
|
||||
}
|
||||
|
||||
cv::Mat GetRgba(absl::string_view path) {
|
||||
cv::Mat bgr = cv::imread(file::JoinPath("./", path));
|
||||
cv::Mat rgba(bgr.rows, bgr.cols, CV_8UC4, cv::Scalar(0, 0, 0, 0));
|
||||
int from_to[] = {0, 2, 1, 1, 2, 0};
|
||||
cv::mixChannels(&bgr, 1, &bgr, 1, from_to, 3);
|
||||
return bgr;
|
||||
}
|
||||
|
||||
// Test template.
|
||||
// No processing/assertions should be done after the function is invoked.
|
||||
void RunTest(const std::string& graph_text, const std::string& tag,
|
||||
const cv::Mat& input, cv::Mat expected_result,
|
||||
float similarity_threshold, std::array<float, 16> matrix,
|
||||
int out_width, int out_height,
|
||||
absl::optional<AffineTransformation::BorderMode> border_mode) {
|
||||
std::string border_mode_str;
|
||||
if (border_mode) {
|
||||
switch (*border_mode) {
|
||||
case AffineTransformation::BorderMode::kReplicate:
|
||||
border_mode_str = "border_mode: BORDER_REPLICATE";
|
||||
break;
|
||||
case AffineTransformation::BorderMode::kZero:
|
||||
border_mode_str = "border_mode: BORDER_ZERO";
|
||||
break;
|
||||
}
|
||||
}
|
||||
auto graph_config = mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
absl::Substitute(graph_text, /*$0=*/border_mode_str));
|
||||
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("output_image", &graph_config, &output_packets);
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
ImageFrame input_image(
|
||||
input.channels() == 4 ? ImageFormat::SRGBA : ImageFormat::SRGB,
|
||||
input.cols, input.rows, input.step, input.data, [](uint8*) {});
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_image",
|
||||
MakePacket<ImageFrame>(std::move(input_image)).At(Timestamp(0))));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"matrix",
|
||||
MakePacket<std::array<float, 16>>(std::move(matrix)).At(Timestamp(0))));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"output_size", MakePacket<std::pair<int, int>>(
|
||||
std::pair<int, int>(out_width, out_height))
|
||||
.At(Timestamp(0))));
|
||||
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
ASSERT_THAT(output_packets, testing::SizeIs(1));
|
||||
|
||||
// Get and process results.
|
||||
const ImageFrame& out_frame = output_packets[0].Get<ImageFrame>();
|
||||
cv::Mat result = formats::MatView(&out_frame);
|
||||
double similarity =
|
||||
1.0 - cv::norm(result, expected_result, cv::NORM_RELATIVE | cv::NORM_L2);
|
||||
EXPECT_GE(similarity, similarity_threshold);
|
||||
|
||||
// Fully close graph at end, otherwise calculator+tensors are destroyed
|
||||
// after calling WaitUntilDone().
|
||||
MP_ASSERT_OK(graph.CloseInputStream("input_image"));
|
||||
MP_ASSERT_OK(graph.CloseInputStream("matrix"));
|
||||
MP_ASSERT_OK(graph.CloseInputStream("output_size"));
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
enum class InputType { kImageFrame, kImage };
|
||||
|
||||
// Similarity is checked against OpenCV results always, and due to differences
|
||||
// on how OpenCV and GL treats pixels there are two thresholds.
|
||||
// TODO: update to have just one threshold when OpenCV
|
||||
// implementation is updated.
|
||||
struct SimilarityConfig {
|
||||
double threshold_on_cpu;
|
||||
double threshold_on_gpu;
|
||||
};
|
||||
|
||||
void RunTest(cv::Mat input, cv::Mat expected_result,
|
||||
const SimilarityConfig& similarity, std::array<float, 16> matrix,
|
||||
int out_width, int out_height,
|
||||
absl::optional<AffineTransformation::BorderMode> border_mode) {
|
||||
RunTest(R"(
|
||||
input_stream: "input_image"
|
||||
input_stream: "output_size"
|
||||
input_stream: "matrix"
|
||||
node {
|
||||
calculator: "WarpAffineCalculatorCpu"
|
||||
input_stream: "IMAGE:input_image"
|
||||
input_stream: "MATRIX:matrix"
|
||||
input_stream: "OUTPUT_SIZE:output_size"
|
||||
output_stream: "IMAGE:output_image"
|
||||
options {
|
||||
[mediapipe.WarpAffineCalculatorOptions.ext] {
|
||||
$0 # border mode
|
||||
}
|
||||
}
|
||||
}
|
||||
)",
|
||||
"cpu", input, expected_result, similarity.threshold_on_cpu, matrix,
|
||||
out_width, out_height, border_mode);
|
||||
|
||||
RunTest(R"(
|
||||
input_stream: "input_image"
|
||||
input_stream: "output_size"
|
||||
input_stream: "matrix"
|
||||
node {
|
||||
calculator: "ToImageCalculator"
|
||||
input_stream: "IMAGE_CPU:input_image"
|
||||
output_stream: "IMAGE:input_image_unified"
|
||||
}
|
||||
node {
|
||||
calculator: "WarpAffineCalculator"
|
||||
input_stream: "IMAGE:input_image_unified"
|
||||
input_stream: "MATRIX:matrix"
|
||||
input_stream: "OUTPUT_SIZE:output_size"
|
||||
output_stream: "IMAGE:output_image_unified"
|
||||
options {
|
||||
[mediapipe.WarpAffineCalculatorOptions.ext] {
|
||||
$0 # border mode
|
||||
}
|
||||
}
|
||||
}
|
||||
node {
|
||||
calculator: "FromImageCalculator"
|
||||
input_stream: "IMAGE:output_image_unified"
|
||||
output_stream: "IMAGE_CPU:output_image"
|
||||
}
|
||||
)",
|
||||
"cpu_image", input, expected_result, similarity.threshold_on_cpu,
|
||||
matrix, out_width, out_height, border_mode);
|
||||
|
||||
RunTest(R"(
|
||||
input_stream: "input_image"
|
||||
input_stream: "output_size"
|
||||
input_stream: "matrix"
|
||||
node {
|
||||
calculator: "ImageFrameToGpuBufferCalculator"
|
||||
input_stream: "input_image"
|
||||
output_stream: "input_image_gpu"
|
||||
}
|
||||
node {
|
||||
calculator: "WarpAffineCalculatorGpu"
|
||||
input_stream: "IMAGE:input_image_gpu"
|
||||
input_stream: "MATRIX:matrix"
|
||||
input_stream: "OUTPUT_SIZE:output_size"
|
||||
output_stream: "IMAGE:output_image_gpu"
|
||||
options {
|
||||
[mediapipe.WarpAffineCalculatorOptions.ext] {
|
||||
$0 # border mode
|
||||
gpu_origin: TOP_LEFT
|
||||
}
|
||||
}
|
||||
}
|
||||
node {
|
||||
calculator: "GpuBufferToImageFrameCalculator"
|
||||
input_stream: "output_image_gpu"
|
||||
output_stream: "output_image"
|
||||
}
|
||||
)",
|
||||
"gpu", input, expected_result, similarity.threshold_on_gpu, matrix,
|
||||
out_width, out_height, border_mode);
|
||||
|
||||
RunTest(R"(
|
||||
input_stream: "input_image"
|
||||
input_stream: "output_size"
|
||||
input_stream: "matrix"
|
||||
node {
|
||||
calculator: "ImageFrameToGpuBufferCalculator"
|
||||
input_stream: "input_image"
|
||||
output_stream: "input_image_gpu"
|
||||
}
|
||||
node {
|
||||
calculator: "ToImageCalculator"
|
||||
input_stream: "IMAGE_GPU:input_image_gpu"
|
||||
output_stream: "IMAGE:input_image_unified"
|
||||
}
|
||||
node {
|
||||
calculator: "WarpAffineCalculator"
|
||||
input_stream: "IMAGE:input_image_unified"
|
||||
input_stream: "MATRIX:matrix"
|
||||
input_stream: "OUTPUT_SIZE:output_size"
|
||||
output_stream: "IMAGE:output_image_unified"
|
||||
options {
|
||||
[mediapipe.WarpAffineCalculatorOptions.ext] {
|
||||
$0 # border mode
|
||||
gpu_origin: TOP_LEFT
|
||||
}
|
||||
}
|
||||
}
|
||||
node {
|
||||
calculator: "FromImageCalculator"
|
||||
input_stream: "IMAGE:output_image_unified"
|
||||
output_stream: "IMAGE_GPU:output_image_gpu"
|
||||
}
|
||||
node {
|
||||
calculator: "GpuBufferToImageFrameCalculator"
|
||||
input_stream: "output_image_gpu"
|
||||
output_stream: "output_image"
|
||||
}
|
||||
)",
|
||||
"gpu_image", input, expected_result, similarity.threshold_on_gpu,
|
||||
matrix, out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
std::array<float, 16> GetMatrix(cv::Mat input, mediapipe::NormalizedRect roi,
|
||||
bool keep_aspect_ratio, int out_width,
|
||||
int out_height) {
|
||||
std::array<float, 16> transform_mat;
|
||||
mediapipe::RotatedRect roi_absolute =
|
||||
mediapipe::GetRoi(input.cols, input.rows, roi);
|
||||
mediapipe::PadRoi(out_width, out_height, keep_aspect_ratio, &roi_absolute)
|
||||
.IgnoreError();
|
||||
mediapipe::GetRotatedSubRectToRectTransformMatrix(
|
||||
roi_absolute, input.cols, input.rows,
|
||||
/*flip_horizontaly=*/false, &transform_mat);
|
||||
return transform_mat;
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, MediumSubRectKeepAspect) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.65f);
|
||||
roi.set_y_center(0.4f);
|
||||
roi.set_width(0.5f);
|
||||
roi.set_height(0.5f);
|
||||
roi.set_rotation(0);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/medium_sub_rect_keep_aspect.png");
|
||||
int out_width = 256;
|
||||
int out_height = 256;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode = {};
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.82},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, MediumSubRectKeepAspectBorderZero) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.65f);
|
||||
roi.set_y_center(0.4f);
|
||||
roi.set_width(0.5f);
|
||||
roi.set_height(0.5f);
|
||||
roi.set_rotation(0);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/"
|
||||
"medium_sub_rect_keep_aspect_border_zero.png");
|
||||
int out_width = 256;
|
||||
int out_height = 256;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kZero;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.81},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, MediumSubRectKeepAspectWithRotation) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.65f);
|
||||
roi.set_y_center(0.4f);
|
||||
roi.set_width(0.5f);
|
||||
roi.set_height(0.5f);
|
||||
roi.set_rotation(M_PI * 90.0f / 180.0f);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/"
|
||||
"medium_sub_rect_keep_aspect_with_rotation.png");
|
||||
int out_width = 256;
|
||||
int out_height = 256;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kReplicate;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.77},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, MediumSubRectKeepAspectWithRotationBorderZero) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.65f);
|
||||
roi.set_y_center(0.4f);
|
||||
roi.set_width(0.5f);
|
||||
roi.set_height(0.5f);
|
||||
roi.set_rotation(M_PI * 90.0f / 180.0f);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/"
|
||||
"medium_sub_rect_keep_aspect_with_rotation_border_zero.png");
|
||||
int out_width = 256;
|
||||
int out_height = 256;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kZero;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.75},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, MediumSubRectWithRotation) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.65f);
|
||||
roi.set_y_center(0.4f);
|
||||
roi.set_width(0.5f);
|
||||
roi.set_height(0.5f);
|
||||
roi.set_rotation(M_PI * -45.0f / 180.0f);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/medium_sub_rect_with_rotation.png");
|
||||
int out_width = 256;
|
||||
int out_height = 256;
|
||||
bool keep_aspect_ratio = false;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kReplicate;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.81},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, MediumSubRectWithRotationBorderZero) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.65f);
|
||||
roi.set_y_center(0.4f);
|
||||
roi.set_width(0.5f);
|
||||
roi.set_height(0.5f);
|
||||
roi.set_rotation(M_PI * -45.0f / 180.0f);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/"
|
||||
"medium_sub_rect_with_rotation_border_zero.png");
|
||||
int out_width = 256;
|
||||
int out_height = 256;
|
||||
bool keep_aspect_ratio = false;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kZero;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.80},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, LargeSubRect) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.5f);
|
||||
roi.set_y_center(0.5f);
|
||||
roi.set_width(1.5f);
|
||||
roi.set_height(1.1f);
|
||||
roi.set_rotation(0);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/large_sub_rect.png");
|
||||
int out_width = 128;
|
||||
int out_height = 128;
|
||||
bool keep_aspect_ratio = false;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kReplicate;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.95},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, LargeSubRectBorderZero) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.5f);
|
||||
roi.set_y_center(0.5f);
|
||||
roi.set_width(1.5f);
|
||||
roi.set_height(1.1f);
|
||||
roi.set_rotation(0);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/large_sub_rect_border_zero.png");
|
||||
int out_width = 128;
|
||||
int out_height = 128;
|
||||
bool keep_aspect_ratio = false;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kZero;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.92},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, LargeSubRectKeepAspect) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.5f);
|
||||
roi.set_y_center(0.5f);
|
||||
roi.set_width(1.5f);
|
||||
roi.set_height(1.1f);
|
||||
roi.set_rotation(0);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/large_sub_rect_keep_aspect.png");
|
||||
int out_width = 128;
|
||||
int out_height = 128;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kReplicate;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.97},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, LargeSubRectKeepAspectBorderZero) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.5f);
|
||||
roi.set_y_center(0.5f);
|
||||
roi.set_width(1.5f);
|
||||
roi.set_height(1.1f);
|
||||
roi.set_rotation(0);
|
||||
auto input = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgb(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/"
|
||||
"large_sub_rect_keep_aspect_border_zero.png");
|
||||
int out_width = 128;
|
||||
int out_height = 128;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kZero;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.97},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, LargeSubRectKeepAspectWithRotation) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.5f);
|
||||
roi.set_y_center(0.5f);
|
||||
roi.set_width(1.5f);
|
||||
roi.set_height(1.1f);
|
||||
roi.set_rotation(M_PI * -15.0f / 180.0f);
|
||||
auto input = GetRgba(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgba(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/"
|
||||
"large_sub_rect_keep_aspect_with_rotation.png");
|
||||
int out_width = 128;
|
||||
int out_height = 128;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode = {};
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.91},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, LargeSubRectKeepAspectWithRotationBorderZero) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.5f);
|
||||
roi.set_y_center(0.5f);
|
||||
roi.set_width(1.5f);
|
||||
roi.set_height(1.1f);
|
||||
roi.set_rotation(M_PI * -15.0f / 180.0f);
|
||||
auto input = GetRgba(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgba(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/"
|
||||
"large_sub_rect_keep_aspect_with_rotation_border_zero.png");
|
||||
int out_width = 128;
|
||||
int out_height = 128;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kZero;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.88},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, NoOp) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.5f);
|
||||
roi.set_y_center(0.5f);
|
||||
roi.set_width(1.0f);
|
||||
roi.set_height(1.0f);
|
||||
roi.set_rotation(0);
|
||||
auto input = GetRgba(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgba(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/noop_except_range.png");
|
||||
int out_width = 64;
|
||||
int out_height = 128;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kReplicate;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.99},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
TEST(WarpAffineCalculatorTest, NoOpBorderZero) {
|
||||
mediapipe::NormalizedRect roi;
|
||||
roi.set_x_center(0.5f);
|
||||
roi.set_y_center(0.5f);
|
||||
roi.set_width(1.0f);
|
||||
roi.set_height(1.0f);
|
||||
roi.set_rotation(0);
|
||||
auto input = GetRgba(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/input.jpg");
|
||||
auto expected_output = GetRgba(
|
||||
"/mediapipe/calculators/"
|
||||
"tensor/testdata/image_to_tensor/noop_except_range.png");
|
||||
int out_width = 64;
|
||||
int out_height = 128;
|
||||
bool keep_aspect_ratio = true;
|
||||
std::optional<AffineTransformation::BorderMode> border_mode =
|
||||
AffineTransformation::BorderMode::kZero;
|
||||
RunTest(input, expected_output,
|
||||
{.threshold_on_cpu = 0.99, .threshold_on_gpu = 0.99},
|
||||
GetMatrix(input, roi, keep_aspect_ratio, out_width, out_height),
|
||||
out_width, out_height, border_mode);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -26,6 +26,11 @@ licenses(["notice"])
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
exports_files(
|
||||
glob(["testdata/image_to_tensor/*"]),
|
||||
visibility = ["//mediapipe/calculators/image:__subpackages__"],
|
||||
)
|
||||
|
||||
selects.config_setting_group(
|
||||
name = "compute_shader_unavailable",
|
||||
match_any = [
|
||||
|
||||
@@ -87,9 +87,9 @@ using GpuBuffer = mediapipe::GpuBuffer;
|
||||
// TENSORS - std::vector<Tensor>
|
||||
// Vector containing a single Tensor populated with an extrated RGB image.
|
||||
// MATRIX - std::array<float, 16> @Optional
|
||||
// An std::array<float, 16> representing a 4x4 row-major-order matrix which
|
||||
// can be used to map a point on the output tensor to a point on the input
|
||||
// image.
|
||||
// An std::array<float, 16> representing a 4x4 row-major-order matrix that
|
||||
// maps a point on the input image to a point on the output tensor, and
|
||||
// can be used to reverse the mapping by inverting the matrix.
|
||||
// LETTERBOX_PADDING - std::array<float, 4> @Optional
|
||||
// An std::array<float, 4> representing the letterbox padding from the 4
|
||||
// sides ([left, top, right, bottom]) of the output image, normalized to
|
||||
|
||||
@@ -517,8 +517,8 @@ absl::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
|
||||
uniform sampler2D frame;
|
||||
|
||||
void main() {
|
||||
$1 // flip
|
||||
vec4 pixel = texture2D(frame, sample_coordinate);
|
||||
vec2 coord = $1
|
||||
vec4 pixel = texture2D(frame, coord);
|
||||
$2 // normalize [-1,1]
|
||||
fragColor.r = pixel.r; // r channel
|
||||
$3 // g & b channels
|
||||
@@ -526,8 +526,9 @@ absl::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
|
||||
})",
|
||||
/*$0=*/single_channel ? "vec1" : "vec4",
|
||||
/*$1=*/
|
||||
flip_vertically_ ? "sample_coordinate.y = 1.0 - sample_coordinate.y;"
|
||||
: "",
|
||||
flip_vertically_
|
||||
? "vec2(sample_coordinate.x, 1.0 - sample_coordinate.y);"
|
||||
: "sample_coordinate;",
|
||||
/*$2=*/output_range_.has_value()
|
||||
? absl::Substitute("pixel = pixel * float($0) + float($1);",
|
||||
(output_range_->second - output_range_->first),
|
||||
|
||||
@@ -587,9 +587,21 @@ cc_library(
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
] + select({
|
||||
"//conditions:default": [
|
||||
"//mediapipe/framework/port:file_helpers",
|
||||
],
|
||||
}),
|
||||
"//mediapipe:android": [],
|
||||
}) + select(
|
||||
{
|
||||
"//conditions:default": [
|
||||
],
|
||||
},
|
||||
) + select(
|
||||
{
|
||||
"//conditions:default": [
|
||||
],
|
||||
"//mediapipe:android": [
|
||||
],
|
||||
},
|
||||
),
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
|
||||
@@ -37,6 +37,7 @@ const char kSequenceExampleTag[] = "SEQUENCE_EXAMPLE";
|
||||
const char kImageTag[] = "IMAGE";
|
||||
const char kFloatContextFeaturePrefixTag[] = "FLOAT_CONTEXT_FEATURE_";
|
||||
const char kFloatFeaturePrefixTag[] = "FLOAT_FEATURE_";
|
||||
const char kBytesFeaturePrefixTag[] = "BYTES_FEATURE_";
|
||||
const char kForwardFlowEncodedTag[] = "FORWARD_FLOW_ENCODED";
|
||||
const char kBBoxTag[] = "BBOX";
|
||||
const char kKeypointsTag[] = "KEYPOINTS";
|
||||
@@ -153,6 +154,9 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
||||
if (absl::StartsWith(tag, kFloatFeaturePrefixTag)) {
|
||||
cc->Inputs().Tag(tag).Set<std::vector<float>>();
|
||||
}
|
||||
if (absl::StartsWith(tag, kBytesFeaturePrefixTag)) {
|
||||
cc->Inputs().Tag(tag).Set<std::vector<std::string>>();
|
||||
}
|
||||
}
|
||||
|
||||
CHECK(cc->Outputs().HasTag(kSequenceExampleTag) ||
|
||||
@@ -231,6 +235,13 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
||||
mpms::ClearFeatureFloats(key, sequence_.get());
|
||||
mpms::ClearFeatureTimestamp(key, sequence_.get());
|
||||
}
|
||||
if (absl::StartsWith(tag, kBytesFeaturePrefixTag)) {
|
||||
std::string key = tag.substr(sizeof(kBytesFeaturePrefixTag) /
|
||||
sizeof(*kBytesFeaturePrefixTag) -
|
||||
1);
|
||||
mpms::ClearFeatureBytes(key, sequence_.get());
|
||||
mpms::ClearFeatureTimestamp(key, sequence_.get());
|
||||
}
|
||||
if (absl::StartsWith(tag, kKeypointsTag)) {
|
||||
std::string key =
|
||||
tag.substr(sizeof(kKeypointsTag) / sizeof(*kKeypointsTag) - 1);
|
||||
@@ -405,6 +416,17 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
||||
cc->Inputs().Tag(tag).Get<std::vector<float>>(),
|
||||
sequence_.get());
|
||||
}
|
||||
if (absl::StartsWith(tag, kBytesFeaturePrefixTag) &&
|
||||
!cc->Inputs().Tag(tag).IsEmpty()) {
|
||||
std::string key = tag.substr(sizeof(kBytesFeaturePrefixTag) /
|
||||
sizeof(*kBytesFeaturePrefixTag) -
|
||||
1);
|
||||
mpms::AddFeatureTimestamp(key, cc->InputTimestamp().Value(),
|
||||
sequence_.get());
|
||||
mpms::AddFeatureBytes(
|
||||
key, cc->Inputs().Tag(tag).Get<std::vector<std::string>>(),
|
||||
sequence_.get());
|
||||
}
|
||||
if (absl::StartsWith(tag, kBBoxTag) && !cc->Inputs().Tag(tag).IsEmpty()) {
|
||||
std::string key = "";
|
||||
if (tag != kBBoxTag) {
|
||||
|
||||
@@ -49,6 +49,8 @@ constexpr char kKeypointsTestTag[] = "KEYPOINTS_TEST";
|
||||
constexpr char kBboxPredictedTag[] = "BBOX_PREDICTED";
|
||||
constexpr char kAudioOtherTag[] = "AUDIO_OTHER";
|
||||
constexpr char kAudioTestTag[] = "AUDIO_TEST";
|
||||
constexpr char kBytesFeatureOtherTag[] = "BYTES_FEATURE_OTHER";
|
||||
constexpr char kBytesFeatureTestTag[] = "BYTES_FEATURE_TEST";
|
||||
constexpr char kForwardFlowEncodedTag[] = "FORWARD_FLOW_ENCODED";
|
||||
constexpr char kFloatContextFeatureOtherTag[] = "FLOAT_CONTEXT_FEATURE_OTHER";
|
||||
constexpr char kFloatContextFeatureTestTag[] = "FLOAT_CONTEXT_FEATURE_TEST";
|
||||
@@ -215,6 +217,54 @@ TEST_F(PackMediaSequenceCalculatorTest, PacksTwoFloatLists) {
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(PackMediaSequenceCalculatorTest, PacksTwoBytesLists) {
|
||||
SetUpCalculator({"BYTES_FEATURE_TEST:test", "BYTES_FEATURE_OTHER:test2"}, {},
|
||||
false, true);
|
||||
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
|
||||
|
||||
int num_timesteps = 2;
|
||||
for (int i = 0; i < num_timesteps; ++i) {
|
||||
auto vs_ptr = ::absl::make_unique<std::vector<std::string>>(
|
||||
2, absl::StrCat("foo", 2 << i));
|
||||
runner_->MutableInputs()
|
||||
->Tag(kBytesFeatureTestTag)
|
||||
.packets.push_back(Adopt(vs_ptr.release()).At(Timestamp(i)));
|
||||
vs_ptr = ::absl::make_unique<std::vector<std::string>>(
|
||||
2, absl::StrCat("bar", 2 << i));
|
||||
runner_->MutableInputs()
|
||||
->Tag(kBytesFeatureOtherTag)
|
||||
.packets.push_back(Adopt(vs_ptr.release()).At(Timestamp(i)));
|
||||
}
|
||||
|
||||
runner_->MutableSidePackets()->Tag(kSequenceExampleTag) =
|
||||
Adopt(input_sequence.release());
|
||||
|
||||
MP_ASSERT_OK(runner_->Run());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag(kSequenceExampleTag).packets;
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
const tf::SequenceExample& output_sequence =
|
||||
output_packets[0].Get<tf::SequenceExample>();
|
||||
|
||||
ASSERT_EQ(num_timesteps,
|
||||
mpms::GetFeatureTimestampSize("TEST", output_sequence));
|
||||
ASSERT_EQ(num_timesteps, mpms::GetFeatureBytesSize("TEST", output_sequence));
|
||||
ASSERT_EQ(num_timesteps,
|
||||
mpms::GetFeatureTimestampSize("OTHER", output_sequence));
|
||||
ASSERT_EQ(num_timesteps, mpms::GetFeatureBytesSize("OTHER", output_sequence));
|
||||
for (int i = 0; i < num_timesteps; ++i) {
|
||||
ASSERT_EQ(i, mpms::GetFeatureTimestampAt("TEST", output_sequence, i));
|
||||
ASSERT_THAT(mpms::GetFeatureBytesAt("TEST", output_sequence, i),
|
||||
::testing::ElementsAreArray(
|
||||
std::vector<std::string>(2, absl::StrCat("foo", 2 << i))));
|
||||
ASSERT_EQ(i, mpms::GetFeatureTimestampAt("OTHER", output_sequence, i));
|
||||
ASSERT_THAT(mpms::GetFeatureBytesAt("OTHER", output_sequence, i),
|
||||
::testing::ElementsAreArray(
|
||||
std::vector<std::string>(2, absl::StrCat("bar", 2 << i))));
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(PackMediaSequenceCalculatorTest, OutputAsZeroTimestamp) {
|
||||
SetUpCalculator({"FLOAT_FEATURE_TEST:test"}, {}, false, true, true);
|
||||
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
|
||||
@@ -829,6 +879,45 @@ TEST_F(PackMediaSequenceCalculatorTest, TestReplacingFloatVectors) {
|
||||
ASSERT_EQ(0, mpms::GetFeatureFloatsSize("OTHER", output_sequence));
|
||||
}
|
||||
|
||||
TEST_F(PackMediaSequenceCalculatorTest, TestReplacingBytesVectors) {
|
||||
SetUpCalculator({"BYTES_FEATURE_TEST:test", "BYTES_FEATURE_OTHER:test2"}, {},
|
||||
false, true);
|
||||
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
|
||||
|
||||
int num_timesteps = 2;
|
||||
for (int i = 0; i < num_timesteps; ++i) {
|
||||
auto vs_ptr = ::absl::make_unique<std::vector<std::string>>(
|
||||
2, absl::StrCat("foo", 2 << i));
|
||||
mpms::AddFeatureBytes("TEST", *vs_ptr, input_sequence.get());
|
||||
mpms::AddFeatureTimestamp("TEST", i, input_sequence.get());
|
||||
vs_ptr = ::absl::make_unique<std::vector<std::string>>(
|
||||
2, absl::StrCat("bar", 2 << i));
|
||||
mpms::AddFeatureBytes("OTHER", *vs_ptr, input_sequence.get());
|
||||
mpms::AddFeatureTimestamp("OTHER", i, input_sequence.get());
|
||||
}
|
||||
ASSERT_EQ(num_timesteps,
|
||||
mpms::GetFeatureTimestampSize("TEST", *input_sequence));
|
||||
ASSERT_EQ(num_timesteps, mpms::GetFeatureBytesSize("TEST", *input_sequence));
|
||||
ASSERT_EQ(num_timesteps,
|
||||
mpms::GetFeatureTimestampSize("OTHER", *input_sequence));
|
||||
ASSERT_EQ(num_timesteps, mpms::GetFeatureBytesSize("OTHER", *input_sequence));
|
||||
runner_->MutableSidePackets()->Tag(kSequenceExampleTag) =
|
||||
Adopt(input_sequence.release());
|
||||
|
||||
MP_ASSERT_OK(runner_->Run());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag(kSequenceExampleTag).packets;
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
const tf::SequenceExample& output_sequence =
|
||||
output_packets[0].Get<tf::SequenceExample>();
|
||||
|
||||
ASSERT_EQ(0, mpms::GetFeatureTimestampSize("TEST", output_sequence));
|
||||
ASSERT_EQ(0, mpms::GetFeatureFloatsSize("TEST", output_sequence));
|
||||
ASSERT_EQ(0, mpms::GetFeatureTimestampSize("OTHER", output_sequence));
|
||||
ASSERT_EQ(0, mpms::GetFeatureFloatsSize("OTHER", output_sequence));
|
||||
}
|
||||
|
||||
TEST_F(PackMediaSequenceCalculatorTest, TestReconcilingAnnotations) {
|
||||
SetUpCalculator({"IMAGE:images"}, {}, false, true);
|
||||
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
|
||||
|
||||
@@ -162,6 +162,27 @@ selects.config_setting_group(
|
||||
],
|
||||
)
|
||||
|
||||
config_setting(
|
||||
name = "edge_tpu_usb",
|
||||
define_values = {
|
||||
"MEDIAPIPE_EDGE_TPU": "usb",
|
||||
},
|
||||
)
|
||||
|
||||
config_setting(
|
||||
name = "edge_tpu_pci",
|
||||
define_values = {
|
||||
"MEDIAPIPE_EDGE_TPU": "pci",
|
||||
},
|
||||
)
|
||||
|
||||
config_setting(
|
||||
name = "edge_tpu_all",
|
||||
define_values = {
|
||||
"MEDIAPIPE_EDGE_TPU": "all",
|
||||
},
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tflite_inference_calculator",
|
||||
srcs = ["tflite_inference_calculator.cc"],
|
||||
@@ -172,6 +193,12 @@ cc_library(
|
||||
],
|
||||
"//conditions:default": [],
|
||||
}),
|
||||
defines = select({
|
||||
"//conditions:default": [],
|
||||
":edge_tpu_usb": ["MEDIAPIPE_EDGE_TPU=usb"],
|
||||
":edge_tpu_pci": ["MEDIAPIPE_EDGE_TPU=pci"],
|
||||
":edge_tpu_all": ["MEDIAPIPE_EDGE_TPU=all"],
|
||||
}),
|
||||
linkopts = select({
|
||||
"//mediapipe:ios": [
|
||||
"-framework CoreVideo",
|
||||
@@ -223,6 +250,20 @@ cc_library(
|
||||
"//conditions:default": [
|
||||
"//mediapipe/util:cpu_util",
|
||||
],
|
||||
}) + select({
|
||||
"//conditions:default": [],
|
||||
":edge_tpu_usb": [
|
||||
"@libedgetpu//tflite/public:edgetpu",
|
||||
"@libedgetpu//tflite/public:oss_edgetpu_direct_usb",
|
||||
],
|
||||
":edge_tpu_pci": [
|
||||
"@libedgetpu//tflite/public:edgetpu",
|
||||
"@libedgetpu//tflite/public:oss_edgetpu_direct_pci",
|
||||
],
|
||||
":edge_tpu_all": [
|
||||
"@libedgetpu//tflite/public:edgetpu",
|
||||
"@libedgetpu//tflite/public:oss_edgetpu_direct_all",
|
||||
],
|
||||
}),
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
@@ -85,7 +85,22 @@ constexpr char kTensorsGpuTag[] = "TENSORS_GPU";
|
||||
} // namespace
|
||||
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
#include "edgetpu.h"
|
||||
#include "tflite/public/edgetpu.h"
|
||||
|
||||
// Checkes whether model contains Edge TPU custom op or not.
|
||||
bool ContainsEdgeTpuCustomOp(const tflite::FlatBufferModel& model) {
|
||||
const auto* opcodes = model.GetModel()->operator_codes();
|
||||
for (const auto* subgraph : *model.GetModel()->subgraphs()) {
|
||||
for (const auto* op : *subgraph->operators()) {
|
||||
const auto* opcode = opcodes->Get(op->opcode_index());
|
||||
if (opcode->custom_code() &&
|
||||
opcode->custom_code()->str() == edgetpu::kCustomOp) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// Creates and returns an Edge TPU interpreter to run the given edgetpu model.
|
||||
std::unique_ptr<tflite::Interpreter> BuildEdgeTpuInterpreter(
|
||||
@@ -94,14 +109,9 @@ std::unique_ptr<tflite::Interpreter> BuildEdgeTpuInterpreter(
|
||||
edgetpu::EdgeTpuContext* edgetpu_context) {
|
||||
resolver->AddCustom(edgetpu::kCustomOp, edgetpu::RegisterCustomOp());
|
||||
std::unique_ptr<tflite::Interpreter> interpreter;
|
||||
if (tflite::InterpreterBuilder(model, *resolver)(&interpreter) != kTfLiteOk) {
|
||||
std::cerr << "Failed to build edge TPU interpreter." << std::endl;
|
||||
}
|
||||
CHECK_EQ(tflite::InterpreterBuilder(model, *resolver)(&interpreter),
|
||||
kTfLiteOk);
|
||||
interpreter->SetExternalContext(kTfLiteEdgeTpuContext, edgetpu_context);
|
||||
interpreter->SetNumThreads(1);
|
||||
if (interpreter->AllocateTensors() != kTfLiteOk) {
|
||||
std::cerr << "Failed to allocate edge TPU tensors." << std::endl;
|
||||
}
|
||||
return interpreter;
|
||||
}
|
||||
#endif // MEDIAPIPE_EDGE_TPU
|
||||
@@ -279,8 +289,7 @@ class TfLiteInferenceCalculator : public CalculatorBase {
|
||||
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
|
||||
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
std::shared_ptr<edgetpu::EdgeTpuContext> edgetpu_context_ =
|
||||
edgetpu::EdgeTpuManager::GetSingleton()->OpenDevice();
|
||||
std::shared_ptr<edgetpu::EdgeTpuContext> edgetpu_context_;
|
||||
#endif
|
||||
|
||||
bool gpu_inference_ = false;
|
||||
@@ -303,6 +312,10 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
// Calculator Core Section
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kCustomOpResolverTag[] = "CUSTOM_OP_RESOLVER";
|
||||
constexpr char kModelTag[] = "MODEL";
|
||||
|
||||
template <class CC>
|
||||
bool ShouldUseGpu(CC* cc) {
|
||||
#if MEDIAPIPE_TFLITE_GPU_SUPPORTED
|
||||
@@ -327,7 +340,7 @@ absl::Status TfLiteInferenceCalculator::GetContract(CalculatorContract* cc) {
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::TfLiteInferenceCalculatorOptions>();
|
||||
RET_CHECK(!options.model_path().empty() ^
|
||||
cc->InputSidePackets().HasTag("MODEL"))
|
||||
cc->InputSidePackets().HasTag(kModelTag))
|
||||
<< "Either model as side packet or model path in options is required.";
|
||||
|
||||
if (cc->Inputs().HasTag(kTensorsTag))
|
||||
@@ -340,13 +353,13 @@ absl::Status TfLiteInferenceCalculator::GetContract(CalculatorContract* cc) {
|
||||
if (cc->Outputs().HasTag(kTensorsGpuTag))
|
||||
cc->Outputs().Tag(kTensorsGpuTag).Set<std::vector<GpuTensor>>();
|
||||
|
||||
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
||||
if (cc->InputSidePackets().HasTag(kCustomOpResolverTag)) {
|
||||
cc->InputSidePackets()
|
||||
.Tag("CUSTOM_OP_RESOLVER")
|
||||
.Tag(kCustomOpResolverTag)
|
||||
.Set<tflite::ops::builtin::BuiltinOpResolver>();
|
||||
}
|
||||
if (cc->InputSidePackets().HasTag("MODEL")) {
|
||||
cc->InputSidePackets().Tag("MODEL").Set<TfLiteModelPtr>();
|
||||
if (cc->InputSidePackets().HasTag(kModelTag)) {
|
||||
cc->InputSidePackets().Tag(kModelTag).Set<TfLiteModelPtr>();
|
||||
}
|
||||
|
||||
if (ShouldUseGpu(cc)) {
|
||||
@@ -486,8 +499,8 @@ absl::Status TfLiteInferenceCalculator::Close(CalculatorContext* cc) {
|
||||
MP_RETURN_IF_ERROR(WriteKernelsToFile());
|
||||
|
||||
return RunInContextIfNeeded([this]() -> absl::Status {
|
||||
interpreter_ = nullptr;
|
||||
if (delegate_) {
|
||||
interpreter_ = nullptr;
|
||||
delegate_ = nullptr;
|
||||
#if MEDIAPIPE_TFLITE_GPU_SUPPORTED
|
||||
if (gpu_inference_) {
|
||||
@@ -501,7 +514,7 @@ absl::Status TfLiteInferenceCalculator::Close(CalculatorContext* cc) {
|
||||
#endif // MEDIAPIPE_TFLITE_GPU_SUPPORTED
|
||||
}
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
edgetpu_context_.reset();
|
||||
edgetpu_context_ = nullptr;
|
||||
#endif
|
||||
return absl::OkStatus();
|
||||
});
|
||||
@@ -723,9 +736,9 @@ absl::Status TfLiteInferenceCalculator::InitTFLiteGPURunner(
|
||||
auto op_resolver_ptr =
|
||||
static_cast<const tflite::ops::builtin::BuiltinOpResolver*>(
|
||||
&default_op_resolver);
|
||||
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
||||
if (cc->InputSidePackets().HasTag(kCustomOpResolverTag)) {
|
||||
op_resolver_ptr = &(cc->InputSidePackets()
|
||||
.Tag("CUSTOM_OP_RESOLVER")
|
||||
.Tag(kCustomOpResolverTag)
|
||||
.Get<tflite::ops::builtin::BuiltinOpResolver>());
|
||||
}
|
||||
|
||||
@@ -825,21 +838,26 @@ absl::Status TfLiteInferenceCalculator::LoadModel(CalculatorContext* cc) {
|
||||
|
||||
tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates
|
||||
default_op_resolver;
|
||||
auto op_resolver_ptr =
|
||||
static_cast<const tflite::ops::builtin::BuiltinOpResolver*>(
|
||||
&default_op_resolver);
|
||||
|
||||
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
||||
op_resolver_ptr = &(cc->InputSidePackets()
|
||||
.Tag("CUSTOM_OP_RESOLVER")
|
||||
.Get<tflite::ops::builtin::BuiltinOpResolver>());
|
||||
}
|
||||
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
interpreter_ =
|
||||
BuildEdgeTpuInterpreter(model, op_resolver_ptr, edgetpu_context_.get());
|
||||
#else
|
||||
tflite::InterpreterBuilder(model, *op_resolver_ptr)(&interpreter_);
|
||||
if (ContainsEdgeTpuCustomOp(model)) {
|
||||
edgetpu_context_ = edgetpu::EdgeTpuManager::GetSingleton()->OpenDevice();
|
||||
interpreter_ = BuildEdgeTpuInterpreter(model, &default_op_resolver,
|
||||
edgetpu_context_.get());
|
||||
} else {
|
||||
#endif // MEDIAPIPE_EDGE_TPU
|
||||
auto op_resolver_ptr =
|
||||
static_cast<const tflite::ops::builtin::BuiltinOpResolver*>(
|
||||
&default_op_resolver);
|
||||
|
||||
if (cc->InputSidePackets().HasTag(kCustomOpResolverTag)) {
|
||||
op_resolver_ptr = &(cc->InputSidePackets()
|
||||
.Tag(kCustomOpResolverTag)
|
||||
.Get<tflite::ops::builtin::BuiltinOpResolver>());
|
||||
}
|
||||
|
||||
tflite::InterpreterBuilder(model, *op_resolver_ptr)(&interpreter_);
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
}
|
||||
#endif // MEDIAPIPE_EDGE_TPU
|
||||
|
||||
RET_CHECK(interpreter_);
|
||||
@@ -872,8 +890,8 @@ absl::StatusOr<Packet> TfLiteInferenceCalculator::GetModelAsPacket(
|
||||
if (!options.model_path().empty()) {
|
||||
return TfLiteModelLoader::LoadFromPath(options.model_path());
|
||||
}
|
||||
if (cc.InputSidePackets().HasTag("MODEL")) {
|
||||
return cc.InputSidePackets().Tag("MODEL");
|
||||
if (cc.InputSidePackets().HasTag(kModelTag)) {
|
||||
return cc.InputSidePackets().Tag(kModelTag);
|
||||
}
|
||||
return absl::Status(absl::StatusCode::kNotFound,
|
||||
"Must specify TFLite model as path or loaded model.");
|
||||
@@ -929,6 +947,8 @@ absl::Status TfLiteInferenceCalculator::LoadDelegate(CalculatorContext* cc) {
|
||||
kTfLiteOk);
|
||||
return absl::OkStatus();
|
||||
}
|
||||
#else
|
||||
(void)use_xnnpack;
|
||||
#endif // !EDGETPU
|
||||
|
||||
// Return and use default tflite infernece (on CPU). No need for GPU
|
||||
|
||||
@@ -1353,3 +1353,34 @@ cc_test(
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "inverse_matrix_calculator",
|
||||
srcs = ["inverse_matrix_calculator.cc"],
|
||||
hdrs = ["inverse_matrix_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/api2:node",
|
||||
"//mediapipe/framework/api2:port",
|
||||
"@com_google_absl//absl/status",
|
||||
"@eigen_archive//:eigen3",
|
||||
],
|
||||
alwayslink = True,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "inverse_matrix_calculator_test",
|
||||
srcs = ["inverse_matrix_calculator_test.cc"],
|
||||
tags = ["desktop_only_test"],
|
||||
deps = [
|
||||
":inverse_matrix_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -33,6 +33,7 @@ namespace {
|
||||
constexpr char kImageFrameTag[] = "IMAGE_CPU";
|
||||
constexpr char kGpuBufferTag[] = "IMAGE_GPU";
|
||||
constexpr char kImageTag[] = "IMAGE";
|
||||
constexpr char kSourceOnGpuTag[] = "SOURCE_ON_GPU";
|
||||
} // namespace
|
||||
|
||||
// A calculator for converting the unified image container into
|
||||
@@ -46,6 +47,8 @@ constexpr char kImageTag[] = "IMAGE";
|
||||
// IMAGE_CPU: An ImageFrame containing output image.
|
||||
// IMAGE_GPU: A GpuBuffer containing output image.
|
||||
//
|
||||
// SOURCE_ON_GPU: The source Image is stored on GPU or CPU.
|
||||
//
|
||||
// Note:
|
||||
// Data is automatically transferred to/from the CPU or GPU
|
||||
// depending on output type.
|
||||
@@ -66,6 +69,7 @@ class FromImageCalculator : public CalculatorBase {
|
||||
absl::Status RenderGpu(CalculatorContext* cc);
|
||||
absl::Status RenderCpu(CalculatorContext* cc);
|
||||
|
||||
bool check_image_source_ = false;
|
||||
bool gpu_output_ = false;
|
||||
bool gpu_initialized_ = false;
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
@@ -102,6 +106,9 @@ absl::Status FromImageCalculator::GetContract(CalculatorContract* cc) {
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag(kSourceOnGpuTag)) {
|
||||
cc->Outputs().Tag(kSourceOnGpuTag).Set<bool>();
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
@@ -111,7 +118,9 @@ absl::Status FromImageCalculator::Open(CalculatorContext* cc) {
|
||||
if (cc->Outputs().HasTag(kGpuBufferTag)) {
|
||||
gpu_output_ = true;
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag(kSourceOnGpuTag)) {
|
||||
check_image_source_ = true;
|
||||
}
|
||||
if (gpu_output_) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
@@ -122,6 +131,13 @@ absl::Status FromImageCalculator::Open(CalculatorContext* cc) {
|
||||
}
|
||||
|
||||
absl::Status FromImageCalculator::Process(CalculatorContext* cc) {
|
||||
if (check_image_source_) {
|
||||
auto& input = cc->Inputs().Tag(kImageTag).Get<mediapipe::Image>();
|
||||
cc->Outputs()
|
||||
.Tag(kSourceOnGpuTag)
|
||||
.AddPacket(MakePacket<bool>(input.UsesGpu()).At(cc->InputTimestamp()));
|
||||
}
|
||||
|
||||
if (gpu_output_) {
|
||||
#if !MEDIAPIPE_DISABLE_GPU
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([&cc]() -> absl::Status {
|
||||
|
||||
@@ -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.
|
||||
|
||||
#include "mediapipe/calculators/util/inverse_matrix_calculator.h"
|
||||
|
||||
#include "Eigen/Core"
|
||||
#include "Eigen/Geometry"
|
||||
#include "Eigen/LU"
|
||||
#include "absl/status/status.h"
|
||||
#include "mediapipe/framework/api2/node.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace api2 {
|
||||
|
||||
class InverseMatrixCalculatorImpl : public NodeImpl<InverseMatrixCalculator> {
|
||||
absl::Status Process(mediapipe::CalculatorContext* cc) override {
|
||||
if (kInputMatrix(cc).IsEmpty()) {
|
||||
return absl::OkStatus();
|
||||
}
|
||||
Eigen::Matrix<float, 4, 4, Eigen::RowMajor> matrix(
|
||||
kInputMatrix(cc).Get().data());
|
||||
|
||||
Eigen::Matrix<float, 4, 4, Eigen::RowMajor> inverse_matrix;
|
||||
bool inverse_check;
|
||||
matrix.computeInverseWithCheck(inverse_matrix, inverse_check);
|
||||
RET_CHECK(inverse_check) << "Inverse matrix cannot be calculated.";
|
||||
|
||||
std::array<float, 16> output;
|
||||
Eigen::Map<Eigen::Matrix<float, 4, 4, Eigen::RowMajor>>(
|
||||
output.data(), 4, 4) = inverse_matrix.matrix();
|
||||
kOutputMatrix(cc).Send(std::move(output));
|
||||
return absl::OkStatus();
|
||||
}
|
||||
};
|
||||
MEDIAPIPE_NODE_IMPLEMENTATION(InverseMatrixCalculatorImpl);
|
||||
|
||||
} // namespace api2
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,51 @@
|
||||
// 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_INVERSE_MATRIX_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_UTIL_INVERSE_MATRIX_CALCULATOR_H_
|
||||
|
||||
#include "mediapipe/framework/api2/node.h"
|
||||
#include "mediapipe/framework/api2/port.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Runs affine transformation.
|
||||
//
|
||||
// Input:
|
||||
// MATRIX - std::array<float, 16>
|
||||
// Row major 4x4 matrix to inverse.
|
||||
//
|
||||
// Output:
|
||||
// MATRIX - std::array<float, 16>
|
||||
// Row major 4x4 inversed matrix.
|
||||
//
|
||||
// Usage example:
|
||||
// node {
|
||||
// calculator: "dishti.aimatter.InverseMatrixCalculator"
|
||||
// input_stream: "MATRIX:input_matrix"
|
||||
// output_stream: "MATRIX:output_matrix"
|
||||
// }
|
||||
class InverseMatrixCalculator : public mediapipe::api2::NodeIntf {
|
||||
public:
|
||||
static constexpr mediapipe::api2::Input<std::array<float, 16>> kInputMatrix{
|
||||
"MATRIX"};
|
||||
static constexpr mediapipe::api2::Output<std::array<float, 16>> kOutputMatrix{
|
||||
"MATRIX"};
|
||||
MEDIAPIPE_NODE_INTERFACE(InverseMatrixCalculator, kInputMatrix,
|
||||
kOutputMatrix);
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_UTIL_INVERSE_MATRIX_CALCULATOR_H_
|
||||
@@ -0,0 +1,126 @@
|
||||
#include "mediapipe/calculators/util/inverse_matrix_calculator.h"
|
||||
|
||||
#include <array>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
void RunTest(const std::array<float, 16>& matrix,
|
||||
const std::array<float, 16>& expected_inverse_matrix) {
|
||||
auto graph_config = mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"pb(
|
||||
input_stream: "matrix"
|
||||
node {
|
||||
calculator: "InverseMatrixCalculator"
|
||||
input_stream: "MATRIX:matrix"
|
||||
output_stream: "MATRIX:inverse_matrix"
|
||||
}
|
||||
)pb");
|
||||
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("inverse_matrix", &graph_config, &output_packets);
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"matrix",
|
||||
MakePacket<std::array<float, 16>>(std::move(matrix)).At(Timestamp(0))));
|
||||
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
ASSERT_THAT(output_packets, testing::SizeIs(1));
|
||||
|
||||
const auto& inverse_matrix = output_packets[0].Get<std::array<float, 16>>();
|
||||
|
||||
EXPECT_THAT(inverse_matrix, testing::Eq(expected_inverse_matrix));
|
||||
|
||||
// Fully close graph at end, otherwise calculator+tensors are destroyed
|
||||
// after calling WaitUntilDone().
|
||||
MP_ASSERT_OK(graph.CloseInputStream("matrix"));
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST(InverseMatrixCalculatorTest, Identity) {
|
||||
// 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,
|
||||
};
|
||||
std::array<float, 16> expected_inverse_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
|
||||
RunTest(matrix, expected_inverse_matrix);
|
||||
}
|
||||
|
||||
TEST(InverseMatrixCalculatorTest, Translation) {
|
||||
// clang-format off
|
||||
std::array<float, 16> matrix = {
|
||||
1.0f, 0.0f, 0.0f, 2.0f,
|
||||
0.0f, 1.0f, 0.0f, -5.0f,
|
||||
0.0f, 0.0f, 1.0f, 0.0f,
|
||||
0.0f, 0.0f, 0.0f, 1.0f,
|
||||
};
|
||||
std::array<float, 16> expected_inverse_matrix = {
|
||||
1.0f, 0.0f, 0.0f, -2.0f,
|
||||
0.0f, 1.0f, 0.0f, 5.0f,
|
||||
0.0f, 0.0f, 1.0f, 0.0f,
|
||||
0.0f, 0.0f, 0.0f, 1.0f,
|
||||
};
|
||||
// clang-format on
|
||||
RunTest(matrix, expected_inverse_matrix);
|
||||
}
|
||||
|
||||
TEST(InverseMatrixCalculatorTest, Scale) {
|
||||
// clang-format off
|
||||
std::array<float, 16> matrix = {
|
||||
5.0f, 0.0f, 0.0f, 0.0f,
|
||||
0.0f, 2.0f, 0.0f, 0.0f,
|
||||
0.0f, 0.0f, 1.0f, 0.0f,
|
||||
0.0f, 0.0f, 0.0f, 1.0f,
|
||||
};
|
||||
std::array<float, 16> expected_inverse_matrix = {
|
||||
0.2f, 0.0f, 0.0f, 0.0f,
|
||||
0.0f, 0.5f, 0.0f, 0.0f,
|
||||
0.0f, 0.0f, 1.0f, 0.0f,
|
||||
0.0f, 0.0f, 0.0f, 1.0f,
|
||||
};
|
||||
// clang-format on
|
||||
RunTest(matrix, expected_inverse_matrix);
|
||||
}
|
||||
|
||||
TEST(InverseMatrixCalculatorTest, Rotation90) {
|
||||
// clang-format off
|
||||
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,
|
||||
};
|
||||
std::array<float, 16> expected_inverse_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
|
||||
RunTest(matrix, expected_inverse_matrix);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
Reference in New Issue
Block a user