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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 19a829ffd755edb43e54d20c0e7b9348512d5108
This commit is contained in:
MediaPipe Team
2022-05-05 19:57:20 +00:00
committed by schmidt-sebastian
parent c6c80c3745
commit 7fb37c80e8
136 changed files with 2572 additions and 555 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/core/BUILD
+34
View File
@@ -214,6 +214,7 @@ cc_library(
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework:packet",
"//mediapipe/framework/formats:classification_cc_proto",
"//mediapipe/framework/formats:detection_cc_proto",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/framework/port:integral_types",
@@ -1257,3 +1258,36 @@ cc_test(
"@com_google_absl//absl/time",
],
)
cc_library(
name = "get_vector_item_calculator",
srcs = ["get_vector_item_calculator.cc"],
hdrs = ["get_vector_item_calculator.h"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:packet",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/api2:port",
"//mediapipe/framework/formats:classification_cc_proto",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
],
alwayslink = 1,
)
cc_library(
name = "vector_size_calculator",
srcs = ["vector_size_calculator.cc"],
hdrs = ["vector_size_calculator.h"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/formats:classification_cc_proto",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/port:status",
],
alwayslink = 1,
)
mediapipe/calculators/core/begin_loop_calculator.cc
+4
View File
@@ -28,6 +28,10 @@ typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedLandmarkList>>
BeginLoopNormalizedLandmarkListVectorCalculator;
REGISTER_CALCULATOR(BeginLoopNormalizedLandmarkListVectorCalculator);
// A calculator to process std::vector<int>.
typedef BeginLoopCalculator<std::vector<int>> BeginLoopIntCalculator;
REGISTER_CALCULATOR(BeginLoopIntCalculator);
// A calculator to process std::vector<NormalizedRect>.
typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedRect>>
BeginLoopNormalizedRectCalculator;
mediapipe/calculators/core/end_loop_calculator.cc
+5
View File
@@ -17,6 +17,7 @@
#include <vector>
#include "mediapipe/framework/formats/classification.pb.h"
#include "mediapipe/framework/formats/detection.pb.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/util/render_data.pb.h"
@@ -50,4 +51,8 @@ REGISTER_CALCULATOR(EndLoopClassificationListCalculator);
typedef EndLoopCalculator<std::vector<TfLiteTensor>> EndLoopTensorCalculator;
REGISTER_CALCULATOR(EndLoopTensorCalculator);
typedef EndLoopCalculator<std::vector<::mediapipe::Detection>>
EndLoopDetectionCalculator;
REGISTER_CALCULATOR(EndLoopDetectionCalculator);
} // namespace mediapipe
mediapipe/calculators/core/get_vector_item_calculator.cc
+32
View File
@@ -0,0 +1,32 @@
// Copyright 2022 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/core/get_vector_item_calculator.h"
#include "mediapipe/framework/formats/classification.pb.h"
#include "mediapipe/framework/formats/landmark.pb.h"
namespace mediapipe {
namespace api2 {
using GetLandmarkListVectorItemCalculator =
GetVectorItemCalculator<mediapipe::LandmarkList>;
REGISTER_CALCULATOR(GetLandmarkListVectorItemCalculator);
using GetClassificationListVectorItemCalculator =
GetVectorItemCalculator<mediapipe::ClassificationList>;
REGISTER_CALCULATOR(GetClassificationListVectorItemCalculator);
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/core/get_vector_item_calculator.h
+77
View File
@@ -0,0 +1,77 @@
// Copyright 2019 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MEDIAPIPE_CALCULATORS_CORE_GET_VECTOR_ITEM_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_CORE_GET_VECTOR_ITEM_CALCULATOR_H_
#include <optional>
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/api2/port.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
namespace mediapipe {
namespace api2 {
// A calcutlator to return an item from the vector by its index.
//
// Inputs:
// VECTOR - std::vector<T>
// Vector to take an item from.
// INDEX - int
// Index of the item to return.
//
// Outputs:
// ITEM - T
// Item from the vector at given index.
//
// Example config:
// node {
// calculator: "Get{SpecificType}VectorItemCalculator"
// input_stream: "VECTOR:vector"
// input_stream: "INDEX:index"
// input_stream: "ITEM:item"
// }
//
template <typename T>
class GetVectorItemCalculator : public Node {
public:
static constexpr Input<std::vector<T>> kIn{"VECTOR"};
static constexpr Input<int> kIdx{"INDEX"};
static constexpr Output<T> kOut{"ITEM"};
MEDIAPIPE_NODE_CONTRACT(kIn, kIdx, kOut);
absl::Status Process(CalculatorContext* cc) final {
if (kIn(cc).IsEmpty() || kIdx(cc).IsEmpty()) {
return absl::OkStatus();
}
const std::vector<T>& items = kIn(cc).Get();
const int idx = kIdx(cc).Get();
RET_CHECK_LT(idx, items.size());
kOut(cc).Send(items[idx]);
return absl::OkStatus();
}
};
} // namespace api2
} // namespace mediapipe
#endif // MEDIAPIPE_CALCULATORS_CORE_GET_VECTOR_ITEM_CALCULATOR_H_
mediapipe/calculators/core/split_vector_calculator.cc
+3
View File
@@ -83,4 +83,7 @@ REGISTER_CALCULATOR(SplitClassificationListVectorCalculator);
typedef SplitVectorCalculator<uint64_t, false> SplitUint64tVectorCalculator;
REGISTER_CALCULATOR(SplitUint64tVectorCalculator);
typedef SplitVectorCalculator<float, false> SplitFloatVectorCalculator;
REGISTER_CALCULATOR(SplitFloatVectorCalculator);
} // namespace mediapipe
mediapipe/calculators/core/vector_size_calculator.cc
+32
View File
@@ -0,0 +1,32 @@
// Copyright 2022 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/core/vector_size_calculator.h"
#include "mediapipe/framework/formats/classification.pb.h"
#include "mediapipe/framework/formats/landmark.pb.h"
namespace mediapipe {
namespace api2 {
using LandmarkListVectorSizeCalculator =
VectorSizeCalculator<mediapipe::LandmarkList>;
REGISTER_CALCULATOR(LandmarkListVectorSizeCalculator);
using ClassificationListVectorSizeCalculator =
VectorSizeCalculator<mediapipe::ClassificationList>;
REGISTER_CALCULATOR(ClassificationListVectorSizeCalculator);
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/core/vector_size_calculator.h
+64
View File
@@ -0,0 +1,64 @@
// Copyright 2022 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_CORE_VECTOR_SIZE_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_CORE_VECTOR_SIZE_CALCULATOR_H_
#include <optional>
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/status.h"
namespace mediapipe {
namespace api2 {
// A calcutlator to return vector size.
//
// Inputs:
// VECTOR - std::vector<T>
// Vector which size to return.
//
// Outputs:
// SIZE - int
// Size of the input vector.
//
// Example config:
// node {
// calculator: "{SpecificType}VectorSizeCalculator"
// input_stream: "VECTOR:vector"
// output_stream: "SIZE:vector_size"
// }
//
template <typename T>
class VectorSizeCalculator : public Node {
public:
static constexpr Input<std::vector<T>> kIn{"VECTOR"};
static constexpr Output<int> kOut{"SIZE"};
MEDIAPIPE_NODE_CONTRACT(kIn, kOut);
absl::Status Process(CalculatorContext* cc) final {
if (kIn(cc).IsEmpty()) {
return absl::OkStatus();
}
kOut(cc).Send(kIn(cc).Get().size());
return absl::OkStatus();
}
};
} // namespace api2
} // namespace mediapipe
#endif // MEDIAPIPE_CALCULATORS_CORE_VECTOR_SIZE_CALCULATOR_H_
mediapipe/calculators/image/scale_image_calculator.cc
+19 -2
View File
@@ -421,6 +421,10 @@ absl::Status ScaleImageCalculator::InitializeFromOptions() {
alignment_boundary_ = options_.alignment_boundary();
}
if (options_.has_output_format()) {
output_format_ = options_.output_format();
}
downscaler_.reset(new ImageResizer(options_.post_sharpening_coefficient()));
return absl::OkStatus();
@@ -433,13 +437,17 @@ absl::Status ScaleImageCalculator::ValidateImageFormats() const {
<< "The output image format was set to UNKNOWN.";
// TODO Remove these conditions.
RET_CHECK(output_format_ == ImageFormat::SRGB ||
output_format_ == ImageFormat::SRGBA ||
(input_format_ == output_format_ &&
output_format_ == ImageFormat::YCBCR420P))
<< "Outputting YCbCr420P images from SRGB input is not yet supported";
RET_CHECK(input_format_ == output_format_ ||
input_format_ == ImageFormat::YCBCR420P)
(input_format_ == ImageFormat::YCBCR420P &&
output_format_ == ImageFormat::SRGB) ||
(input_format_ == ImageFormat::SRGB &&
output_format_ == ImageFormat::SRGBA))
<< "Conversion of the color space (except from "
"YCbCr420P to SRGB) is not yet supported.";
"YCbCr420P to SRGB or SRGB to SRBGA) is not yet supported.";
return absl::OkStatus();
}
@@ -604,6 +612,15 @@ absl::Status ScaleImageCalculator::Process(CalculatorContext* cc) {
.Add(output_image.release(), cc->InputTimestamp());
return absl::OkStatus();
}
} else if (input_format_ == ImageFormat::SRGB &&
output_format_ == ImageFormat::SRGBA) {
image_frame = &cc->Inputs().Get(input_data_id_).Get<ImageFrame>();
cv::Mat input_mat = ::mediapipe::formats::MatView(image_frame);
converted_image_frame.Reset(ImageFormat::SRGBA, image_frame->Width(),
image_frame->Height(), alignment_boundary_);
cv::Mat output_mat = ::mediapipe::formats::MatView(&converted_image_frame);
cv::cvtColor(input_mat, output_mat, cv::COLOR_RGB2RGBA, 4);
image_frame = &converted_image_frame;
} else {
image_frame = &cc->Inputs().Get(input_data_id_).Get<ImageFrame>();
MP_RETURN_IF_ERROR(ValidateImageFrame(cc, *image_frame));
mediapipe/calculators/tensor/BUILD
+8 -2
View File
@@ -28,7 +28,9 @@ package(default_visibility = ["//visibility:private"])
exports_files(
glob(["testdata/image_to_tensor/*"]),
visibility = ["//mediapipe/calculators/image:__subpackages__"],
visibility = [
"//mediapipe/calculators/image:__subpackages__",
],
)
selects.config_setting_group(
@@ -64,15 +66,16 @@ cc_library(
":inference_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/api2:packet",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
"//mediapipe/framework/tool:subgraph_expansion",
"//mediapipe/util/tflite:config",
"//mediapipe/util/tflite:tflite_model_loader",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@org_tensorflow//tensorflow/lite:framework",
"@org_tensorflow//tensorflow/lite/core/api:op_resolver",
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
],
alwayslink = 1,
@@ -91,6 +94,7 @@ cc_library(
"//mediapipe/util/tflite:tflite_gpu_runner",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@org_tensorflow//tensorflow/lite:framework_stable",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
],
@@ -142,6 +146,8 @@ cc_library(
":inference_calculator_interface",
"@com_google_absl//absl/memory",
"@org_tensorflow//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
"@org_tensorflow//tensorflow/lite:framework_stable",
"@org_tensorflow//tensorflow/lite/c:c_api_types",
] + select({
"//conditions:default": [
"//mediapipe/util:cpu_util",
mediapipe/calculators/tensor/image_to_tensor_calculator.cc
+47 -21
View File
@@ -142,22 +142,35 @@ class ImageToTensorCalculator : public Node {
cc->Options<mediapipe::ImageToTensorCalculatorOptions>();
RET_CHECK(options.has_output_tensor_float_range() ||
options.has_output_tensor_int_range())
options.has_output_tensor_int_range() ||
options.has_output_tensor_uint_range())
<< "Output tensor range is required.";
if (options.has_output_tensor_float_range()) {
RET_CHECK_LT(options.output_tensor_float_range().min(),
options.output_tensor_float_range().max())
<< "Valid output float tensor range is required.";
}
if (options.has_output_tensor_uint_range()) {
RET_CHECK_LT(options.output_tensor_uint_range().min(),
options.output_tensor_uint_range().max())
<< "Valid output uint tensor range is required.";
RET_CHECK_GE(options.output_tensor_uint_range().min(), 0)
<< "The minimum of the output uint tensor range must be "
"non-negative.";
RET_CHECK_LE(options.output_tensor_uint_range().max(), 255)
<< "The maximum of the output uint tensor range must be less than or "
"equal to 255.";
}
if (options.has_output_tensor_int_range()) {
RET_CHECK_LT(options.output_tensor_int_range().min(),
options.output_tensor_int_range().max())
<< "Valid output int tensor range is required.";
RET_CHECK_GE(options.output_tensor_int_range().min(), 0)
<< "The minimum of the output int tensor range must be non-negative.";
RET_CHECK_LE(options.output_tensor_int_range().max(), 255)
RET_CHECK_GE(options.output_tensor_int_range().min(), -128)
<< "The minimum of the output int tensor range must be greater than "
"or equal to -128.";
RET_CHECK_LE(options.output_tensor_int_range().max(), 127)
<< "The maximum of the output int tensor range must be less than or "
"equal to 255.";
"equal to 127.";
}
RET_CHECK_GT(options.output_tensor_width(), 0)
<< "Valid output tensor width is required.";
@@ -187,15 +200,19 @@ class ImageToTensorCalculator : public Node {
options_ = cc->Options<mediapipe::ImageToTensorCalculatorOptions>();
output_width_ = options_.output_tensor_width();
output_height_ = options_.output_tensor_height();
is_int_output_ = options_.has_output_tensor_int_range();
range_min_ =
is_int_output_
? static_cast<float>(options_.output_tensor_int_range().min())
: options_.output_tensor_float_range().min();
range_max_ =
is_int_output_
? static_cast<float>(options_.output_tensor_int_range().max())
: options_.output_tensor_float_range().max();
is_float_output_ = options_.has_output_tensor_float_range();
if (options_.has_output_tensor_uint_range()) {
range_min_ =
static_cast<float>(options_.output_tensor_uint_range().min());
range_max_ =
static_cast<float>(options_.output_tensor_uint_range().max());
} else if (options_.has_output_tensor_int_range()) {
range_min_ = static_cast<float>(options_.output_tensor_int_range().min());
range_max_ = static_cast<float>(options_.output_tensor_int_range().max());
} else {
range_min_ = options_.output_tensor_float_range().min();
range_max_ = options_.output_tensor_float_range().max();
}
return absl::OkStatus();
}
@@ -275,6 +292,17 @@ class ImageToTensorCalculator : public Node {
}
}
Tensor::ElementType GetOutputTensorType() {
if (is_float_output_) {
return Tensor::ElementType::kFloat32;
}
if (range_min_ < 0) {
return Tensor::ElementType::kInt8;
} else {
return Tensor::ElementType::kUInt8;
}
}
absl::StatusOr<std::shared_ptr<const mediapipe::Image>> GetInputImage(
CalculatorContext* cc) {
if (kIn(cc).IsConnected()) {
@@ -305,7 +333,7 @@ class ImageToTensorCalculator : public Node {
const Image& image) {
// Lazy initialization of the GPU or CPU converter.
if (image.UsesGpu()) {
if (is_int_output_) {
if (!is_float_output_) {
return absl::UnimplementedError(
"ImageToTensorConverter for the input GPU image currently doesn't "
"support quantization.");
@@ -337,11 +365,9 @@ class ImageToTensorCalculator : public Node {
} else {
if (!cpu_converter_) {
#if !MEDIAPIPE_DISABLE_OPENCV
ASSIGN_OR_RETURN(cpu_converter_,
CreateOpenCvConverter(
cc, GetBorderMode(),
is_int_output_ ? Tensor::ElementType::kUInt8
: Tensor::ElementType::kFloat32));
ASSIGN_OR_RETURN(
cpu_converter_,
CreateOpenCvConverter(cc, GetBorderMode(), GetOutputTensorType()));
#else
LOG(FATAL) << "Cannot create image to tensor opencv converter since "
"MEDIAPIPE_DISABLE_OPENCV is defined.";
@@ -356,7 +382,7 @@ class ImageToTensorCalculator : public Node {
mediapipe::ImageToTensorCalculatorOptions options_;
int output_width_ = 0;
int output_height_ = 0;
bool is_int_output_ = false;
bool is_float_output_ = false;
float range_min_ = 0.0f;
float range_max_ = 1.0f;
};
mediapipe/calculators/tensor/image_to_tensor_calculator.proto
+9
View File
@@ -39,6 +39,14 @@ message ImageToTensorCalculatorOptions {
optional int64 max = 2;
}
// Range of uint values [min, max].
// min, must be strictly less than max.
// Please note that UIntRange is supported for CPU tensors only.
message UIntRange {
optional uint64 min = 1;
optional uint64 max = 2;
}
// Pixel extrapolation methods. See @border_mode.
enum BorderMode {
BORDER_UNSPECIFIED = 0;
@@ -58,6 +66,7 @@ message ImageToTensorCalculatorOptions {
oneof range {
FloatRange output_tensor_float_range = 4;
IntRange output_tensor_int_range = 7;
UIntRange output_tensor_uint_range = 8;
}
// For CONVENTIONAL mode for OpenGL, input image starts at bottom and needs
mediapipe/calculators/tensor/image_to_tensor_calculator_test.cc
+68 -50
View File
@@ -76,12 +76,21 @@ void RunTestWithInputImagePacket(const Packet& input_image_packet,
}
std::string output_tensor_range;
if (output_int_tensor) {
output_tensor_range = absl::Substitute(R"(output_tensor_int_range {
if (range_min < 0) {
output_tensor_range = absl::Substitute(R"(output_tensor_int_range {
min: $0
max: $1
})",
static_cast<int>(range_min),
static_cast<int>(range_max));
static_cast<int>(range_min),
static_cast<int>(range_max));
} else {
output_tensor_range = absl::Substitute(R"(output_tensor_uint_range {
min: $0
max: $1
})",
static_cast<uint>(range_min),
static_cast<uint>(range_max));
}
} else {
output_tensor_range = absl::Substitute(R"(output_tensor_float_range {
min: $0
@@ -141,9 +150,15 @@ void RunTestWithInputImagePacket(const Packet& input_image_packet,
auto view = tensor.GetCpuReadView();
cv::Mat tensor_mat;
if (output_int_tensor) {
EXPECT_EQ(tensor.element_type(), Tensor::ElementType::kUInt8);
tensor_mat = cv::Mat(tensor_height, tensor_width, CV_8UC3,
const_cast<uint8*>(view.buffer<uint8>()));
if (range_min < 0) {
EXPECT_EQ(tensor.element_type(), Tensor::ElementType::kInt8);
tensor_mat = cv::Mat(tensor_height, tensor_width, CV_8SC3,
const_cast<int8*>(view.buffer<int8>()));
} else {
EXPECT_EQ(tensor.element_type(), Tensor::ElementType::kUInt8);
tensor_mat = cv::Mat(tensor_height, tensor_width, CV_8UC3,
const_cast<uint8*>(view.buffer<uint8>()));
}
} else {
EXPECT_EQ(tensor.element_type(), Tensor::ElementType::kFloat32);
tensor_mat = cv::Mat(tensor_height, tensor_width, CV_32FC3,
@@ -190,25 +205,28 @@ const std::vector<InputType> kInputTypesToTest = {InputType::kImageFrame,
InputType::kImage};
void RunTest(cv::Mat input, cv::Mat expected_result,
std::vector<float> float_range, std::vector<int> int_range,
int tensor_width, int tensor_height, bool keep_aspect,
std::vector<std::pair<float, float>> float_ranges,
std::vector<std::pair<int, int>> int_ranges, int tensor_width,
int tensor_height, bool keep_aspect,
absl::optional<BorderMode> border_mode,
const mediapipe::NormalizedRect& roi) {
ASSERT_EQ(2, float_range.size());
ASSERT_EQ(2, int_range.size());
for (auto input_type : kInputTypesToTest) {
RunTestWithInputImagePacket(
input_type == InputType::kImageFrame ? MakeImageFramePacket(input)
: MakeImagePacket(input),
expected_result, float_range[0], float_range[1], tensor_width,
tensor_height, keep_aspect, border_mode, roi,
/*output_int_tensor=*/false);
RunTestWithInputImagePacket(
input_type == InputType::kImageFrame ? MakeImageFramePacket(input)
: MakeImagePacket(input),
expected_result, int_range[0], int_range[1], tensor_width,
tensor_height, keep_aspect, border_mode, roi,
/*output_int_tensor=*/true);
for (auto float_range : float_ranges) {
RunTestWithInputImagePacket(
input_type == InputType::kImageFrame ? MakeImageFramePacket(input)
: MakeImagePacket(input),
expected_result, float_range.first, float_range.second, tensor_width,
tensor_height, keep_aspect, border_mode, roi,
/*output_int_tensor=*/false);
}
for (auto int_range : int_ranges) {
RunTestWithInputImagePacket(
input_type == InputType::kImageFrame ? MakeImageFramePacket(input)
: MakeImagePacket(input),
expected_result, int_range.first, int_range.second, tensor_width,
tensor_height, keep_aspect, border_mode, roi,
/*output_int_tensor=*/true);
}
}
}
@@ -224,8 +242,8 @@ TEST(ImageToTensorCalculatorTest, MediumSubRectKeepAspect) {
"tensor/testdata/image_to_tensor/input.jpg"),
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/medium_sub_rect_keep_aspect.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/256, /*tensor_height=*/256, /*keep_aspect=*/true,
/*border mode*/ {}, roi);
}
@@ -242,8 +260,8 @@ TEST(ImageToTensorCalculatorTest, MediumSubRectKeepAspectBorderZero) {
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/"
"medium_sub_rect_keep_aspect_border_zero.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/256, /*tensor_height=*/256, /*keep_aspect=*/true,
BorderMode::kZero, roi);
}
@@ -260,8 +278,8 @@ TEST(ImageToTensorCalculatorTest, MediumSubRectKeepAspectWithRotation) {
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/"
"medium_sub_rect_keep_aspect_with_rotation.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}},
/*tensor_width=*/256, /*tensor_height=*/256, /*keep_aspect=*/true,
BorderMode::kReplicate, roi);
}
@@ -279,8 +297,8 @@ TEST(ImageToTensorCalculatorTest,
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/"
"medium_sub_rect_keep_aspect_with_rotation_border_zero.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/256, /*tensor_height=*/256, /*keep_aspect=*/true,
BorderMode::kZero, roi);
}
@@ -298,8 +316,8 @@ TEST(ImageToTensorCalculatorTest, MediumSubRectWithRotation) {
GetRgb(
"/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/medium_sub_rect_with_rotation.png"),
/*float_range=*/{-1.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{-1.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/256, /*tensor_height=*/256, /*keep_aspect=*/false,
BorderMode::kReplicate, roi);
}
@@ -316,8 +334,8 @@ TEST(ImageToTensorCalculatorTest, MediumSubRectWithRotationBorderZero) {
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/"
"medium_sub_rect_with_rotation_border_zero.png"),
/*float_range=*/{-1.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{-1.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/256, /*tensor_height=*/256, /*keep_aspect=*/false,
BorderMode::kZero, roi);
}
@@ -333,8 +351,8 @@ TEST(ImageToTensorCalculatorTest, LargeSubRect) {
"tensor/testdata/image_to_tensor/input.jpg"),
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/large_sub_rect.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}},
/*tensor_width=*/128, /*tensor_height=*/128, /*keep_aspect=*/false,
BorderMode::kReplicate, roi);
}
@@ -351,8 +369,8 @@ TEST(ImageToTensorCalculatorTest, LargeSubRectBorderZero) {
"tensor/testdata/image_to_tensor/input.jpg"),
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/large_sub_rect_border_zero.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/128, /*tensor_height=*/128, /*keep_aspect=*/false,
BorderMode::kZero, roi);
}
@@ -369,8 +387,8 @@ TEST(ImageToTensorCalculatorTest, LargeSubRectKeepAspect) {
"tensor/testdata/image_to_tensor/input.jpg"),
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/large_sub_rect_keep_aspect.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/128, /*tensor_height=*/128, /*keep_aspect=*/true,
BorderMode::kReplicate, roi);
}
@@ -387,8 +405,8 @@ TEST(ImageToTensorCalculatorTest, LargeSubRectKeepAspectBorderZero) {
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/"
"large_sub_rect_keep_aspect_border_zero.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/128, /*tensor_height=*/128, /*keep_aspect=*/true,
BorderMode::kZero, roi);
}
@@ -405,8 +423,8 @@ TEST(ImageToTensorCalculatorTest, LargeSubRectKeepAspectWithRotation) {
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/"
"large_sub_rect_keep_aspect_with_rotation.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/128, /*tensor_height=*/128, /*keep_aspect=*/true,
/*border_mode=*/{}, roi);
}
@@ -424,8 +442,8 @@ TEST(ImageToTensorCalculatorTest,
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/"
"large_sub_rect_keep_aspect_with_rotation_border_zero.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}},
/*tensor_width=*/128, /*tensor_height=*/128, /*keep_aspect=*/true,
/*border_mode=*/BorderMode::kZero, roi);
}
@@ -441,8 +459,8 @@ TEST(ImageToTensorCalculatorTest, NoOpExceptRange) {
"tensor/testdata/image_to_tensor/input.jpg"),
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/noop_except_range.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/64, /*tensor_height=*/128, /*keep_aspect=*/true,
BorderMode::kReplicate, roi);
}
@@ -458,8 +476,8 @@ TEST(ImageToTensorCalculatorTest, NoOpExceptRangeBorderZero) {
"tensor/testdata/image_to_tensor/input.jpg"),
GetRgb("/mediapipe/calculators/"
"tensor/testdata/image_to_tensor/noop_except_range.png"),
/*float_range=*/{0.0f, 1.0f},
/*int_range=*/{0, 255},
/*float_ranges=*/{{0.0f, 1.0f}},
/*int_ranges=*/{{0, 255}, {-128, 127}},
/*tensor_width=*/64, /*tensor_height=*/128, /*keep_aspect=*/true,
BorderMode::kZero, roi);
}
mediapipe/calculators/tensor/image_to_tensor_converter_gl_buffer.cc
+6 -4
View File
@@ -268,10 +268,12 @@ class GlProcessor : public ImageToTensorConverter {
const RotatedRect& roi,
const Size& output_dims, float range_min,
float range_max) override {
if (input.format() != mediapipe::GpuBufferFormat::kBGRA32) {
return InvalidArgumentError(
absl::StrCat("Only BGRA/RGBA textures are supported, passed format: ",
static_cast<uint32_t>(input.format())));
if (input.format() != mediapipe::GpuBufferFormat::kBGRA32 &&
input.format() != mediapipe::GpuBufferFormat::kRGBAHalf64 &&
input.format() != mediapipe::GpuBufferFormat::kRGBAFloat128) {
return InvalidArgumentError(absl::StrCat(
"Only 4-channel texture input formats are supported, passed format: ",
static_cast<uint32_t>(input.format())));
}
constexpr int kNumChannels = 3;
mediapipe/calculators/tensor/image_to_tensor_converter_gl_texture.cc
+6 -4
View File
@@ -172,10 +172,12 @@ class GlProcessor : public ImageToTensorConverter {
const RotatedRect& roi,
const Size& output_dims, float range_min,
float range_max) override {
if (input.format() != mediapipe::GpuBufferFormat::kBGRA32) {
return InvalidArgumentError(
absl::StrCat("Only BGRA/RGBA textures are supported, passed format: ",
static_cast<uint32_t>(input.format())));
if (input.format() != mediapipe::GpuBufferFormat::kBGRA32 &&
input.format() != mediapipe::GpuBufferFormat::kRGBAHalf64 &&
input.format() != mediapipe::GpuBufferFormat::kRGBAFloat128) {
return InvalidArgumentError(absl::StrCat(
"Only 4-channel texture input formats are supported, passed format: ",
static_cast<uint32_t>(input.format())));
}
constexpr int kNumChannels = 3;
mediapipe/calculators/tensor/image_to_tensor_converter_metal.cc
+6 -5
View File
@@ -352,11 +352,12 @@ class MetalProcessor : public ImageToTensorConverter {
const RotatedRect& roi,
const Size& output_dims, float range_min,
float range_max) override {
if (input.format() != mediapipe::GpuBufferFormat::kBGRA32) {
return InvalidArgumentError(
absl::StrCat("Only BGRA/RGBA textures are supported, passed "
"format: ",
static_cast<uint32_t>(input.format())));
if (input.format() != mediapipe::GpuBufferFormat::kBGRA32 &&
input.format() != mediapipe::GpuBufferFormat::kRGBAHalf64 &&
input.format() != mediapipe::GpuBufferFormat::kRGBAFloat128) {
return InvalidArgumentError(absl::StrCat(
"Only 4-channel texture input formats are supported, passed format: ",
static_cast<uint32_t>(input.format())));
}
@autoreleasepool {
mediapipe/calculators/tensor/image_to_tensor_converter_opencv.cc
+36 -7
View File
@@ -45,7 +45,19 @@ class OpenCvProcessor : public ImageToTensorConverter {
border_mode_ = cv::BORDER_CONSTANT;
break;
}
mat_type_ = tensor_type == Tensor::ElementType::kUInt8 ? CV_8UC3 : CV_32FC3;
switch (tensor_type_) {
case Tensor::ElementType::kInt8:
mat_type_ = CV_8SC3;
break;
case Tensor::ElementType::kFloat32:
mat_type_ = CV_32FC3;
break;
case Tensor::ElementType::kUInt8:
mat_type_ = CV_8UC3;
break;
default:
mat_type_ = -1;
}
}
absl::StatusOr<Tensor> Convert(const mediapipe::Image& input,
@@ -65,12 +77,22 @@ class OpenCvProcessor : public ImageToTensorConverter {
output_dims.width, kNumChannels});
auto buffer_view = tensor.GetCpuWriteView();
cv::Mat dst;
if (tensor_type_ == Tensor::ElementType::kUInt8) {
dst = cv::Mat(output_dims.height, output_dims.width, mat_type_,
buffer_view.buffer<uint8>());
} else {
dst = cv::Mat(output_dims.height, output_dims.width, mat_type_,
buffer_view.buffer<float>());
switch (tensor_type_) {
case Tensor::ElementType::kInt8:
dst = cv::Mat(output_dims.height, output_dims.width, mat_type_,
buffer_view.buffer<int8>());
break;
case Tensor::ElementType::kFloat32:
dst = cv::Mat(output_dims.height, output_dims.width, mat_type_,
buffer_view.buffer<float>());
break;
case Tensor::ElementType::kUInt8:
dst = cv::Mat(output_dims.height, output_dims.width, mat_type_,
buffer_view.buffer<uint8>());
break;
default:
return InvalidArgumentError(
absl::StrCat("Unsupported tensor type: ", tensor_type_));
}
const cv::RotatedRect rotated_rect(cv::Point2f(roi.center_x, roi.center_y),
@@ -124,6 +146,13 @@ class OpenCvProcessor : public ImageToTensorConverter {
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>> CreateOpenCvConverter(
CalculatorContext* cc, BorderMode border_mode,
Tensor::ElementType tensor_type) {
if (tensor_type != Tensor::ElementType::kInt8 &&
tensor_type != Tensor::ElementType::kFloat32 &&
tensor_type != Tensor::ElementType::kUInt8) {
return absl::InvalidArgumentError(absl::StrCat(
"Tensor type is currently not supported by OpenCvProcessor, type: ",
tensor_type));
}
return absl::make_unique<OpenCvProcessor>(border_mode, tensor_type);
}
mediapipe/calculators/tensor/inference_calculator.cc
+14
View File
@@ -21,7 +21,9 @@
#include "absl/memory/memory.h"
#include "absl/strings/string_view.h"
#include "mediapipe/framework/api2/packet.h"
#include "mediapipe/framework/tool/subgraph_expansion.h"
#include "tensorflow/lite/core/api/op_resolver.h"
namespace mediapipe {
namespace api2 {
@@ -67,5 +69,17 @@ absl::StatusOr<Packet<TfLiteModelPtr>> InferenceCalculator::GetModelAsPacket(
"Must specify TFLite model as path or loaded model.");
}
absl::StatusOr<Packet<tflite::OpResolver>>
InferenceCalculator::GetOpResolverAsPacket(CalculatorContext* cc) {
if (kSideInOpResolver(cc).IsConnected()) {
return kSideInOpResolver(cc).As<tflite::OpResolver>();
} else if (kSideInCustomOpResolver(cc).IsConnected()) {
return kSideInCustomOpResolver(cc).As<tflite::OpResolver>();
}
return PacketAdopting<tflite::OpResolver>(
std::make_unique<
tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates>());
}
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/inference_calculator.h
+15 -2
View File
@@ -27,6 +27,7 @@
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/tflite/tflite_model_loader.h"
#include "tensorflow/lite/core/api/op_resolver.h"
#include "tensorflow/lite/error_reporter.h"
#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/kernels/register.h"
@@ -55,8 +56,11 @@ namespace api2 {
// TENSORS - Vector of Tensors
//
// Input side packet:
// DEPRECATED: Prefer to use the "OP_RESOLVER" input side packet instead.
// CUSTOM_OP_RESOLVER (optional) - Use a custom op resolver,
// instead of the builtin one.
// OP_RESOLVER (optional) - Use to provide tflite op resolver
// (tflite::OpResolver)
// MODEL (optional) - Use to specify TfLite model
// (std::unique_ptr<tflite::FlatBufferModel,
// std::function<void(tflite::FlatBufferModel*)>>)
@@ -95,15 +99,21 @@ namespace api2 {
class InferenceCalculator : public NodeIntf {
public:
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
// Deprecated. Prefers to use "OP_RESOLVER" input side packet instead.
// TODO: Removes the "CUSTOM_OP_RESOLVER" side input after the
// migration.
static constexpr SideInput<tflite::ops::builtin::BuiltinOpResolver>::Optional
kSideInCustomOpResolver{"CUSTOM_OP_RESOLVER"};
static constexpr SideInput<tflite::OpResolver>::Optional kSideInOpResolver{
"OP_RESOLVER"};
static constexpr SideInput<TfLiteModelPtr>::Optional kSideInModel{"MODEL"};
static constexpr Output<std::vector<Tensor>> kOutTensors{"TENSORS"};
static constexpr SideInput<
mediapipe::InferenceCalculatorOptions::Delegate>::Optional kDelegate{
"DELEGATE"};
MEDIAPIPE_NODE_CONTRACT(kInTensors, kSideInCustomOpResolver, kSideInModel,
kOutTensors, kDelegate);
MEDIAPIPE_NODE_CONTRACT(kInTensors, kSideInCustomOpResolver,
kSideInOpResolver, kSideInModel, kOutTensors,
kDelegate);
protected:
using TfLiteDelegatePtr =
@@ -111,6 +121,9 @@ class InferenceCalculator : public NodeIntf {
absl::StatusOr<Packet<TfLiteModelPtr>> GetModelAsPacket(
CalculatorContext* cc);
absl::StatusOr<Packet<tflite::OpResolver>> GetOpResolverAsPacket(
CalculatorContext* cc);
};
struct InferenceCalculatorSelector : public InferenceCalculator {
mediapipe/calculators/tensor/inference_calculator.proto
+3
View File
@@ -116,6 +116,9 @@ message InferenceCalculatorOptions {
// to ensure there is no clash of the tokens. If unspecified, NNAPI will
// not try caching the compilation.
optional string model_token = 2;
// The name of an accelerator to be used for NNAPI delegate, e.g.
// "google-edgetpu". When not specified, it will be selected by NNAPI.
optional string accelerator_name = 3;
}
message Xnnpack {
// Number of threads for XNNPACK delegate. (By default, calculator tries
mediapipe/calculators/tensor/inference_calculator_cpu.cc
+57 -45
View File
@@ -19,7 +19,7 @@
#include "absl/memory/memory.h"
#include "mediapipe/calculators/tensor/inference_calculator.h"
#include "tensorflow/lite/interpreter_builder.h"
#if defined(MEDIAPIPE_ANDROID)
#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
#endif // ANDROID
@@ -28,6 +28,7 @@
#include "mediapipe/util/cpu_util.h"
#endif // !__EMSCRIPTEN__ || __EMSCRIPTEN_PTHREADS__
#include "tensorflow/lite/c/c_api_types.h"
#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
namespace mediapipe {
@@ -61,6 +62,17 @@ int GetXnnpackNumThreads(
return GetXnnpackDefaultNumThreads();
}
template <typename T>
void CopyTensorBuffer(const Tensor& input_tensor,
tflite::Interpreter* interpreter,
int input_tensor_index) {
auto input_tensor_view = input_tensor.GetCpuReadView();
auto input_tensor_buffer = input_tensor_view.buffer<T>();
T* local_tensor_buffer =
interpreter->typed_input_tensor<T>(input_tensor_index);
std::memcpy(local_tensor_buffer, input_tensor_buffer, input_tensor.bytes());
}
} // namespace
class InferenceCalculatorCpuImpl
@@ -73,15 +85,16 @@ class InferenceCalculatorCpuImpl
absl::Status Close(CalculatorContext* cc) override;
private:
absl::Status LoadModel(CalculatorContext* cc);
absl::Status LoadDelegate(CalculatorContext* cc);
absl::Status LoadDelegateAndAllocateTensors(CalculatorContext* cc);
absl::Status InitInterpreter(CalculatorContext* cc);
absl::Status LoadDelegate(CalculatorContext* cc,
tflite::InterpreterBuilder* interpreter_builder);
absl::Status AllocateTensors();
// TfLite requires us to keep the model alive as long as the interpreter is.
Packet<TfLiteModelPtr> model_packet_;
std::unique_ptr<tflite::Interpreter> interpreter_;
TfLiteDelegatePtr delegate_;
bool has_quantized_input_;
TfLiteType input_tensor_type_ = TfLiteType::kTfLiteNoType;
};
absl::Status InferenceCalculatorCpuImpl::UpdateContract(
@@ -94,8 +107,7 @@ absl::Status InferenceCalculatorCpuImpl::UpdateContract(
}
absl::Status InferenceCalculatorCpuImpl::Open(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(LoadModel(cc));
return LoadDelegateAndAllocateTensors(cc);
return InitInterpreter(cc);
}
absl::Status InferenceCalculatorCpuImpl::Process(CalculatorContext* cc) {
@@ -108,19 +120,23 @@ absl::Status InferenceCalculatorCpuImpl::Process(CalculatorContext* cc) {
// Read CPU input into tensors.
for (int i = 0; i < input_tensors.size(); ++i) {
const Tensor* input_tensor = &input_tensors[i];
auto input_tensor_view = input_tensor->GetCpuReadView();
if (has_quantized_input_) {
// TODO: Support more quantized tensor types.
auto input_tensor_buffer = input_tensor_view.buffer<uint8>();
uint8* local_tensor_buffer = interpreter_->typed_input_tensor<uint8>(i);
std::memcpy(local_tensor_buffer, input_tensor_buffer,
input_tensor->bytes());
} else {
auto input_tensor_buffer = input_tensor_view.buffer<float>();
float* local_tensor_buffer = interpreter_->typed_input_tensor<float>(i);
std::memcpy(local_tensor_buffer, input_tensor_buffer,
input_tensor->bytes());
switch (input_tensor_type_) {
case TfLiteType::kTfLiteFloat16:
case TfLiteType::kTfLiteFloat32: {
CopyTensorBuffer<float>(input_tensors[i], interpreter_.get(), i);
break;
}
case TfLiteType::kTfLiteUInt8: {
CopyTensorBuffer<uint8>(input_tensors[i], interpreter_.get(), i);
break;
}
case TfLiteType::kTfLiteInt8: {
CopyTensorBuffer<int8>(input_tensors[i], interpreter_.get(), i);
break;
}
default:
return absl::InvalidArgumentError(
absl::StrCat("Unsupported input tensor type:", input_tensor_type_));
}
}
@@ -150,39 +166,34 @@ absl::Status InferenceCalculatorCpuImpl::Close(CalculatorContext* cc) {
return absl::OkStatus();
}
absl::Status InferenceCalculatorCpuImpl::LoadModel(CalculatorContext* cc) {
absl::Status InferenceCalculatorCpuImpl::InitInterpreter(
CalculatorContext* cc) {
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(cc));
const auto& model = *model_packet_.Get();
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates());
tflite::InterpreterBuilder(model, op_resolver)(&interpreter_);
RET_CHECK(interpreter_);
ASSIGN_OR_RETURN(auto op_resolver_packet, GetOpResolverAsPacket(cc));
const auto& op_resolver = op_resolver_packet.Get();
tflite::InterpreterBuilder interpreter_builder(model, op_resolver);
MP_RETURN_IF_ERROR(LoadDelegate(cc, &interpreter_builder));
#if defined(__EMSCRIPTEN__)
interpreter_->SetNumThreads(1);
interpreter_builder.SetNumThreads(1);
#else
interpreter_->SetNumThreads(
interpreter_builder.SetNumThreads(
cc->Options<mediapipe::InferenceCalculatorOptions>().cpu_num_thread());
#endif // __EMSCRIPTEN__
return absl::OkStatus();
RET_CHECK_EQ(interpreter_builder(&interpreter_), kTfLiteOk);
RET_CHECK(interpreter_);
return AllocateTensors();
}
absl::Status InferenceCalculatorCpuImpl::LoadDelegateAndAllocateTensors(
CalculatorContext* cc) {
MP_RETURN_IF_ERROR(LoadDelegate(cc));
// AllocateTensors() can be called only after ModifyGraphWithDelegate.
absl::Status InferenceCalculatorCpuImpl::AllocateTensors() {
RET_CHECK_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
has_quantized_input_ =
interpreter_->tensor(interpreter_->inputs()[0])->quantization.type ==
kTfLiteAffineQuantization;
input_tensor_type_ = interpreter_->tensor(interpreter_->inputs()[0])->type;
return absl::OkStatus();
}
absl::Status InferenceCalculatorCpuImpl::LoadDelegate(CalculatorContext* cc) {
absl::Status InferenceCalculatorCpuImpl::LoadDelegate(
CalculatorContext* cc, tflite::InterpreterBuilder* interpreter_builder) {
const auto& calculator_opts =
cc->Options<mediapipe::InferenceCalculatorOptions>();
auto opts_delegate = calculator_opts.delegate();
@@ -211,18 +222,20 @@ absl::Status InferenceCalculatorCpuImpl::LoadDelegate(CalculatorContext* cc) {
if (nnapi_requested) {
// Attempt to use NNAPI.
// If not supported, the default CPU delegate will be created and used.
interpreter_->SetAllowFp16PrecisionForFp32(1);
tflite::StatefulNnApiDelegate::Options options;
const auto& nnapi = opts_delegate.nnapi();
options.allow_fp16 = true;
// Set up cache_dir and model_token for NNAPI compilation cache.
options.cache_dir =
nnapi.has_cache_dir() ? nnapi.cache_dir().c_str() : nullptr;
options.model_token =
nnapi.has_model_token() ? nnapi.model_token().c_str() : nullptr;
options.accelerator_name = nnapi.has_accelerator_name()
? nnapi.accelerator_name().c_str()
: nullptr;
delegate_ = TfLiteDelegatePtr(new tflite::StatefulNnApiDelegate(options),
[](TfLiteDelegate*) {});
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
interpreter_builder->AddDelegate(delegate_.get());
return absl::OkStatus();
}
#endif // MEDIAPIPE_ANDROID
@@ -239,8 +252,7 @@ absl::Status InferenceCalculatorCpuImpl::LoadDelegate(CalculatorContext* cc) {
GetXnnpackNumThreads(opts_has_delegate, opts_delegate);
delegate_ = TfLiteDelegatePtr(TfLiteXNNPackDelegateCreate(&xnnpack_opts),
&TfLiteXNNPackDelegateDelete);
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
interpreter_builder->AddDelegate(delegate_.get());
}
return absl::OkStatus();
mediapipe/calculators/tensor/inference_calculator_gl.cc
+29 -38
View File
@@ -22,6 +22,7 @@
#include "mediapipe/calculators/tensor/inference_calculator.h"
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/util/tflite/config.h"
#include "tensorflow/lite/interpreter_builder.h"
#if MEDIAPIPE_TFLITE_GL_INFERENCE
#include "mediapipe/gpu/gl_calculator_helper.h"
@@ -52,9 +53,11 @@ class InferenceCalculatorGlImpl
private:
absl::Status ReadGpuCaches();
absl::Status SaveGpuCaches();
absl::Status LoadModel(CalculatorContext* cc);
absl::Status LoadDelegate(CalculatorContext* cc);
absl::Status LoadDelegateAndAllocateTensors(CalculatorContext* cc);
absl::Status InitInterpreter(CalculatorContext* cc);
absl::Status LoadDelegate(CalculatorContext* cc,
tflite::InterpreterBuilder* interpreter_builder);
absl::Status BindBuffersToTensors();
absl::Status AllocateTensors();
absl::Status InitTFLiteGPURunner(CalculatorContext* cc);
// TfLite requires us to keep the model alive as long as the interpreter is.
@@ -137,17 +140,11 @@ absl::Status InferenceCalculatorGlImpl::Open(CalculatorContext* cc) {
#endif // MEDIAPIPE_ANDROID
}
// When use_advanced_gpu_api_, model loading is handled in InitTFLiteGPURunner
// for everything.
if (!use_advanced_gpu_api_) {
MP_RETURN_IF_ERROR(LoadModel(cc));
}
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, &cc]() -> ::mediapipe::Status {
return use_advanced_gpu_api_ ? InitTFLiteGPURunner(cc)
: LoadDelegateAndAllocateTensors(cc);
: InitInterpreter(cc);
}));
return absl::OkStatus();
}
@@ -292,12 +289,6 @@ absl::Status InferenceCalculatorGlImpl::ReadGpuCaches() {
absl::Status InferenceCalculatorGlImpl::InitTFLiteGPURunner(
CalculatorContext* cc) {
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(cc));
const auto& model = *model_packet_.Get();
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates());
// Create runner
tflite::gpu::InferenceOptions options;
options.priority1 = allow_precision_loss_
@@ -335,6 +326,10 @@ absl::Status InferenceCalculatorGlImpl::InitTFLiteGPURunner(
break;
}
}
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(cc));
const auto& model = *model_packet_.Get();
ASSIGN_OR_RETURN(auto op_resolver_packet, GetOpResolverAsPacket(cc));
const auto& op_resolver = op_resolver_packet.Get();
MP_RETURN_IF_ERROR(tflite_gpu_runner_->InitializeWithModel(
model, op_resolver, /*allow_quant_ops=*/true));
@@ -355,31 +350,27 @@ absl::Status InferenceCalculatorGlImpl::InitTFLiteGPURunner(
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::LoadModel(CalculatorContext* cc) {
absl::Status InferenceCalculatorGlImpl::InitInterpreter(CalculatorContext* cc) {
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(cc));
const auto& model = *model_packet_.Get();
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates());
tflite::InterpreterBuilder(model, op_resolver)(&interpreter_);
RET_CHECK(interpreter_);
ASSIGN_OR_RETURN(auto op_resolver_packet, GetOpResolverAsPacket(cc));
const auto& op_resolver = op_resolver_packet.Get();
tflite::InterpreterBuilder interpreter_builder(model, op_resolver);
MP_RETURN_IF_ERROR(LoadDelegate(cc, &interpreter_builder));
#if defined(__EMSCRIPTEN__)
interpreter_->SetNumThreads(1);
interpreter_builder.SetNumThreads(1);
#else
interpreter_->SetNumThreads(
interpreter_builder.SetNumThreads(
cc->Options<mediapipe::InferenceCalculatorOptions>().cpu_num_thread());
#endif // __EMSCRIPTEN__
RET_CHECK_EQ(interpreter_builder(&interpreter_), kTfLiteOk);
RET_CHECK(interpreter_);
MP_RETURN_IF_ERROR(BindBuffersToTensors());
MP_RETURN_IF_ERROR(AllocateTensors());
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::LoadDelegateAndAllocateTensors(
CalculatorContext* cc) {
MP_RETURN_IF_ERROR(LoadDelegate(cc));
// AllocateTensors() can be called only after ModifyGraphWithDelegate.
absl::Status InferenceCalculatorGlImpl::AllocateTensors() {
RET_CHECK_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
// TODO: Support quantized tensors.
RET_CHECK_NE(
@@ -388,7 +379,8 @@ absl::Status InferenceCalculatorGlImpl::LoadDelegateAndAllocateTensors(
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::LoadDelegate(CalculatorContext* cc) {
absl::Status InferenceCalculatorGlImpl::LoadDelegate(
CalculatorContext* cc, tflite::InterpreterBuilder* interpreter_builder) {
// Configure and create the delegate.
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
options.compile_options.precision_loss_allowed =
@@ -399,7 +391,11 @@ absl::Status InferenceCalculatorGlImpl::LoadDelegate(CalculatorContext* cc) {
options.compile_options.inline_parameters = 1;
delegate_ = TfLiteDelegatePtr(TfLiteGpuDelegateCreate(&options),
&TfLiteGpuDelegateDelete);
interpreter_builder->AddDelegate(delegate_.get());
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::BindBuffersToTensors() {
// Get input image sizes.
const auto& input_indices = interpreter_->inputs();
for (int i = 0; i < input_indices.size(); ++i) {
@@ -431,11 +427,6 @@ absl::Status InferenceCalculatorGlImpl::LoadDelegate(CalculatorContext* cc) {
output_indices[i]),
kTfLiteOk);
}
// Must call this last.
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
return absl::OkStatus();
}
mediapipe/calculators/tensor/inference_calculator_metal.cc
+21 -26
View File
@@ -90,9 +90,10 @@ class InferenceCalculatorMetalImpl
absl::Status Close(CalculatorContext* cc) override;
private:
absl::Status LoadModel(CalculatorContext* cc);
absl::Status LoadDelegate(CalculatorContext* cc);
absl::Status LoadDelegateAndAllocateTensors(CalculatorContext* cc);
absl::Status InitInterpreter(CalculatorContext* cc);
void AddDelegate(CalculatorContext* cc,
tflite::InterpreterBuilder* interpreter_builder);
absl::Status CreateConverters(CalculatorContext* cc);
// TfLite requires us to keep the model alive as long as the interpreter is.
Packet<TfLiteModelPtr> model_packet_;
@@ -127,11 +128,9 @@ absl::Status InferenceCalculatorMetalImpl::Open(CalculatorContext* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
allow_precision_loss_ = options.delegate().gpu().allow_precision_loss();
MP_RETURN_IF_ERROR(LoadModel(cc));
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(gpu_helper_);
return LoadDelegateAndAllocateTensors(cc);
return InitInterpreter(cc);
}
absl::Status InferenceCalculatorMetalImpl::Process(CalculatorContext* cc) {
@@ -199,27 +198,20 @@ absl::Status InferenceCalculatorMetalImpl::Close(CalculatorContext* cc) {
return absl::OkStatus();
}
absl::Status InferenceCalculatorMetalImpl::LoadModel(CalculatorContext* cc) {
absl::Status InferenceCalculatorMetalImpl::InitInterpreter(
CalculatorContext* cc) {
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(cc));
const auto& model = *model_packet_.Get();
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates());
tflite::InterpreterBuilder(model, op_resolver)(&interpreter_);
ASSIGN_OR_RETURN(auto op_resolver_packet, GetOpResolverAsPacket(cc));
const auto& op_resolver = op_resolver_packet.Get();
tflite::InterpreterBuilder interpreter_builder(model, op_resolver);
AddDelegate(cc, &interpreter_builder);
interpreter_builder.SetNumThreads(
cc->Options<mediapipe::InferenceCalculatorOptions>().cpu_num_thread());
RET_CHECK_EQ(interpreter_builder(&interpreter_), kTfLiteOk);
RET_CHECK(interpreter_);
interpreter_->SetNumThreads(
cc->Options<mediapipe::InferenceCalculatorOptions>().cpu_num_thread());
return absl::OkStatus();
}
absl::Status InferenceCalculatorMetalImpl::LoadDelegateAndAllocateTensors(
CalculatorContext* cc) {
MP_RETURN_IF_ERROR(LoadDelegate(cc));
// AllocateTensors() can be called only after ModifyGraphWithDelegate.
MP_RETURN_IF_ERROR(CreateConverters(cc));
RET_CHECK_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
// TODO: Support quantized tensors.
RET_CHECK_NE(
@@ -228,7 +220,8 @@ absl::Status InferenceCalculatorMetalImpl::LoadDelegateAndAllocateTensors(
return absl::OkStatus();
}
absl::Status InferenceCalculatorMetalImpl::LoadDelegate(CalculatorContext* cc) {
void InferenceCalculatorMetalImpl::AddDelegate(
CalculatorContext* cc, tflite::InterpreterBuilder* interpreter_builder) {
const auto& calculator_opts =
cc->Options<mediapipe::InferenceCalculatorOptions>();
@@ -242,9 +235,11 @@ absl::Status InferenceCalculatorMetalImpl::LoadDelegate(CalculatorContext* cc) {
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypeDoNotWait;
delegate_ =
TfLiteDelegatePtr(TFLGpuDelegateCreate(&options), &TFLGpuDelegateDelete);
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
interpreter_builder->AddDelegate(delegate_.get());
}
absl::Status InferenceCalculatorMetalImpl::CreateConverters(
CalculatorContext* cc) {
id<MTLDevice> device = gpu_helper_.mtlDevice;
// Get input image sizes.
mediapipe/calculators/tensor/tensors_to_detections_calculator.cc
+184 -28
View File
@@ -91,6 +91,40 @@ void ConvertAnchorsToRawValues(const std::vector<Anchor>& anchors,
}
}
absl::Status CheckCustomTensorMapping(
const TensorsToDetectionsCalculatorOptions::TensorMapping& tensor_mapping) {
RET_CHECK(tensor_mapping.has_detections_tensor_index() &&
tensor_mapping.has_scores_tensor_index());
int bitmap = 0;
bitmap |= 1 << tensor_mapping.detections_tensor_index();
bitmap |= 1 << tensor_mapping.scores_tensor_index();
if (!tensor_mapping.has_num_detections_tensor_index() &&
!tensor_mapping.has_classes_tensor_index() &&
!tensor_mapping.has_anchors_tensor_index()) {
// Only allows the output tensor index 0 and 1 to be occupied.
RET_CHECK_EQ(3, bitmap) << "The custom output tensor indices should only "
"cover index 0 and 1.";
} else if (tensor_mapping.has_anchors_tensor_index()) {
RET_CHECK(!tensor_mapping.has_classes_tensor_index() &&
!tensor_mapping.has_num_detections_tensor_index());
bitmap |= 1 << tensor_mapping.anchors_tensor_index();
// If the"anchors" tensor will be available, only allows the output tensor
// index 0, 1, 2 to be occupied.
RET_CHECK_EQ(7, bitmap) << "The custom output tensor indices should only "
"cover index 0, 1 and 2.";
} else {
RET_CHECK(tensor_mapping.has_classes_tensor_index() &&
tensor_mapping.has_num_detections_tensor_index());
// If the "classes" and the "number of detections" tensors will be
// available, only allows the output tensor index 0, 1, 2, 3 to be occupied.
bitmap |= 1 << tensor_mapping.classes_tensor_index();
bitmap |= 1 << tensor_mapping.num_detections_tensor_index();
RET_CHECK_EQ(15, bitmap) << "The custom output tensor indices should only "
"cover index 0, 1, 2 and 3.";
}
return absl::OkStatus();
}
} // namespace
// Convert result Tensors from object detection models into MediaPipe
@@ -170,13 +204,27 @@ class TensorsToDetectionsCalculator : public Node {
Detection ConvertToDetection(float box_ymin, float box_xmin, float box_ymax,
float box_xmax, float score, int class_id,
bool flip_vertically);
bool IsClassIndexAllowed(int class_index);
int num_classes_ = 0;
int num_boxes_ = 0;
int num_coords_ = 0;
std::set<int> ignore_classes_;
int max_results_ = -1;
::mediapipe::TensorsToDetectionsCalculatorOptions options_;
// Set of allowed or ignored class indices.
struct ClassIndexSet {
absl::flat_hash_set<int> values;
bool is_allowlist;
};
// Allowed or ignored class indices based on provided options or side packet.
// These are used to filter out the output detection results.
ClassIndexSet class_index_set_;
TensorsToDetectionsCalculatorOptions options_;
bool scores_tensor_index_is_set_ = false;
TensorsToDetectionsCalculatorOptions::TensorMapping tensor_mapping_;
std::vector<int> box_indices_ = {0, 1, 2, 3};
bool has_custom_box_indices_ = false;
std::vector<Anchor> anchors_;
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
@@ -239,6 +287,21 @@ absl::Status TensorsToDetectionsCalculator::Process(CalculatorContext* cc) {
}
}
}
const int num_input_tensors = kInTensors(cc)->size();
if (!scores_tensor_index_is_set_) {
if (num_input_tensors == 2 ||
num_input_tensors == kNumInputTensorsWithAnchors) {
tensor_mapping_.set_scores_tensor_index(1);
} else {
tensor_mapping_.set_scores_tensor_index(2);
}
scores_tensor_index_is_set_ = true;
}
if (gpu_processing || num_input_tensors != 4) {
// Allows custom bounding box indices when receiving 4 cpu tensors.
// Uses the default bbox indices in other cases.
RET_CHECK(!has_custom_box_indices_);
}
if (gpu_processing) {
if (!gpu_inited_) {
@@ -263,13 +326,15 @@ absl::Status TensorsToDetectionsCalculator::ProcessCPU(
// Postprocessing on CPU for model without postprocessing op. E.g. output
// raw score tensor and box tensor. Anchor decoding will be handled below.
// TODO: Add flexible input tensor size handling.
auto raw_box_tensor = &input_tensors[0];
auto raw_box_tensor =
&input_tensors[tensor_mapping_.detections_tensor_index()];
RET_CHECK_EQ(raw_box_tensor->shape().dims.size(), 3);
RET_CHECK_EQ(raw_box_tensor->shape().dims[0], 1);
RET_CHECK_GT(num_boxes_, 0) << "Please set num_boxes in calculator options";
RET_CHECK_EQ(raw_box_tensor->shape().dims[1], num_boxes_);
RET_CHECK_EQ(raw_box_tensor->shape().dims[2], num_coords_);
auto raw_score_tensor = &input_tensors[1];
auto raw_score_tensor =
&input_tensors[tensor_mapping_.scores_tensor_index()];
RET_CHECK_EQ(raw_score_tensor->shape().dims.size(), 3);
RET_CHECK_EQ(raw_score_tensor->shape().dims[0], 1);
RET_CHECK_EQ(raw_score_tensor->shape().dims[1], num_boxes_);
@@ -282,7 +347,8 @@ absl::Status TensorsToDetectionsCalculator::ProcessCPU(
// TODO: Support other options to load anchors.
if (!anchors_init_) {
if (input_tensors.size() == kNumInputTensorsWithAnchors) {
auto anchor_tensor = &input_tensors[2];
auto anchor_tensor =
&input_tensors[tensor_mapping_.anchors_tensor_index()];
RET_CHECK_EQ(anchor_tensor->shape().dims.size(), 2);
RET_CHECK_EQ(anchor_tensor->shape().dims[0], num_boxes_);
RET_CHECK_EQ(anchor_tensor->shape().dims[1], kNumCoordsPerBox);
@@ -308,7 +374,7 @@ absl::Status TensorsToDetectionsCalculator::ProcessCPU(
float max_score = -std::numeric_limits<float>::max();
// Find the top score for box i.
for (int score_idx = 0; score_idx < num_classes_; ++score_idx) {
if (ignore_classes_.find(score_idx) == ignore_classes_.end()) {
if (IsClassIndexAllowed(score_idx)) {
auto score = raw_scores[i * num_classes_ + score_idx];
if (options_.sigmoid_score()) {
if (options_.has_score_clipping_thresh()) {
@@ -338,23 +404,26 @@ absl::Status TensorsToDetectionsCalculator::ProcessCPU(
// Postprocessing on CPU with postprocessing op (e.g. anchor decoding and
// non-maximum suppression) within the model.
RET_CHECK_EQ(input_tensors.size(), 4);
auto num_boxes_tensor = &input_tensors[3];
auto num_boxes_tensor =
&input_tensors[tensor_mapping_.num_detections_tensor_index()];
RET_CHECK_EQ(num_boxes_tensor->shape().dims.size(), 1);
RET_CHECK_EQ(num_boxes_tensor->shape().dims[0], 1);
auto detection_boxes_tensor = &input_tensors[0];
auto detection_boxes_tensor =
&input_tensors[tensor_mapping_.detections_tensor_index()];
RET_CHECK_EQ(detection_boxes_tensor->shape().dims.size(), 3);
RET_CHECK_EQ(detection_boxes_tensor->shape().dims[0], 1);
const int max_detections = detection_boxes_tensor->shape().dims[1];
RET_CHECK_EQ(detection_boxes_tensor->shape().dims[2], num_coords_);
auto detection_classes_tensor = &input_tensors[1];
auto detection_classes_tensor =
&input_tensors[tensor_mapping_.classes_tensor_index()];
RET_CHECK_EQ(detection_classes_tensor->shape().dims.size(), 2);
RET_CHECK_EQ(detection_classes_tensor->shape().dims[0], 1);
RET_CHECK_EQ(detection_classes_tensor->shape().dims[1], max_detections);
auto detection_scores_tensor = &input_tensors[2];
auto detection_scores_tensor =
&input_tensors[tensor_mapping_.scores_tensor_index()];
RET_CHECK_EQ(detection_scores_tensor->shape().dims.size(), 2);
RET_CHECK_EQ(detection_scores_tensor->shape().dims[0], 1);
RET_CHECK_EQ(detection_scores_tensor->shape().dims[1], max_detections);
@@ -394,12 +463,14 @@ absl::Status TensorsToDetectionsCalculator::ProcessGPU(
-> absl::Status {
if (!anchors_init_) {
if (input_tensors.size() == kNumInputTensorsWithAnchors) {
auto read_view = input_tensors[2].GetOpenGlBufferReadView();
auto read_view = input_tensors[tensor_mapping_.anchors_tensor_index()]
.GetOpenGlBufferReadView();
glBindBuffer(GL_COPY_READ_BUFFER, read_view.name());
auto write_view = raw_anchors_buffer_->GetOpenGlBufferWriteView();
glBindBuffer(GL_COPY_WRITE_BUFFER, write_view.name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0,
input_tensors[2].bytes());
glCopyBufferSubData(
GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0,
input_tensors[tensor_mapping_.anchors_tensor_index()].bytes());
} else if (!kInAnchors(cc).IsEmpty()) {
const auto& anchors = *kInAnchors(cc);
auto anchors_view = raw_anchors_buffer_->GetCpuWriteView();
@@ -418,7 +489,9 @@ absl::Status TensorsToDetectionsCalculator::ProcessGPU(
auto decoded_boxes_view =
decoded_boxes_buffer_->GetOpenGlBufferWriteView();
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, decoded_boxes_view.name());
auto input0_view = input_tensors[0].GetOpenGlBufferReadView();
auto input0_view =
input_tensors[tensor_mapping_.detections_tensor_index()]
.GetOpenGlBufferReadView();
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, input0_view.name());
auto raw_anchors_view = raw_anchors_buffer_->GetOpenGlBufferReadView();
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, raw_anchors_view.name());
@@ -427,7 +500,8 @@ absl::Status TensorsToDetectionsCalculator::ProcessGPU(
// Score boxes.
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, scored_boxes_view.name());
auto input1_view = input_tensors[1].GetOpenGlBufferReadView();
auto input1_view = input_tensors[tensor_mapping_.scores_tensor_index()]
.GetOpenGlBufferReadView();
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, input1_view.name());
glUseProgram(score_program_);
glDispatchCompute(num_boxes_, 1, 1);
@@ -459,7 +533,8 @@ absl::Status TensorsToDetectionsCalculator::ProcessGPU(
if (input_tensors.size() == kNumInputTensorsWithAnchors) {
RET_CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
auto command_buffer = [gpu_helper_ commandBuffer];
auto src_buffer = input_tensors[2].GetMtlBufferReadView(command_buffer);
auto src_buffer = input_tensors[tensor_mapping_.anchors_tensor_index()]
.GetMtlBufferReadView(command_buffer);
auto dest_buffer =
raw_anchors_buffer_->GetMtlBufferWriteView(command_buffer);
id<MTLBlitCommandEncoder> blit_command =
@@ -468,7 +543,9 @@ absl::Status TensorsToDetectionsCalculator::ProcessGPU(
sourceOffset:0
toBuffer:dest_buffer.buffer()
destinationOffset:0
size:input_tensors[2].bytes()];
size:input_tensors[tensor_mapping_
.anchors_tensor_index()]
.bytes()];
[blit_command endEncoding];
[command_buffer commit];
} else if (!kInAnchors(cc).IsEmpty()) {
@@ -495,7 +572,8 @@ absl::Status TensorsToDetectionsCalculator::ProcessGPU(
auto decoded_boxes_view =
decoded_boxes_buffer_->GetMtlBufferWriteView(command_buffer);
[command_encoder setBuffer:decoded_boxes_view.buffer() offset:0 atIndex:0];
auto input0_view = input_tensors[0].GetMtlBufferReadView(command_buffer);
auto input0_view = input_tensors[tensor_mapping_.detections_tensor_index()]
.GetMtlBufferReadView(command_buffer);
[command_encoder setBuffer:input0_view.buffer() offset:0 atIndex:1];
auto raw_anchors_view =
raw_anchors_buffer_->GetMtlBufferReadView(command_buffer);
@@ -507,7 +585,8 @@ absl::Status TensorsToDetectionsCalculator::ProcessGPU(
[command_encoder setComputePipelineState:score_program_];
[command_encoder setBuffer:scored_boxes_view.buffer() offset:0 atIndex:0];
auto input1_view = input_tensors[1].GetMtlBufferReadView(command_buffer);
auto input1_view = input_tensors[tensor_mapping_.scores_tensor_index()]
.GetMtlBufferReadView(command_buffer);
[command_encoder setBuffer:input1_view.buffer() offset:0 atIndex:1];
MTLSize score_threads_per_group = MTLSizeMake(1, num_classes_, 1);
MTLSize score_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
@@ -570,6 +649,10 @@ absl::Status TensorsToDetectionsCalculator::LoadOptions(CalculatorContext* cc) {
num_classes_ = options_.num_classes();
num_boxes_ = options_.num_boxes();
num_coords_ = options_.num_coords();
CHECK_NE(options_.max_results(), 0)
<< "The maximum number of the top-scored detection results must be "
"non-zero.";
max_results_ = options_.max_results();
// Currently only support 2D when num_values_per_keypoint equals to 2.
CHECK_EQ(options_.num_values_per_keypoint(), 2);
@@ -581,15 +664,55 @@ absl::Status TensorsToDetectionsCalculator::LoadOptions(CalculatorContext* cc) {
if (kSideInIgnoreClasses(cc).IsConnected()) {
RET_CHECK(!kSideInIgnoreClasses(cc).IsEmpty());
RET_CHECK(options_.allow_classes().empty());
class_index_set_.is_allowlist = false;
for (int ignore_class : *kSideInIgnoreClasses(cc)) {
ignore_classes_.insert(ignore_class);
class_index_set_.values.insert(ignore_class);
}
} else if (!options_.allow_classes().empty()) {
RET_CHECK(options_.ignore_classes().empty());
class_index_set_.is_allowlist = true;
for (int i = 0; i < options_.allow_classes_size(); ++i) {
class_index_set_.values.insert(options_.allow_classes(i));
}
} else {
class_index_set_.is_allowlist = false;
for (int i = 0; i < options_.ignore_classes_size(); ++i) {
ignore_classes_.insert(options_.ignore_classes(i));
class_index_set_.values.insert(options_.ignore_classes(i));
}
}
if (options_.has_tensor_mapping()) {
RET_CHECK_OK(CheckCustomTensorMapping(options_.tensor_mapping()));
tensor_mapping_ = options_.tensor_mapping();
scores_tensor_index_is_set_ = true;
} else {
// Assigns the default tensor indices.
tensor_mapping_.set_detections_tensor_index(0);
tensor_mapping_.set_classes_tensor_index(1);
tensor_mapping_.set_anchors_tensor_index(2);
tensor_mapping_.set_num_detections_tensor_index(3);
// The scores tensor index needs to be determined based on the number of
// model's output tensors, which will be available in the first invocation
// of the Process() method.
tensor_mapping_.set_scores_tensor_index(-1);
scores_tensor_index_is_set_ = false;
}
if (options_.has_box_boundaries_indices()) {
box_indices_ = {options_.box_boundaries_indices().ymin(),
options_.box_boundaries_indices().xmin(),
options_.box_boundaries_indices().ymax(),
options_.box_boundaries_indices().xmax()};
int bitmap = 0;
for (int i : box_indices_) {
bitmap |= 1 << i;
}
RET_CHECK_EQ(bitmap, 15) << "The custom box boundaries indices should only "
"cover index 0, 1, 2, and 3.";
has_custom_box_indices_ = true;
}
return absl::OkStatus();
}
@@ -661,14 +784,22 @@ absl::Status TensorsToDetectionsCalculator::ConvertToDetections(
const float* detection_boxes, const float* detection_scores,
const int* detection_classes, std::vector<Detection>* output_detections) {
for (int i = 0; i < num_boxes_; ++i) {
if (max_results_ > 0 && output_detections->size() == max_results_) {
break;
}
if (options_.has_min_score_thresh() &&
detection_scores[i] < options_.min_score_thresh()) {
continue;
}
if (!IsClassIndexAllowed(detection_classes[i])) {
continue;
}
const int box_offset = i * num_coords_;
Detection detection = ConvertToDetection(
detection_boxes[box_offset + 0], detection_boxes[box_offset + 1],
detection_boxes[box_offset + 2], detection_boxes[box_offset + 3],
/*box_ymin=*/detection_boxes[box_offset + box_indices_[0]],
/*box_xmin=*/detection_boxes[box_offset + box_indices_[1]],
/*box_ymax=*/detection_boxes[box_offset + box_indices_[2]],
/*box_xmax=*/detection_boxes[box_offset + box_indices_[3]],
detection_scores[i], detection_classes[i], options_.flip_vertically());
const auto& bbox = detection.location_data().relative_bounding_box();
if (bbox.width() < 0 || bbox.height() < 0 || std::isnan(bbox.width()) ||
@@ -910,7 +1041,7 @@ void main() {
options_.has_score_clipping_thresh() ? 1 : 0,
options_.has_score_clipping_thresh() ? options_.score_clipping_thresh()
: 0,
!ignore_classes_.empty() ? 1 : 0);
!IsClassIndexAllowed(0));
// # filter classes supported is hardware dependent.
int max_wg_size; // typically <= 1024
@@ -919,7 +1050,14 @@ void main() {
CHECK_LT(num_classes_, max_wg_size)
<< "# classes must be < " << max_wg_size;
// TODO support better filtering.
CHECK_LE(ignore_classes_.size(), 1) << "Only ignore class 0 is allowed";
if (class_index_set_.is_allowlist) {
CHECK_EQ(class_index_set_.values.size(),
IsClassIndexAllowed(0) ? num_classes_ : num_classes_ - 1)
<< "Only all classes >= class 0 or >= class 1";
} else {
CHECK_EQ(class_index_set_.values.size(), IsClassIndexAllowed(0) ? 0 : 1)
<< "Only ignore class 0 is allowed";
}
// Shader program
{
@@ -1126,10 +1264,17 @@ kernel void scoreKernel(
options_.has_score_clipping_thresh() ? 1 : 0,
options_.has_score_clipping_thresh() ? options_.score_clipping_thresh()
: 0,
ignore_classes_.size() ? 1 : 0);
!IsClassIndexAllowed(0));
// TODO support better filtering.
CHECK_LE(ignore_classes_.size(), 1) << "Only ignore class 0 is allowed";
if (class_index_set_.is_allowlist) {
CHECK_EQ(class_index_set_.values.size(),
IsClassIndexAllowed(0) ? num_classes_ : num_classes_ - 1)
<< "Only all classes >= class 0 or >= class 1";
} else {
CHECK_EQ(class_index_set_.values.size(), IsClassIndexAllowed(0) ? 0 : 1)
<< "Only ignore class 0 is allowed";
}
{
// Shader program
@@ -1161,5 +1306,16 @@ kernel void scoreKernel(
return absl::OkStatus();
}
bool TensorsToDetectionsCalculator::IsClassIndexAllowed(int class_index) {
if (class_index_set_.values.empty()) {
return true;
}
if (class_index_set_.is_allowlist) {
return class_index_set_.values.contains(class_index);
} else {
return !class_index_set_.values.contains(class_index);
}
}
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/tensors_to_detections_calculator.proto
+41
View File
@@ -57,7 +57,12 @@ message TensorsToDetectionsCalculatorOptions {
optional bool reverse_output_order = 14 [default = false];
// The ids of classes that should be ignored during decoding the score for
// each predicted box. Can be overridden with IGNORE_CLASSES side packet.
// `ignore_classes` and `allow_classes` are mutually exclusive.
repeated int32 ignore_classes = 8;
// The ids of classes that should be allowed during decoding the score for
// each predicted box. `ignore_classes` and `allow_classes` are mutually
// exclusive.
repeated int32 allow_classes = 21 [packed = true];
optional bool sigmoid_score = 15 [default = false];
optional float score_clipping_thresh = 16;
@@ -71,4 +76,40 @@ message TensorsToDetectionsCalculatorOptions {
// Score threshold for perserving decoded detections.
optional float min_score_thresh = 19;
// The maximum number of the detection results to return. If < 0, all
// available results will be returned.
// For the detection models that have built-in non max suppression op, the
// output detections are the top-scored results. Otherwise, the output
// detections are the first N results that have higher scores than
// `min_score_thresh`.
optional int32 max_results = 20 [default = -1];
// The custom model output tensor mapping.
// The indices of the "detections" tensor and the "scores" tensor are always
// required. If the model outputs an "anchors" tensor, `anchors_tensor_index`
// must be specified. If the model outputs both "classes" tensor and "number
// of detections" tensors, `classes_tensor_index` and
// `num_detections_tensor_index` must be set.
message TensorMapping {
optional int32 detections_tensor_index = 1;
optional int32 classes_tensor_index = 2;
optional int32 scores_tensor_index = 3;
optional int32 num_detections_tensor_index = 4;
optional int32 anchors_tensor_index = 5;
}
optional TensorMapping tensor_mapping = 22;
// Represents the bounding box by using the combination of boundaries,
// {ymin, xmin, ymax, xmax}.
// The default order is {ymin, xmin, ymax, xmax}.
message BoxBoundariesIndices {
optional int32 ymin = 1 [default = 0];
optional int32 xmin = 2 [default = 1];
optional int32 ymax = 3 [default = 2];
optional int32 xmax = 4 [default = 3];
}
oneof box_indices {
BoxBoundariesIndices box_boundaries_indices = 23;
}
}
mediapipe/calculators/tensorflow/tensor_to_image_frame_calculator.cc
+6 -2
View File
@@ -121,8 +121,12 @@ absl::Status TensorToImageFrameCalculator::Process(CalculatorContext* cc) {
if (d > 255) d = 255;
buffer[i] = d;
}
output = ::absl::make_unique<ImageFrame>(format, width, height,
width * depth, buffer.release());
output = ::absl::make_unique<ImageFrame>(
format, width, height, width * depth, buffer.release(),
[total_size](uint8* ptr) {
::operator delete[](ptr, total_size,
std::align_val_t(EIGEN_MAX_ALIGN_BYTES));
});
} else if (input_tensor.dtype() == tensorflow::DT_UINT8) {
if (scale_factor_ != 1.0) {
return absl::InvalidArgumentError("scale_factor_ given for uint8 tensor");
mediapipe/calculators/tflite/BUILD
+2 -1
View File
@@ -121,10 +121,11 @@ cc_library(
deps = [
":tflite_custom_op_resolver_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:packet",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/util/tflite:cpu_op_resolver",
"//mediapipe/util/tflite:op_resolver",
"@org_tensorflow//tensorflow/lite/core/api:op_resolver",
],
alwayslink = 1,
)
mediapipe/calculators/tflite/tflite_custom_op_resolver_calculator.cc
+37 -5
View File
@@ -12,14 +12,22 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <memory>
#include "mediapipe/calculators/tflite/tflite_custom_op_resolver_calculator.pb.h"
#include "mediapipe/framework/api2/packet.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/tflite/cpu_op_resolver.h"
#include "mediapipe/util/tflite/op_resolver.h"
#include "tensorflow/lite/core/api/op_resolver.h"
namespace mediapipe {
namespace {
constexpr char kOpResolverTag[] = "OP_RESOLVER";
} // namespace
// This calculator creates a custom op resolver as a side packet that can be
// used in TfLiteInferenceCalculator. Current custom op resolver supports the
// following custom op on CPU and GPU:
@@ -27,7 +35,9 @@ namespace mediapipe {
// MaxPoolArgmax
// MaxUnpooling
//
// Usage example:
// Usage examples:
//
// For using with TfliteInferenceCalculator:
// node {
// calculator: "TfLiteCustomOpResolverCalculator"
// output_side_packet: "op_resolver"
@@ -37,12 +47,27 @@ namespace mediapipe {
// }
// }
// }
//
// For using with InferenceCalculator:
// node {
// calculator: "TfLiteCustomOpResolverCalculator"
// output_side_packet: "OP_RESOLVER:op_resolver"
// node_options: {
// [type.googleapis.com/mediapipe.TfLiteCustomOpResolverCalculatorOptions] {
// use_gpu: true
// }
// }
// }
class TfLiteCustomOpResolverCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
cc->OutputSidePackets()
.Index(0)
.Set<tflite::ops::builtin::BuiltinOpResolver>();
if (cc->OutputSidePackets().HasTag(kOpResolverTag)) {
cc->OutputSidePackets().Tag(kOpResolverTag).Set<tflite::OpResolver>();
} else {
cc->OutputSidePackets()
.Index(0)
.Set<tflite::ops::builtin::BuiltinOpResolver>();
}
return absl::OkStatus();
}
@@ -59,7 +84,14 @@ class TfLiteCustomOpResolverCalculator : public CalculatorBase {
op_resolver = absl::make_unique<mediapipe::CpuOpResolver>();
}
cc->OutputSidePackets().Index(0).Set(Adopt(op_resolver.release()));
if (cc->OutputSidePackets().HasTag(kOpResolverTag)) {
cc->OutputSidePackets()
.Tag(kOpResolverTag)
.Set(mediapipe::api2::PacketAdopting<tflite::OpResolver>(
std::move(op_resolver)));
} else {
cc->OutputSidePackets().Index(0).Set(Adopt(op_resolver.release()));
}
return absl::OkStatus();
}
mediapipe/calculators/util/BUILD
+36
View File
@@ -54,6 +54,7 @@ mediapipe_proto_library(
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
"//mediapipe/util:label_map_proto",
],
)
@@ -304,6 +305,7 @@ cc_library(
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:packet",
"//mediapipe/util:resource_util",
"//mediapipe/util:label_map_cc_proto",
] + select({
"//mediapipe:android": [
"//mediapipe/util/android/file/base",
@@ -350,6 +352,40 @@ cc_library(
alwayslink = 1,
)
cc_library(
name = "detection_transformation_calculator",
srcs = ["detection_transformation_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/api2:packet",
"//mediapipe/framework/api2:port",
"//mediapipe/framework/formats:detection_cc_proto",
"//mediapipe/framework/formats:location_data_cc_proto",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
],
alwayslink = 1,
)
cc_test(
name = "detection_transformation_calculator_test",
size = "small",
srcs = ["detection_transformation_calculator_test.cc"],
deps = [
":detection_transformation_calculator",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework:packet",
"//mediapipe/framework/formats:detection_cc_proto",
"//mediapipe/framework/formats:location_data_cc_proto",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
],
)
cc_library(
name = "non_max_suppression_calculator",
srcs = ["non_max_suppression_calculator.cc"],
mediapipe/calculators/util/detection_label_id_to_text_calculator.cc
+41 -13
View File
@@ -12,12 +12,12 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include "absl/container/node_hash_map.h"
#include "mediapipe/calculators/util/detection_label_id_to_text_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/detection.pb.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/util/label_map.pb.h"
#include "mediapipe/util/resource_util.h"
#if defined(MEDIAPIPE_MOBILE)
@@ -53,8 +53,11 @@ class DetectionLabelIdToTextCalculator : public CalculatorBase {
absl::Status Process(CalculatorContext* cc) override;
private:
absl::node_hash_map<int, std::string> label_map_;
::mediapipe::DetectionLabelIdToTextCalculatorOptions options_;
// Local label map built from the calculator options' `label_map_path` or
// `label` field.
LabelMap local_label_map_;
bool keep_label_id_;
const LabelMap& GetLabelMap(CalculatorContext* cc);
};
REGISTER_CALCULATOR(DetectionLabelIdToTextCalculator);
@@ -69,13 +72,16 @@ absl::Status DetectionLabelIdToTextCalculator::GetContract(
absl::Status DetectionLabelIdToTextCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
options_ =
const auto& options =
cc->Options<::mediapipe::DetectionLabelIdToTextCalculatorOptions>();
if (options_.has_label_map_path()) {
if (options.has_label_map_path()) {
RET_CHECK(!options.has_label_map() && options.label().empty())
<< "Only can set one of the following fields in the CalculatorOptions: "
"label_map_path, label, and label_map.";
std::string string_path;
ASSIGN_OR_RETURN(string_path,
PathToResourceAsFile(options_.label_map_path()));
PathToResourceAsFile(options.label_map_path()));
std::string label_map_string;
MP_RETURN_IF_ERROR(file::GetContents(string_path, &label_map_string));
@@ -83,13 +89,21 @@ absl::Status DetectionLabelIdToTextCalculator::Open(CalculatorContext* cc) {
std::string line;
int i = 0;
while (std::getline(stream, line)) {
label_map_[i++] = line;
LabelMapItem item;
item.set_name(line);
(*local_label_map_.mutable_index_to_item())[i++] = item;
}
} else {
for (int i = 0; i < options_.label_size(); ++i) {
label_map_[i] = options_.label(i);
} else if (!options.label().empty()) {
RET_CHECK(!options.has_label_map())
<< "Only can set one of the following fields in the CalculatorOptions: "
"label_map_path, label, and label_map.";
for (int i = 0; i < options.label_size(); ++i) {
LabelMapItem item;
item.set_name(options.label(i));
(*local_label_map_.mutable_index_to_item())[i] = item;
}
}
keep_label_id_ = options.keep_label_id();
return absl::OkStatus();
}
@@ -101,13 +115,18 @@ absl::Status DetectionLabelIdToTextCalculator::Process(CalculatorContext* cc) {
Detection& output_detection = output_detections.back();
bool has_text_label = false;
for (const int32 label_id : output_detection.label_id()) {
if (label_map_.find(label_id) != label_map_.end()) {
output_detection.add_label(label_map_[label_id]);
if (GetLabelMap(cc).index_to_item().find(label_id) !=
GetLabelMap(cc).index_to_item().end()) {
auto item = GetLabelMap(cc).index_to_item().at(label_id);
output_detection.add_label(item.name());
if (item.has_display_name()) {
output_detection.add_display_name(item.display_name());
}
has_text_label = true;
}
}
// Remove label_id field if text labels exist.
if (has_text_label && !options_.keep_label_id()) {
if (has_text_label && !keep_label_id_) {
output_detection.clear_label_id();
}
}
@@ -117,4 +136,13 @@ absl::Status DetectionLabelIdToTextCalculator::Process(CalculatorContext* cc) {
return absl::OkStatus();
}
const LabelMap& DetectionLabelIdToTextCalculator::GetLabelMap(
CalculatorContext* cc) {
return !local_label_map_.index_to_item().empty()
? local_label_map_
: cc->Options<
::mediapipe::DetectionLabelIdToTextCalculatorOptions>()
.label_map();
}
} // namespace mediapipe
mediapipe/calculators/util/detection_label_id_to_text_calculator.proto
+5 -1
View File
@@ -17,6 +17,7 @@ syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
import "mediapipe/util/label_map.proto";
message DetectionLabelIdToTextCalculatorOptions {
extend mediapipe.CalculatorOptions {
@@ -26,7 +27,7 @@ message DetectionLabelIdToTextCalculatorOptions {
// Path to a label map file for getting the actual name of detected classes.
optional string label_map_path = 1;
// Alternative way to specify label map
// Alternative way to specify label map.
// label: "label for id 0"
// label: "label for id 1"
// ...
@@ -36,4 +37,7 @@ message DetectionLabelIdToTextCalculatorOptions {
// could be found. By setting this field to true, it is always copied to the
// output detections.
optional bool keep_label_id = 3;
// Label map.
optional LabelMap label_map = 4;
}
mediapipe/calculators/util/detection_transformation_calculator.cc
+298
View File
@@ -0,0 +1,298 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <memory>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/api2/packet.h"
#include "mediapipe/framework/api2/port.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/detection.pb.h"
#include "mediapipe/framework/formats/location_data.pb.h"
namespace mediapipe {
namespace api2 {
namespace {
template <typename T>
T BoundedValue(T value, T upper_bound) {
T output = std::min(value, upper_bound);
if (output < 0) {
return 0;
}
return output;
}
absl::Status ConvertRelativeBoundingBoxToBoundingBox(
const std::pair<int, int>& image_size, Detection* detection) {
const int image_width = image_size.first;
const int image_height = image_size.second;
const auto& relative_bbox =
detection->location_data().relative_bounding_box();
auto* bbox = detection->mutable_location_data()->mutable_bounding_box();
bbox->set_xmin(
BoundedValue<int>(relative_bbox.xmin() * image_width, image_width));
bbox->set_ymin(
BoundedValue<int>(relative_bbox.ymin() * image_height, image_height));
bbox->set_width(
BoundedValue<int>(relative_bbox.width() * image_width, image_width));
bbox->set_height(
BoundedValue<int>(relative_bbox.height() * image_height, image_height));
detection->mutable_location_data()->set_format(LocationData::BOUNDING_BOX);
detection->mutable_location_data()->clear_relative_bounding_box();
return absl::OkStatus();
}
absl::Status ConvertBoundingBoxToRelativeBoundingBox(
const std::pair<int, int>& image_size, Detection* detection) {
int image_width = image_size.first;
int image_height = image_size.second;
const auto& bbox = detection->location_data().bounding_box();
auto* relative_bbox =
detection->mutable_location_data()->mutable_relative_bounding_box();
relative_bbox->set_xmin(
BoundedValue<float>((float)bbox.xmin() / image_width, 1.0f));
relative_bbox->set_ymin(
BoundedValue<float>((float)bbox.ymin() / image_height, 1.0f));
relative_bbox->set_width(
BoundedValue<float>((float)bbox.width() / image_width, 1.0f));
relative_bbox->set_height(
BoundedValue<float>((float)bbox.height() / image_height, 1.0f));
detection->mutable_location_data()->clear_bounding_box();
detection->mutable_location_data()->set_format(
LocationData::RELATIVE_BOUNDING_BOX);
return absl::OkStatus();
}
absl::StatusOr<LocationData::Format> GetLocationDataFormat(
const Detection& detection) {
if (!detection.has_location_data()) {
return absl::InvalidArgumentError("Detection must have location data.");
}
LocationData::Format format = detection.location_data().format();
RET_CHECK(format == LocationData::RELATIVE_BOUNDING_BOX ||
format == LocationData::BOUNDING_BOX)
<< "Detection's location data format must be either "
"RELATIVE_BOUNDING_BOX or BOUNDING_BOX";
return format;
}
absl::StatusOr<LocationData::Format> GetLocationDataFormat(
std::vector<Detection>& detections) {
RET_CHECK(!detections.empty());
LocationData::Format output_format;
ASSIGN_OR_RETURN(output_format, GetLocationDataFormat(detections[0]));
for (int i = 1; i < detections.size(); ++i) {
ASSIGN_OR_RETURN(LocationData::Format format,
GetLocationDataFormat(detections[i]));
if (output_format != format) {
return absl::InvalidArgumentError(
"Input detections have different location data formats.");
}
}
return output_format;
}
absl::Status ConvertBoundingBox(const std::pair<int, int>& image_size,
Detection* detection) {
if (!detection->has_location_data()) {
return absl::InvalidArgumentError("Detection must have location data.");
}
switch (detection->location_data().format()) {
case LocationData::RELATIVE_BOUNDING_BOX:
return ConvertRelativeBoundingBoxToBoundingBox(image_size, detection);
case LocationData::BOUNDING_BOX:
return ConvertBoundingBoxToRelativeBoundingBox(image_size, detection);
default:
return absl::InvalidArgumentError(
"Detection's location data format must be either "
"RELATIVE_BOUNDING_BOX or BOUNDING_BOX.");
}
}
} // namespace
// Transforms relative bounding box(es) to pixel bounding box(es) in a detection
// proto/detection list/detection vector, or vice versa.
//
// Inputs:
// One of the following:
// DETECTION: A Detection proto.
// DETECTIONS: An std::vector<Detection>/ a DetectionList proto.
// IMAGE_SIZE: A std::pair<int, int> represention image width and height.
//
// Outputs:
// At least one of the following:
// PIXEL_DETECTION: A Detection proto with pixel bounding box.
// PIXEL_DETECTIONS: An std::vector<Detection> with pixel bounding boxes.
// PIXEL_DETECTION_LIST: A DetectionList proto with pixel bounding boxes.
// RELATIVE_DETECTION: A Detection proto with relative bounding box.
// RELATIVE_DETECTIONS: An std::vector<Detection> with relative bounding boxes.
// RELATIVE_DETECTION_LIST: A DetectionList proto with relative bounding boxes.
//
// Example config:
// For input detection(s) with relative bounding box(es):
// node {
// calculator: "DetectionTransformationCalculator"
// input_stream: "DETECTION:input_detection"
// input_stream: "IMAGE_SIZE:image_size"
// output_stream: "PIXEL_DETECTION:output_detection"
// output_stream: "PIXEL_DETECTIONS:output_detections"
// output_stream: "PIXEL_DETECTION_LIST:output_detection_list"
// }
//
// For input detection(s) with pixel bounding box(es):
// node {
// calculator: "DetectionTransformationCalculator"
// input_stream: "DETECTION:input_detection"
// input_stream: "IMAGE_SIZE:image_size"
// output_stream: "RELATIVE_DETECTION:output_detection"
// output_stream: "RELATIVE_DETECTIONS:output_detections"
// output_stream: "RELATIVE_DETECTION_LIST:output_detection_list"
// }
class DetectionTransformationCalculator : public Node {
public:
static constexpr Input<Detection>::Optional kInDetection{"DETECTION"};
static constexpr Input<OneOf<DetectionList, std::vector<Detection>>>::Optional
kInDetections{"DETECTIONS"};
static constexpr Input<std::pair<int, int>> kInImageSize{"IMAGE_SIZE"};
static constexpr Output<Detection>::Optional kOutPixelDetection{
"PIXEL_DETECTION"};
static constexpr Output<std::vector<Detection>>::Optional kOutPixelDetections{
"PIXEL_DETECTIONS"};
static constexpr Output<DetectionList>::Optional kOutPixelDetectionList{
"PIXEL_DETECTION_LIST"};
static constexpr Output<Detection>::Optional kOutRelativeDetection{
"RELATIVE_DETECTION"};
static constexpr Output<std::vector<Detection>>::Optional
kOutRelativeDetections{"RELATIVE_DETECTIONS"};
static constexpr Output<DetectionList>::Optional kOutRelativeDetectionList{
"RELATIVE_DETECTION_LIST"};
MEDIAPIPE_NODE_CONTRACT(kInDetection, kInDetections, kInImageSize,
kOutPixelDetection, kOutPixelDetections,
kOutPixelDetectionList, kOutRelativeDetection,
kOutRelativeDetections, kOutRelativeDetectionList);
static absl::Status UpdateContract(CalculatorContract* cc) {
RET_CHECK(kInImageSize(cc).IsConnected()) << "Image size must be provided.";
RET_CHECK(kInDetections(cc).IsConnected() ^ kInDetection(cc).IsConnected());
if (kInDetections(cc).IsConnected()) {
RET_CHECK(kOutPixelDetections(cc).IsConnected() ||
kOutPixelDetectionList(cc).IsConnected() ||
kOutRelativeDetections(cc).IsConnected() ||
kOutRelativeDetectionList(cc).IsConnected())
<< "Output must be a container of detections.";
}
RET_CHECK(kOutPixelDetections(cc).IsConnected() ||
kOutPixelDetectionList(cc).IsConnected() ||
kOutPixelDetection(cc).IsConnected() ||
kOutRelativeDetections(cc).IsConnected() ||
kOutRelativeDetectionList(cc).IsConnected() ||
kOutRelativeDetection(cc).IsConnected())
<< "Must connect at least one output stream.";
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) override {
output_pixel_bounding_boxes_ = kOutPixelDetections(cc).IsConnected() ||
kOutPixelDetectionList(cc).IsConnected() ||
kOutPixelDetection(cc).IsConnected();
output_relative_bounding_boxes_ =
kOutRelativeDetections(cc).IsConnected() ||
kOutRelativeDetectionList(cc).IsConnected() ||
kOutRelativeDetection(cc).IsConnected();
RET_CHECK(output_pixel_bounding_boxes_ ^ output_relative_bounding_boxes_)
<< "All output streams must have the same stream tag prefix, either "
"\"PIXEL\" or \"RELATIVE_\".";
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) override {
std::pair<int, int> image_size = kInImageSize(cc).Get();
std::vector<Detection> transformed_detections;
LocationData::Format input_location_data_format;
if (kInDetections(cc).IsConnected()) {
transformed_detections = kInDetections(cc).Visit(
[&](const DetectionList& detection_list) {
return std::vector<Detection>(detection_list.detection().begin(),
detection_list.detection().end());
},
[&](const std::vector<Detection>& detection_vector) {
return detection_vector;
});
ASSIGN_OR_RETURN(input_location_data_format,
GetLocationDataFormat(transformed_detections));
for (Detection& detection : transformed_detections) {
MP_RETURN_IF_ERROR(ConvertBoundingBox(image_size, &detection));
}
} else {
ASSIGN_OR_RETURN(input_location_data_format,
GetLocationDataFormat(kInDetection(cc).Get()));
Detection transformed_detection(kInDetection(cc).Get());
MP_RETURN_IF_ERROR(
ConvertBoundingBox(image_size, &transformed_detection));
transformed_detections.push_back(transformed_detection);
}
if (input_location_data_format == LocationData::RELATIVE_BOUNDING_BOX) {
RET_CHECK(!output_relative_bounding_boxes_)
<< "Input detections are with relative bounding box(es), and the "
"output detections must have pixel bounding box(es).";
if (kOutPixelDetection(cc).IsConnected()) {
kOutPixelDetection(cc).Send(transformed_detections[0]);
}
if (kOutPixelDetections(cc).IsConnected()) {
kOutPixelDetections(cc).Send(transformed_detections);
}
if (kOutPixelDetectionList(cc).IsConnected()) {
DetectionList detection_list;
for (const auto& detection : transformed_detections) {
detection_list.add_detection()->CopyFrom(detection);
}
kOutPixelDetectionList(cc).Send(detection_list);
}
} else {
RET_CHECK(!output_pixel_bounding_boxes_)
<< "Input detections are with pixel bounding box(es), and the "
"output detections must have relative bounding box(es).";
if (kOutRelativeDetection(cc).IsConnected()) {
kOutRelativeDetection(cc).Send(transformed_detections[0]);
}
if (kOutRelativeDetections(cc).IsConnected()) {
kOutRelativeDetections(cc).Send(transformed_detections);
}
if (kOutRelativeDetectionList(cc).IsConnected()) {
DetectionList detection_list;
for (const auto& detection : transformed_detections) {
detection_list.add_detection()->CopyFrom(detection);
}
kOutRelativeDetectionList(cc).Send(detection_list);
}
}
return absl::OkStatus();
}
private:
bool output_relative_bounding_boxes_;
bool output_pixel_bounding_boxes_;
};
MEDIAPIPE_REGISTER_NODE(DetectionTransformationCalculator);
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/util/detection_transformation_calculator_test.cc
+287
View File
@@ -0,0 +1,287 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <algorithm>
#include <memory>
#include <vector>
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_runner.h"
#include "mediapipe/framework/formats/detection.pb.h"
#include "mediapipe/framework/formats/location_data.pb.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
namespace {
constexpr char kDetectionTag[] = "DETECTION";
constexpr char kDetectionsTag[] = "DETECTIONS";
constexpr char kImageSizeTag[] = "IMAGE_SIZE";
constexpr char kPixelDetectionTag[] = "PIXEL_DETECTION";
constexpr char kPixelDetectionListTag[] = "PIXEL_DETECTION_LIST";
constexpr char kPixelDetectionsTag[] = "PIXEL_DETECTIONS";
constexpr char kRelativeDetectionListTag[] = "RELATIVE_DETECTION_LIST";
constexpr char kRelativeDetectionsTag[] = "RELATIVE_DETECTIONS";
Detection DetectionWithBoundingBox(int32 xmin, int32 ymin, int32 width,
int32 height) {
Detection detection;
LocationData* location_data = detection.mutable_location_data();
location_data->set_format(LocationData::BOUNDING_BOX);
location_data->mutable_bounding_box()->set_xmin(xmin);
location_data->mutable_bounding_box()->set_ymin(ymin);
location_data->mutable_bounding_box()->set_width(width);
location_data->mutable_bounding_box()->set_height(height);
return detection;
}
Detection DetectionWithRelativeBoundingBox(float xmin, float ymin, float width,
float height) {
Detection detection;
LocationData* location_data = detection.mutable_location_data();
location_data->set_format(LocationData::RELATIVE_BOUNDING_BOX);
location_data->mutable_relative_bounding_box()->set_xmin(xmin);
location_data->mutable_relative_bounding_box()->set_ymin(ymin);
location_data->mutable_relative_bounding_box()->set_width(width);
location_data->mutable_relative_bounding_box()->set_height(height);
return detection;
}
std::vector<Detection> ConvertToDetectionVector(
const DetectionList& detection_list) {
std::vector<Detection> output;
for (const auto& detection : detection_list.detection()) {
output.push_back(detection);
}
return output;
}
void CheckBoundingBox(const Detection& output, const Detection& expected) {
const auto& output_bbox = output.location_data().bounding_box();
const auto& expected_bbox = output.location_data().bounding_box();
EXPECT_THAT(output_bbox.xmin(), testing::Eq(expected_bbox.xmin()));
EXPECT_THAT(output_bbox.ymin(), testing::Eq(expected_bbox.ymin()));
EXPECT_THAT(output_bbox.width(), testing::Eq(expected_bbox.width()));
EXPECT_THAT(output_bbox.height(), testing::Eq(expected_bbox.height()));
}
void CheckRelativeBoundingBox(const Detection& output,
const Detection& expected) {
const auto& output_bbox = output.location_data().relative_bounding_box();
const auto& expected_bbox = output.location_data().relative_bounding_box();
EXPECT_THAT(output_bbox.xmin(), testing::FloatEq(expected_bbox.xmin()));
EXPECT_THAT(output_bbox.ymin(), testing::FloatEq(expected_bbox.ymin()));
EXPECT_THAT(output_bbox.width(), testing::FloatEq(expected_bbox.width()));
EXPECT_THAT(output_bbox.height(), testing::FloatEq(expected_bbox.height()));
}
void CheckOutputDetections(const std::vector<Detection>& expected,
const std::vector<Detection>& output) {
ASSERT_EQ(output.size(), expected.size());
for (int i = 0; i < output.size(); ++i) {
auto output_format = output[i].location_data().format();
ASSERT_TRUE(output_format == LocationData::RELATIVE_BOUNDING_BOX ||
output_format == LocationData::BOUNDING_BOX);
ASSERT_EQ(output_format, expected[i].location_data().format());
if (output_format == LocationData::RELATIVE_BOUNDING_BOX) {
CheckRelativeBoundingBox(output[i], expected[i]);
}
if (output_format == LocationData::BOUNDING_BOX) {
CheckBoundingBox(output[i], expected[i]);
}
}
}
TEST(DetectionsTransformationCalculatorTest, MissingImageSize) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
calculator: "DetectionTransformationCalculator"
input_stream: "DETECTIONS:detections"
output_stream: "PIXEL_DETECTION:detection"
)pb"));
auto status = runner.Run();
ASSERT_FALSE(status.ok());
EXPECT_THAT(status.message(),
testing::HasSubstr("Image size must be provided"));
}
TEST(DetectionsTransformationCalculatorTest, WrongOutputType) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
calculator: "DetectionTransformationCalculator"
input_stream: "DETECTIONS:detections"
input_stream: "IMAGE_SIZE:image_size"
output_stream: "PIXEL_DETECTION:detection"
)pb"));
auto status = runner.Run();
ASSERT_FALSE(status.ok());
EXPECT_THAT(status.message(),
testing::HasSubstr("Output must be a container of detections"));
}
TEST(DetectionsTransformationCalculatorTest, WrongLocationDataFormat) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
calculator: "DetectionTransformationCalculator"
input_stream: "DETECTION:input_detection"
input_stream: "IMAGE_SIZE:image_size"
output_stream: "PIXEL_DETECTION:output_detection"
)pb"));
Detection detection;
detection.mutable_location_data()->set_format(LocationData::GLOBAL);
runner.MutableInputs()
->Tag(kDetectionTag)
.packets.push_back(MakePacket<Detection>(detection).At(Timestamp(0)));
std::pair<int, int> image_size({2000, 1000});
runner.MutableInputs()
->Tag(kImageSizeTag)
.packets.push_back(
MakePacket<std::pair<int, int>>(image_size).At(Timestamp(0)));
auto status = runner.Run();
ASSERT_FALSE(status.ok());
EXPECT_THAT(status.message(),
testing::HasSubstr("location data format must be either "
"RELATIVE_BOUNDING_BOX or BOUNDING_BOX"));
}
TEST(DetectionsTransformationCalculatorTest,
ConvertBoundingBoxToRelativeBoundingBox) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
calculator: "DetectionTransformationCalculator"
input_stream: "DETECTIONS:input_detections"
input_stream: "IMAGE_SIZE:image_size"
output_stream: "RELATIVE_DETECTIONS:output_detections"
output_stream: "RELATIVE_DETECTION_LIST:output_detection_list"
)pb"));
auto detections(absl::make_unique<std::vector<Detection>>());
detections->push_back(DetectionWithBoundingBox(100, 200, 400, 300));
detections->push_back(DetectionWithBoundingBox(0, 0, 2000, 1000));
std::pair<int, int> image_size({2000, 1000});
runner.MutableInputs()
->Tag(kDetectionsTag)
.packets.push_back(Adopt(detections.release()).At(Timestamp(0)));
runner.MutableInputs()
->Tag(kImageSizeTag)
.packets.push_back(
MakePacket<std::pair<int, int>>(image_size).At(Timestamp(0)));
MP_ASSERT_OK(runner.Run());
std::vector<Detection> expected(
{DetectionWithRelativeBoundingBox(0.05, 0.2, 0.2, 0.3),
DetectionWithRelativeBoundingBox(0, 0, 1, 1)});
const std::vector<Packet>& detections_output =
runner.Outputs().Tag(kRelativeDetectionsTag).packets;
ASSERT_EQ(1, detections_output.size());
CheckOutputDetections(expected,
detections_output[0].Get<std::vector<Detection>>());
const std::vector<Packet>& detection_list_output =
runner.Outputs().Tag(kRelativeDetectionListTag).packets;
ASSERT_EQ(1, detection_list_output.size());
CheckOutputDetections(
expected,
ConvertToDetectionVector(detection_list_output[0].Get<DetectionList>()));
}
TEST(DetectionsTransformationCalculatorTest,
ConvertRelativeBoundingBoxToBoundingBox) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
calculator: "DetectionTransformationCalculator"
input_stream: "DETECTIONS:input_detections"
input_stream: "IMAGE_SIZE:image_size"
output_stream: "PIXEL_DETECTIONS:output_detections"
output_stream: "PIXEL_DETECTION_LIST:output_detection_list"
)pb"));
auto detections(absl::make_unique<std::vector<Detection>>());
detections->push_back(DetectionWithRelativeBoundingBox(0.1, 0.2, 0.3, 0.4));
detections->push_back(DetectionWithRelativeBoundingBox(0, 0, 1, 1));
std::pair<int, int> image_size({2000, 1000});
runner.MutableInputs()
->Tag(kDetectionsTag)
.packets.push_back(Adopt(detections.release()).At(Timestamp(0)));
runner.MutableInputs()
->Tag(kImageSizeTag)
.packets.push_back(
MakePacket<std::pair<int, int>>(image_size).At(Timestamp(0)));
MP_ASSERT_OK(runner.Run());
std::vector<Detection> expected({DetectionWithBoundingBox(100, 200, 400, 300),
DetectionWithBoundingBox(0, 0, 2000, 1000)});
const std::vector<Packet>& detections_output =
runner.Outputs().Tag(kPixelDetectionsTag).packets;
ASSERT_EQ(1, detections_output.size());
CheckOutputDetections(expected,
detections_output[0].Get<std::vector<Detection>>());
const std::vector<Packet>& detection_list_output =
runner.Outputs().Tag(kPixelDetectionListTag).packets;
ASSERT_EQ(1, detection_list_output.size());
CheckOutputDetections(
expected,
ConvertToDetectionVector(detection_list_output[0].Get<DetectionList>()));
}
TEST(DetectionsTransformationCalculatorTest, ConvertSingleDetection) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
calculator: "DetectionTransformationCalculator"
input_stream: "DETECTION:input_detection"
input_stream: "IMAGE_SIZE:image_size"
output_stream: "PIXEL_DETECTION:outpu_detection"
output_stream: "PIXEL_DETECTIONS:output_detections"
output_stream: "PIXEL_DETECTION_LIST:output_detection_list"
)pb"));
runner.MutableInputs()
->Tag(kDetectionTag)
.packets.push_back(MakePacket<Detection>(DetectionWithRelativeBoundingBox(
0.05, 0.2, 0.2, 0.3))
.At(Timestamp(0)));
std::pair<int, int> image_size({2000, 1000});
runner.MutableInputs()
->Tag(kImageSizeTag)
.packets.push_back(
MakePacket<std::pair<int, int>>(image_size).At(Timestamp(0)));
MP_ASSERT_OK(runner.Run());
std::vector<Detection> expected(
{DetectionWithBoundingBox(100, 200, 400, 300)});
const std::vector<Packet>& detection_output =
runner.Outputs().Tag(kPixelDetectionTag).packets;
ASSERT_EQ(1, detection_output.size());
CheckOutputDetections(expected, {detection_output[0].Get<Detection>()});
const std::vector<Packet>& detections_output =
runner.Outputs().Tag(kPixelDetectionsTag).packets;
ASSERT_EQ(1, detections_output.size());
CheckOutputDetections(expected,
detections_output[0].Get<std::vector<Detection>>());
const std::vector<Packet>& detection_list_output =
runner.Outputs().Tag(kPixelDetectionListTag).packets;
ASSERT_EQ(1, detection_list_output.size());
CheckOutputDetections(
expected,
ConvertToDetectionVector(detection_list_output[0].Get<DetectionList>()));
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/video/tracking_graph_test.cc
+1 -1
View File
@@ -181,7 +181,7 @@ class TrackingGraphTest : public Test {
// Each image is shifted to the right and bottom by kTranslationStep
// pixels compared with the previous image.
static constexpr int kTranslationStep = 10;
static constexpr float kEqualityTolerance = 3e-4f;
static constexpr float kEqualityTolerance = 1e-3f;
};
void TrackingGraphTest::ExpectBoxAtFrame(const TimedBoxProto& box, float frame,