Project import generated by Copybara.
GitOrigin-RevId: 19a829ffd755edb43e54d20c0e7b9348512d5108
This commit is contained in:
committed by
schmidt-sebastian
parent
c6c80c3745
commit
7fb37c80e8
@@ -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",
|
||||
|
||||
@@ -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;
|
||||
};
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user