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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: d073f8e21be2fcc0e503cb97c6695078b6b75310
This commit is contained in:
MediaPipe Team
2021-02-27 03:30:05 -05:00
committed by chuoling
parent 39309bedba
commit 350fbb2100
755 changed files with 16391 additions and 11075 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/tensor/BUILD
+112 -56
View File
@@ -15,6 +15,12 @@
load("@bazel_skylib//lib:selects.bzl", "selects")
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_proto_library")
load(
"//mediapipe/framework/tool:mediapipe_graph.bzl",
"mediapipe_binary_graph",
)
load("//mediapipe/framework:mediapipe_cc_test.bzl", "mediapipe_cc_test")
load("//mediapipe/framework:encode_binary_proto.bzl", "encode_binary_proto")
licenses(["notice"])
@@ -38,81 +44,128 @@ mediapipe_proto_library(
)
cc_library(
name = "inference_calculator",
name = "inference_calculator_interface",
srcs = ["inference_calculator.cc"],
hdrs = ["inference_calculator.h"],
copts = select({
# TODO: fix tensor.h not to require this, if possible
"//mediapipe:apple": [
"-x objective-c++",
"-fobjc-arc", # enable reference-counting
],
"//conditions:default": [],
}),
features = ["-layering_check"], # allow depending on inference_calculator_gpu_deps
linkopts = select({
deps = [
":inference_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:node",
"//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/kernels:builtin_ops",
],
alwayslink = 1,
)
cc_library(
name = "inference_calculator_gl",
srcs = ["inference_calculator_gl.cc"],
tags = ["nomac"], # config problem with cpuinfo via TF
deps = [
"inference_calculator_interface",
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/util/tflite:tflite_gpu_runner",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
],
alwayslink = 1,
)
cc_library(
name = "inference_calculator_metal",
srcs = ["inference_calculator_metal.cc"],
copts = [
"-x objective-c++",
"-fobjc-arc", # enable reference-counting
],
linkopts = [
"-framework CoreVideo",
"-framework MetalKit",
],
tags = ["ios"],
deps = [
"inference_calculator_interface",
"//mediapipe/gpu:MPPMetalHelper",
"//mediapipe/gpu:MPPMetalUtil",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/objc:mediapipe_framework_ios",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate_internal",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
"@org_tensorflow//tensorflow/lite/delegates/gpu/metal:buffer_convert",
],
alwayslink = 1,
)
cc_library(
name = "inference_calculator_cpu",
srcs = [
"inference_calculator_cpu.cc",
],
copts = select({
# TODO: fix tensor.h not to require this, if possible
"//mediapipe:apple": [
"-framework CoreVideo",
"-framework MetalKit",
"-x objective-c++",
"-fobjc-arc", # enable reference-counting
],
"//conditions:default": [],
}),
visibility = ["//visibility:public"],
deps = [
":inference_calculator_cc_proto",
":inference_calculator_interface",
"@com_google_absl//absl/memory",
"//mediapipe/framework/api2:node",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:tensor",
"//mediapipe/util/tflite:tflite_model_loader",
"//mediapipe/util/tflite:config",
"@org_tensorflow//tensorflow/lite:framework",
"@org_tensorflow//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
"//mediapipe/framework/port:ret_check",
] + select({
":compute_shader_unavailable": [],
"//conditions:default": [":inference_calculator_gpu_deps"],
}) + select({
"//conditions:default": [],
"//mediapipe:android": [
"//mediapipe/util/android/file/base",
"@org_tensorflow//tensorflow/lite/delegates/nnapi:nnapi_delegate",
],
}) + select({
"//conditions:default": [
"//mediapipe/util:cpu_util",
],
}) + select({
"//conditions:default": [],
"//mediapipe:android": ["@org_tensorflow//tensorflow/lite/delegates/nnapi:nnapi_delegate"],
}),
alwayslink = 1,
)
cc_library(
name = "inference_calculator_gpu_deps",
deps = selects.with_or({
"//mediapipe:ios": [
"//mediapipe/gpu:MPPMetalHelper",
"//mediapipe/gpu:MPPMetalUtil",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/objc:mediapipe_framework_ios",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
"@org_tensorflow//tensorflow/lite/delegates/gpu/metal:buffer_convert",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate_internal",
],
"//mediapipe:macos": [],
"//conditions:default": [
"//mediapipe/util/tflite:tflite_gpu_runner",
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gpu_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
],
name = "inference_calculator_gl_if_compute_shader_available",
deps = select({
":compute_shader_unavailable": [],
"//conditions:default": [":inference_calculator_gl"],
}),
)
cc_library(
name = "inference_calculator",
visibility = ["//visibility:public"],
deps = [
":inference_calculator_interface",
":inference_calculator_cpu",
] + select({
"//conditions:default": [":inference_calculator_gl_if_compute_shader_available"],
"//mediapipe:ios": [":inference_calculator_metal"],
}),
alwayslink = 1,
)
mediapipe_proto_library(
name = "tensor_converter_calculator_proto",
srcs = ["tensor_converter_calculator.proto"],
@@ -357,6 +410,7 @@ cc_library(
visibility = ["//visibility:public"],
deps = [
":tensors_to_classification_calculator_cc_proto",
"@com_google_absl//absl/container:node_hash_map",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
"//mediapipe/framework/api2:node",
@@ -427,6 +481,7 @@ cc_library(
":image_to_tensor_converter_opencv",
":image_to_tensor_utils",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/framework/formats:tensor",
@@ -452,6 +507,7 @@ cc_library(
],
"//mediapipe:apple": [
":image_to_tensor_converter_metal",
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:MPPMetalHelper",
"//mediapipe/gpu:gpu_buffer",
],
@@ -499,21 +555,21 @@ cc_test(
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework/deps:file_path",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:image_format_cc_proto",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:image_frame_opencv",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:commandlineflags",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:opencv_core",
"//mediapipe/framework/port:opencv_imgcodecs",
"//mediapipe/framework/port:opencv_imgproc",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/tool:validate_type",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@org_tensorflow//tensorflow/lite:framework",
],
)
@@ -529,7 +585,7 @@ cc_library(
}),
deps = [
":image_to_tensor_utils",
"//mediapipe/framework:packet",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:statusor",
],
@@ -550,9 +606,9 @@ cc_library(
":image_to_tensor_converter",
":image_to_tensor_utils",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:image_format_cc_proto",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:image_frame_opencv",
"//mediapipe/framework/formats:image_opencv",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:opencv_core",
"//mediapipe/framework/port:opencv_imgproc",
@@ -579,7 +635,7 @@ cc_library(
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:statusor",
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/framework/formats:image",
"//mediapipe/gpu:gpu_buffer_format",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:types",
@@ -612,7 +668,7 @@ cc_library(
"//mediapipe/framework/port:statusor",
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/framework/formats:image",
"//mediapipe/gpu:shader_util",
],
}),
@@ -663,7 +719,7 @@ cc_library(
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:statusor",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/framework/formats:image",
"//mediapipe/gpu:gpu_buffer_format",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:types",
mediapipe/calculators/tensor/image_to_tensor_calculator.cc
+105 -58
View File
@@ -22,6 +22,7 @@
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/framework/formats/tensor.h"
@@ -29,6 +30,7 @@
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/statusor.h"
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/gpu/gpu_buffer.h"
@@ -60,11 +62,18 @@ using GpuBuffer = mediapipe::GpuBuffer;
// normalization, according to specified inputs and options.
//
// Inputs:
// IMAGE - ImageFrame [ImageFormat::SRGB/SRGBA]
// Image to extract from.
// IMAGE - Image[ImageFormat::SRGB / SRGBA, GpuBufferFormat::kBGRA32] or
// ImageFrame [ImageFormat::SRGB/SRGBA] (for backward compatibility
// with existing graphs that use IMAGE for ImageFrame input)
// IMAGE_GPU - GpuBuffer [GpuBufferFormat::kBGRA32]
// Image to extract from.
// (Either IMAGE or IMAGE_GPU has to be specified.)
//
// Note:
// - One and only one of IMAGE and IMAGE_GPU should be specified.
// - IMAGE input of type Image is processed on GPU if the data is already on
// GPU (i.e., Image::UsesGpu() returns true), or otherwise processed on CPU.
// - IMAGE input of type ImageFrame is always processed on CPU.
// - IMAGE_GPU input (of type GpuBuffer) is always processed on GPU.
//
// NORM_RECT - NormalizedRect @Optional
// Describes region of image to extract.
@@ -112,7 +121,8 @@ using GpuBuffer = mediapipe::GpuBuffer;
// }
class ImageToTensorCalculator : public Node {
public:
static constexpr Input<mediapipe::ImageFrame>::Optional kInCpu{"IMAGE"};
static constexpr Input<
OneOf<mediapipe::Image, mediapipe::ImageFrame>>::Optional kIn{"IMAGE"};
static constexpr Input<GpuBuffer>::Optional kInGpu{"IMAGE_GPU"};
static constexpr Input<mediapipe::NormalizedRect>::Optional kInNormRect{
"NORM_RECT"};
@@ -121,10 +131,10 @@ class ImageToTensorCalculator : public Node {
"LETTERBOX_PADDING"};
static constexpr Output<std::array<float, 16>>::Optional kOutMatrix{"MATRIX"};
MEDIAPIPE_NODE_CONTRACT(kInCpu, kInGpu, kInNormRect, kOutTensors,
MEDIAPIPE_NODE_CONTRACT(kIn, kInGpu, kInNormRect, kOutTensors,
kOutLetterboxPadding, kOutMatrix);
static ::mediapipe::Status UpdateContract(CalculatorContract* cc) {
static absl::Status UpdateContract(CalculatorContract* cc) {
const auto& options =
cc->Options<mediapipe::ImageToTensorCalculatorOptions>();
@@ -138,69 +148,47 @@ class ImageToTensorCalculator : public Node {
RET_CHECK_GT(options.output_tensor_height(), 0)
<< "Valid output tensor height is required.";
RET_CHECK(kInCpu(cc).IsConnected() ^ kInGpu(cc).IsConnected())
<< "One and only one of CPU or GPU input is expected.";
RET_CHECK(kIn(cc).IsConnected() ^ kInGpu(cc).IsConnected())
<< "One and only one of IMAGE and IMAGE_GPU input is expected.";
if (kInGpu(cc).IsConnected()) {
#if MEDIAPIPE_DISABLE_GPU
return mediapipe::UnimplementedError("GPU processing is disabled");
#else
#if MEDIAPIPE_METAL_ENABLED
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#else
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // MEDIAPIPE_METAL_ENABLED
#endif // MEDIAPIPE_DISABLE_GPU
if (kInGpu(cc).IsConnected()) {
return absl::UnimplementedError(
"GPU processing is disabled in build flags");
}
return mediapipe::OkStatus();
#else // !MEDIAPIPE_DISABLE_GPU
#if MEDIAPIPE_METAL_ENABLED
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#else
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // MEDIAPIPE_METAL_ENABLED
#endif // MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
mediapipe::Status Open(CalculatorContext* cc) {
absl::Status Open(CalculatorContext* cc) {
options_ = cc->Options<mediapipe::ImageToTensorCalculatorOptions>();
output_width_ = options_.output_tensor_width();
output_height_ = options_.output_tensor_height();
range_min_ = options_.output_tensor_float_range().min();
range_max_ = options_.output_tensor_float_range().max();
if (kInCpu(cc).IsConnected()) {
ASSIGN_OR_RETURN(converter_, CreateOpenCvConverter(cc, GetBorderMode()));
} else {
#if MEDIAPIPE_DISABLE_GPU
return mediapipe::UnimplementedError("GPU processing is disabled");
#else
#if MEDIAPIPE_METAL_ENABLED
ASSIGN_OR_RETURN(converter_, CreateMetalConverter(cc, GetBorderMode()));
#elif MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
ASSIGN_OR_RETURN(converter_,
CreateImageToGlBufferTensorConverter(
cc, DoesInputStartAtBottom(), GetBorderMode()));
#else
ASSIGN_OR_RETURN(converter_,
CreateImageToGlTextureTensorConverter(
cc, DoesInputStartAtBottom(), GetBorderMode()));
#endif // MEDIAPIPE_METAL_ENABLED
#endif // MEDIAPIPE_DISABLE_GPU
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status Process(CalculatorContext* cc) {
const PacketBase& image_packet =
kInCpu(cc).IsConnected() ? kInCpu(cc).packet() : kInGpu(cc).packet();
if (image_packet.IsEmpty()) {
// Timestamp bound update happens automatically. (See Open().)
return mediapipe::OkStatus();
absl::Status Process(CalculatorContext* cc) {
if ((kIn(cc).IsConnected() && kIn(cc).IsEmpty()) ||
(kInGpu(cc).IsConnected() && kInGpu(cc).IsEmpty())) {
// Timestamp bound update happens automatically.
return absl::OkStatus();
}
absl::optional<mediapipe::NormalizedRect> norm_rect;
if (kInNormRect(cc).IsConnected()) {
if (kInNormRect(cc).IsEmpty()) {
// Timestamp bound update happens automatically. (See Open().)
return mediapipe::OkStatus();
return absl::OkStatus();
}
norm_rect = *kInNormRect(cc);
if (norm_rect->width() == 0 && norm_rect->height() == 0) {
@@ -211,11 +199,12 @@ class ImageToTensorCalculator : public Node {
// NOTE: usage of sentinel rects should be avoided.
DLOG(WARNING)
<< "Updating timestamp bound in response to a sentinel rect";
return mediapipe::OkStatus();
return absl::OkStatus();
}
}
const Size& size = converter_->GetImageSize(image_packet);
ASSIGN_OR_RETURN(auto image, GetInputImage(cc));
const Size size{image->width(), image->height()};
RotatedRect roi = GetRoi(size.width, size.height, norm_rect);
ASSIGN_OR_RETURN(auto padding, PadRoi(options_.output_tensor_width(),
options_.output_tensor_height(),
@@ -231,16 +220,19 @@ class ImageToTensorCalculator : public Node {
kOutMatrix(cc).Send(std::move(matrix));
}
ASSIGN_OR_RETURN(
Tensor tensor,
converter_->Convert(image_packet, roi, {output_width_, output_height_},
range_min_, range_max_));
// Lazy initialization of the GPU or CPU converter.
MP_RETURN_IF_ERROR(InitConverterIfNecessary(cc, image->UsesGpu()));
ASSIGN_OR_RETURN(Tensor tensor,
(image->UsesGpu() ? gpu_converter_ : cpu_converter_)
->Convert(*image, roi, {output_width_, output_height_},
range_min_, range_max_));
auto result = std::make_unique<std::vector<Tensor>>();
result->push_back(std::move(tensor));
kOutTensors(cc).Send(std::move(result));
return mediapipe::OkStatus();
return absl::OkStatus();
}
private:
@@ -261,7 +253,62 @@ class ImageToTensorCalculator : public Node {
}
}
std::unique_ptr<ImageToTensorConverter> converter_;
absl::StatusOr<std::shared_ptr<const mediapipe::Image>> GetInputImage(
CalculatorContext* cc) {
if (kIn(cc).IsConnected()) {
const auto& packet = kIn(cc).packet();
return kIn(cc).Visit(
[&packet](const mediapipe::Image&) {
return SharedPtrWithPacket<mediapipe::Image>(packet);
},
[&packet](const mediapipe::ImageFrame&) {
return std::make_shared<const mediapipe::Image>(
std::const_pointer_cast<mediapipe::ImageFrame>(
SharedPtrWithPacket<mediapipe::ImageFrame>(packet)));
});
} else { // if (kInGpu(cc).IsConnected())
#if !MEDIAPIPE_DISABLE_GPU
const GpuBuffer& input = *kInGpu(cc);
// A shallow copy is okay since the resulting 'image' object is local in
// Process(), and thus never outlives 'input'.
return std::make_shared<const mediapipe::Image>(input);
#else
return absl::UnimplementedError(
"GPU processing is disabled in build flags");
#endif // !MEDIAPIPE_DISABLE_GPU
}
}
absl::Status InitConverterIfNecessary(CalculatorContext* cc, bool use_gpu) {
// Lazy initialization of the GPU or CPU converter.
if (use_gpu) {
if (!gpu_converter_) {
#if !MEDIAPIPE_DISABLE_GPU
#if MEDIAPIPE_METAL_ENABLED
ASSIGN_OR_RETURN(gpu_converter_,
CreateMetalConverter(cc, GetBorderMode()));
#elif MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
ASSIGN_OR_RETURN(gpu_converter_,
CreateImageToGlBufferTensorConverter(
cc, DoesInputStartAtBottom(), GetBorderMode()));
#else
ASSIGN_OR_RETURN(gpu_converter_,
CreateImageToGlTextureTensorConverter(
cc, DoesInputStartAtBottom(), GetBorderMode()));
#endif // MEDIAPIPE_METAL_ENABLED
#endif // !MEDIAPIPE_DISABLE_GPU
}
} else {
if (!cpu_converter_) {
ASSIGN_OR_RETURN(cpu_converter_,
CreateOpenCvConverter(cc, GetBorderMode()));
}
}
return absl::OkStatus();
}
std::unique_ptr<ImageToTensorConverter> gpu_converter_;
std::unique_ptr<ImageToTensorConverter> cpu_converter_;
mediapipe::ImageToTensorCalculatorOptions options_;
int output_width_ = 0;
int output_height_ = 0;
mediapipe/calculators/tensor/image_to_tensor_calculator_test.cc
+43 -12
View File
@@ -22,11 +22,13 @@
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_runner.h"
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_frame_opencv.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/commandlineflags.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/opencv_core_inc.h"
@@ -54,10 +56,12 @@ cv::Mat GetRgba(absl::string_view path) {
// Image to tensor test template.
// No processing/assertions should be done after the function is invoked.
void RunTest(cv::Mat input, cv::Mat expected_result, float range_min,
float range_max, int tensor_width, int tensor_height,
bool keep_aspect, absl::optional<BorderMode> border_mode,
const mediapipe::NormalizedRect& roi) {
void RunTestWithInputImagePacket(const Packet& input_image_packet,
cv::Mat expected_result, float range_min,
float range_max, int tensor_width,
int tensor_height, bool keep_aspect,
absl::optional<BorderMode> border_mode,
const mediapipe::NormalizedRect& roi) {
std::string border_mode_str;
if (border_mode) {
switch (*border_mode) {
@@ -107,12 +111,8 @@ void RunTest(cv::Mat input, cv::Mat expected_result, float range_min,
MP_ASSERT_OK(graph.Initialize(graph_config));
MP_ASSERT_OK(graph.StartRun({}));
ImageFrame input_image(
input.channels() == 4 ? ImageFormat::SRGBA : ImageFormat::SRGB,
input.cols, input.rows, input.step, input.data, [](uint8*) {});
MP_ASSERT_OK(graph.AddPacketToInputStream(
"input_image",
MakePacket<ImageFrame>(std::move(input_image)).At(Timestamp(0))));
MP_ASSERT_OK(graph.AddPacketToInputStream("input_image", input_image_packet));
MP_ASSERT_OK(graph.AddPacketToInputStream(
"roi",
MakePacket<mediapipe::NormalizedRect>(std::move(roi)).At(Timestamp(0))));
@@ -133,8 +133,7 @@ void RunTest(cv::Mat input, cv::Mat expected_result, float range_min,
const_cast<float*>(view.buffer<float>()));
cv::Mat result_rgb;
auto transformation =
GetValueRangeTransformation(range_min, range_max, 0.0f, 255.0f)
.ValueOrDie();
GetValueRangeTransformation(range_min, range_max, 0.0f, 255.0f).value();
tensor_mat.convertTo(result_rgb, CV_8UC3, transformation.scale,
transformation.offset);
@@ -152,6 +151,38 @@ void RunTest(cv::Mat input, cv::Mat expected_result, float range_min,
MP_ASSERT_OK(graph.WaitUntilDone());
}
Packet MakeImageFramePacket(cv::Mat input) {
ImageFrame input_image(
input.channels() == 4 ? ImageFormat::SRGBA : ImageFormat::SRGB,
input.cols, input.rows, input.step, input.data, [](uint8*) {});
return MakePacket<ImageFrame>(std::move(input_image)).At(Timestamp(0));
}
Packet MakeImagePacket(cv::Mat input) {
mediapipe::Image input_image(std::make_shared<mediapipe::ImageFrame>(
input.channels() == 4 ? ImageFormat::SRGBA : ImageFormat::SRGB,
input.cols, input.rows, input.step, input.data, [](uint8*) {}));
return MakePacket<mediapipe::Image>(std::move(input_image)).At(Timestamp(0));
}
enum class InputType { kImageFrame, kImage };
const std::vector<InputType> kInputTypesToTest = {InputType::kImageFrame,
InputType::kImage};
void RunTest(cv::Mat input, cv::Mat expected_result, float range_min,
float range_max, int tensor_width, int tensor_height,
bool keep_aspect, absl::optional<BorderMode> border_mode,
const mediapipe::NormalizedRect& roi) {
for (auto input_type : kInputTypesToTest) {
RunTestWithInputImagePacket(
input_type == InputType::kImageFrame ? MakeImageFramePacket(input)
: MakeImagePacket(input),
expected_result, range_min, range_max, tensor_width, tensor_height,
keep_aspect, border_mode, roi);
}
}
TEST(ImageToTensorCalculatorTest, MediumSubRectKeepAspect) {
mediapipe::NormalizedRect roi;
roi.set_x_center(0.65f);
mediapipe/calculators/tensor/image_to_tensor_converter.h
+6 -9
View File
@@ -16,8 +16,8 @@
#define MEDIAPIPE_CALCULATORS_TENSOR_IMAGE_TO_TENSOR_CONVERTER_H_
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/statusor.h"
namespace mediapipe {
@@ -38,20 +38,17 @@ class ImageToTensorConverter {
public:
virtual ~ImageToTensorConverter() = default;
virtual Size GetImageSize(const Packet& image_packet) = 0;
// Converts image to tensor.
// @image_packet contains image to extract from.
// @image contains image to extract from.
// @roi describes region of interest within the image to extract (absolute
// values).
// @output_dims dimensions of output tensor.
// @range_min/max describes output tensor range image pixels should converted
// to.
virtual mediapipe::StatusOr<Tensor> Convert(const Packet& image_packet,
const RotatedRect& roi,
const Size& output_dims,
float range_min,
float range_max) = 0;
virtual absl::StatusOr<Tensor> Convert(const mediapipe::Image& input,
const RotatedRect& roi,
const Size& output_dims,
float range_min, float range_max) = 0;
};
} // namespace mediapipe
mediapipe/calculators/tensor/image_to_tensor_converter_gl_buffer.cc
+47 -54
View File
@@ -27,13 +27,13 @@
#include "mediapipe/calculators/tensor/image_to_tensor_converter_gl_utils.h"
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/common/types.h"
#include "tensorflow/lite/delegates/gpu/gl/command_queue.h"
@@ -54,7 +54,7 @@ class SubRectExtractorGl {
public:
// Extracts a region defined by @sub_rect, removes A channel, transforms input
// pixels as alpha * x + beta and resizes result into destination.
mediapipe::Status ExtractSubRectToBuffer(
absl::Status ExtractSubRectToBuffer(
const tflite::gpu::gl::GlTexture& texture,
const tflite::gpu::HW& texture_size, const RotatedRect& sub_rect,
bool flip_horizontaly, float alpha, float beta,
@@ -62,7 +62,7 @@ class SubRectExtractorGl {
tflite::gpu::gl::CommandQueue* command_queue,
tflite::gpu::gl::GlBuffer* destination);
static mediapipe::StatusOr<SubRectExtractorGl> Create(
static absl::StatusOr<SubRectExtractorGl> Create(
const mediapipe::GlContext& gl_context, bool input_starts_at_bottom,
BorderMode border_mode);
@@ -82,8 +82,8 @@ class SubRectExtractorGl {
BorderMode border_mode_ = BorderMode::kReplicate;
};
mediapipe::Status SetMat4x4(const tflite::gpu::gl::GlProgram& program,
const std::string& name, float* data) {
absl::Status SetMat4x4(const tflite::gpu::gl::GlProgram& program,
const std::string& name, float* data) {
GLint uniform_id;
MP_RETURN_IF_ERROR(TFLITE_GPU_CALL_GL(glGetUniformLocation, &uniform_id,
program.id(), name.c_str()));
@@ -151,7 +151,7 @@ void main() {
}
)";
mediapipe::Status SubRectExtractorGl::ExtractSubRectToBuffer(
absl::Status SubRectExtractorGl::ExtractSubRectToBuffer(
const tflite::gpu::gl::GlTexture& texture,
const tflite::gpu::HW& texture_size, const RotatedRect& texture_sub_rect,
bool flip_horizontaly, float alpha, float beta,
@@ -205,10 +205,10 @@ mediapipe::Status SubRectExtractorGl::ExtractSubRectToBuffer(
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::StatusOr<SubRectExtractorGl> SubRectExtractorGl::Create(
absl::StatusOr<SubRectExtractorGl> SubRectExtractorGl::Create(
const mediapipe::GlContext& gl_context, bool input_starts_at_bottom,
BorderMode border_mode) {
bool use_custom_zero_border = border_mode == BorderMode::kZero &&
@@ -244,11 +244,11 @@ mediapipe::StatusOr<SubRectExtractorGl> SubRectExtractorGl::Create(
class GlProcessor : public ImageToTensorConverter {
public:
mediapipe::Status Init(CalculatorContext* cc, bool input_starts_at_bottom,
BorderMode border_mode) {
absl::Status Init(CalculatorContext* cc, bool input_starts_at_bottom,
BorderMode border_mode) {
MP_RETURN_IF_ERROR(gl_helper_.Open(cc));
return gl_helper_.RunInGlContext([this, input_starts_at_bottom,
border_mode]() -> mediapipe::Status {
border_mode]() -> absl::Status {
tflite::gpu::GpuInfo gpu_info;
MP_RETURN_IF_ERROR(tflite::gpu::gl::RequestGpuInfo(&gpu_info));
RET_CHECK(gpu_info.IsApiOpenGl31OrAbove())
@@ -260,20 +260,14 @@ class GlProcessor : public ImageToTensorConverter {
SubRectExtractorGl::Create(gl_helper_.GetGlContext(),
input_starts_at_bottom, border_mode));
extractor_ = absl::make_unique<SubRectExtractorGl>(std::move(extractor));
return mediapipe::OkStatus();
return absl::OkStatus();
});
}
Size GetImageSize(const Packet& image_packet) override {
const auto& image = image_packet.Get<mediapipe::GpuBuffer>();
return {image.width(), image.height()};
}
mediapipe::StatusOr<Tensor> Convert(const Packet& image_packet,
const RotatedRect& roi,
const Size& output_dims, float range_min,
float range_max) override {
const auto& input = image_packet.Get<mediapipe::GpuBuffer>();
absl::StatusOr<Tensor> Convert(const mediapipe::Image& input,
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: ",
@@ -284,40 +278,39 @@ class GlProcessor : public ImageToTensorConverter {
Tensor tensor(Tensor::ElementType::kFloat32,
{1, output_dims.height, output_dims.width, kNumChannels});
MP_RETURN_IF_ERROR(gl_helper_.RunInGlContext(
[this, &tensor, &input, &roi, &output_dims, range_min,
range_max]() -> mediapipe::Status {
constexpr int kRgbaNumChannels = 4;
auto source_texture = gl_helper_.CreateSourceTexture(input);
tflite::gpu::gl::GlTexture input_texture(
GL_TEXTURE_2D, source_texture.name(), GL_RGBA,
source_texture.width() * source_texture.height() *
kRgbaNumChannels * sizeof(uint8_t),
/*layer=*/0,
/*owned=*/false);
MP_RETURN_IF_ERROR(gl_helper_.RunInGlContext([this, &tensor, &input, &roi,
&output_dims, range_min,
range_max]() -> absl::Status {
constexpr int kRgbaNumChannels = 4;
auto source_texture = gl_helper_.CreateSourceTexture(input);
tflite::gpu::gl::GlTexture input_texture(
GL_TEXTURE_2D, source_texture.name(), GL_RGBA,
source_texture.width() * source_texture.height() * kRgbaNumChannels *
sizeof(uint8_t),
/*layer=*/0,
/*owned=*/false);
constexpr float kInputImageRangeMin = 0.0f;
constexpr float kInputImageRangeMax = 1.0f;
ASSIGN_OR_RETURN(auto transform,
GetValueRangeTransformation(kInputImageRangeMin,
kInputImageRangeMax,
range_min, range_max));
constexpr float kInputImageRangeMin = 0.0f;
constexpr float kInputImageRangeMax = 1.0f;
ASSIGN_OR_RETURN(
auto transform,
GetValueRangeTransformation(kInputImageRangeMin, kInputImageRangeMax,
range_min, range_max));
auto buffer_view = tensor.GetOpenGlBufferWriteView();
tflite::gpu::gl::GlBuffer output(GL_SHADER_STORAGE_BUFFER,
buffer_view.name(), tensor.bytes(),
/*offset=*/0,
/*has_ownership=*/false);
MP_RETURN_IF_ERROR(extractor_->ExtractSubRectToBuffer(
input_texture,
tflite::gpu::HW(source_texture.height(), source_texture.width()),
roi,
/*flip_horizontaly=*/false, transform.scale, transform.offset,
tflite::gpu::HW(output_dims.height, output_dims.width),
command_queue_.get(), &output));
auto buffer_view = tensor.GetOpenGlBufferWriteView();
tflite::gpu::gl::GlBuffer output(GL_SHADER_STORAGE_BUFFER,
buffer_view.name(), tensor.bytes(),
/*offset=*/0,
/*has_ownership=*/false);
MP_RETURN_IF_ERROR(extractor_->ExtractSubRectToBuffer(
input_texture,
tflite::gpu::HW(source_texture.height(), source_texture.width()), roi,
/*flip_horizontaly=*/false, transform.scale, transform.offset,
tflite::gpu::HW(output_dims.height, output_dims.width),
command_queue_.get(), &output));
return mediapipe::OkStatus();
}));
return absl::OkStatus();
}));
return tensor;
}
@@ -338,7 +331,7 @@ class GlProcessor : public ImageToTensorConverter {
} // namespace
mediapipe::StatusOr<std::unique_ptr<ImageToTensorConverter>>
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>>
CreateImageToGlBufferTensorConverter(CalculatorContext* cc,
bool input_starts_at_bottom,
BorderMode border_mode) {
mediapipe/calculators/tensor/image_to_tensor_converter_gl_buffer.h
+1 -1
View File
@@ -30,7 +30,7 @@ namespace mediapipe {
// Creates image to tensor (represented as OpenGL buffer) converter.
// NOTE: mediapipe::GlCalculatorHelper::UpdateContract invocation must precede
// converter creation.
mediapipe::StatusOr<std::unique_ptr<ImageToTensorConverter>>
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>>
CreateImageToGlBufferTensorConverter(CalculatorContext* cc,
bool input_starts_at_bottom,
BorderMode border_mode);
mediapipe/calculators/tensor/image_to_tensor_converter_gl_texture.cc
+38 -30
View File
@@ -27,6 +27,7 @@
#include "mediapipe/calculators/tensor/image_to_tensor_converter_gl_utils.h"
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/ret_check.h"
@@ -34,7 +35,6 @@
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/shader_util.h"
namespace mediapipe {
@@ -47,11 +47,11 @@ constexpr int kNumAttributes = 2;
class GlProcessor : public ImageToTensorConverter {
public:
mediapipe::Status Init(CalculatorContext* cc, bool input_starts_at_bottom,
BorderMode border_mode) {
absl::Status Init(CalculatorContext* cc, bool input_starts_at_bottom,
BorderMode border_mode) {
MP_RETURN_IF_ERROR(gl_helper_.Open(cc));
return gl_helper_.RunInGlContext([this, input_starts_at_bottom,
border_mode]() -> mediapipe::Status {
border_mode]() -> absl::Status {
use_custom_zero_border_ =
border_mode == BorderMode::kZero &&
!IsGlClampToBorderSupported(gl_helper_.GetGlContext());
@@ -164,20 +164,14 @@ class GlProcessor : public ImageToTensorConverter {
glBindBuffer(GL_ARRAY_BUFFER, 0);
return mediapipe::OkStatus();
return absl::OkStatus();
});
}
Size GetImageSize(const Packet& image_packet) override {
const auto& image = image_packet.Get<mediapipe::GpuBuffer>();
return {image.width(), image.height()};
}
mediapipe::StatusOr<Tensor> Convert(const Packet& image_packet,
const RotatedRect& roi,
const Size& output_dims, float range_min,
float range_max) override {
const auto& input = image_packet.Get<mediapipe::GpuBuffer>();
absl::StatusOr<Tensor> Convert(const mediapipe::Image& input,
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: ",
@@ -189,9 +183,9 @@ class GlProcessor : public ImageToTensorConverter {
Tensor::ElementType::kFloat32,
Tensor::Shape{1, output_dims.height, output_dims.width, kNumChannels});
MP_RETURN_IF_ERROR(gl_helper_.RunInGlContext(
[this, &tensor, &input, &roi, &output_dims, range_min,
range_max]() -> mediapipe::Status {
MP_RETURN_IF_ERROR(
gl_helper_.RunInGlContext([this, &tensor, &input, &roi, &output_dims,
range_min, range_max]() -> absl::Status {
auto input_texture = gl_helper_.CreateSourceTexture(input);
constexpr float kInputImageRangeMin = 0.0f;
@@ -205,21 +199,18 @@ class GlProcessor : public ImageToTensorConverter {
/*flip_horizontaly=*/false,
transform.scale, transform.offset,
output_dims, &tensor_view));
return mediapipe::OkStatus();
return absl::OkStatus();
}));
return tensor;
}
mediapipe::Status ExtractSubRect(const mediapipe::GlTexture& texture,
const RotatedRect& sub_rect,
bool flip_horizontaly, float alpha,
float beta, const Size& output_dims,
Tensor::OpenGlTexture2dView* output) {
absl::Status ExtractSubRect(const mediapipe::GlTexture& texture,
const RotatedRect& sub_rect,
bool flip_horizontaly, float alpha, float beta,
const Size& output_dims,
Tensor::OpenGlTexture2dView* output) {
std::array<float, 16> transform_mat;
GetRotatedSubRectToRectTransformMatrix(sub_rect, texture.width(),
texture.height(), flip_horizontaly,
&transform_mat);
glDisable(GL_DEPTH_TEST);
glBindFramebuffer(GL_FRAMEBUFFER, framebuffer_);
@@ -258,7 +249,24 @@ class GlProcessor : public ImageToTensorConverter {
glUseProgram(program_);
glUniform1f(alpha_id_, alpha);
glUniform1f(beta_id_, beta);
glUniformMatrix4fv(matrix_id_, 1, GL_TRUE, transform_mat.data());
// If our context is ES2, then we must use GL_FALSE for our 'transpose'
// GLboolean in glUniformMatrix4fv, or else we'll get an INVALID_VALUE
// error. So in that case, we'll grab the transpose of our original matrix
// and send that instead.
const auto gl_context = mediapipe::GlContext::GetCurrent();
LOG_IF(FATAL, !gl_context) << "GlContext is not bound to the thread.";
if (gl_context->GetGlVersion() == mediapipe::GlVersion::kGLES2) {
GetTransposedRotatedSubRectToRectTransformMatrix(
sub_rect, texture.width(), texture.height(), flip_horizontaly,
&transform_mat);
glUniformMatrix4fv(matrix_id_, 1, GL_FALSE, transform_mat.data());
} else {
GetRotatedSubRectToRectTransformMatrix(sub_rect, texture.width(),
texture.height(), flip_horizontaly,
&transform_mat);
glUniformMatrix4fv(matrix_id_, 1, GL_TRUE, transform_mat.data());
}
// vao
glBindVertexArray(vao_);
@@ -292,7 +300,7 @@ class GlProcessor : public ImageToTensorConverter {
glActiveTexture(GL_TEXTURE0);
glBindTexture(GL_TEXTURE_2D, 0);
return mediapipe::OkStatus();
return absl::OkStatus();
}
~GlProcessor() override {
@@ -320,7 +328,7 @@ class GlProcessor : public ImageToTensorConverter {
} // namespace
mediapipe::StatusOr<std::unique_ptr<ImageToTensorConverter>>
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>>
CreateImageToGlTextureTensorConverter(CalculatorContext* cc,
bool input_starts_at_bottom,
BorderMode border_mode) {
mediapipe/calculators/tensor/image_to_tensor_converter_gl_texture.h
+1 -1
View File
@@ -30,7 +30,7 @@ namespace mediapipe {
// Creates image to tensor (represented as OpenGL texture) converter.
// NOTE: mediapipe::GlCalculatorHelper::UpdateContract invocation must precede
// converter creation.
mediapipe::StatusOr<std::unique_ptr<ImageToTensorConverter>>
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>>
CreateImageToGlTextureTensorConverter(CalculatorContext* cc,
bool input_starts_at_bottom,
BorderMode border_mode);
mediapipe/calculators/tensor/image_to_tensor_converter_gl_utils_test.cc
+1 -1
View File
@@ -14,7 +14,7 @@ namespace {
TEST(ImageToTensorConverterGlUtilsTest, GlTexParameteriOverrider) {
auto status_or_context = mediapipe::GlContext::Create(nullptr, false);
MP_ASSERT_OK(status_or_context);
auto context = status_or_context.ValueOrDie();
auto context = status_or_context.value();
std::vector<GLint> min_filter_changes;
context->Run([&min_filter_changes]() {
mediapipe/calculators/tensor/image_to_tensor_converter_metal.cc
+27 -33
View File
@@ -26,13 +26,13 @@
#include "mediapipe/calculators/tensor/image_to_tensor_converter.h"
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/gpu/MPPMetalHelper.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/common/types.h"
@@ -146,7 +146,7 @@ int GetBytesPerRaw(OutputFormat output_format, const tflite::gpu::HW& size) {
class SubRectExtractorMetal {
public:
static mediapipe::StatusOr<std::unique_ptr<SubRectExtractorMetal>> Make(
static absl::StatusOr<std::unique_ptr<SubRectExtractorMetal>> Make(
id<MTLDevice> device, OutputFormat output_format,
BorderMode border_mode) {
id<MTLRenderPipelineState> pipeline_state;
@@ -174,12 +174,12 @@ class SubRectExtractorMetal {
options:MTLResourceOptionCPUCacheModeDefault];
}
mediapipe::Status Execute(id<MTLTexture> input_texture,
const RotatedRect& sub_rect, bool flip_horizontaly,
float alpha, float beta,
const tflite::gpu::HW& destination_size,
id<MTLCommandBuffer> command_buffer,
id<MTLBuffer> destination) {
absl::Status Execute(id<MTLTexture> input_texture,
const RotatedRect& sub_rect, bool flip_horizontaly,
float alpha, float beta,
const tflite::gpu::HW& destination_size,
id<MTLCommandBuffer> command_buffer,
id<MTLBuffer> destination) {
auto output_texture = MTLTextureWithBuffer(destination_size, destination);
return InternalExecute(input_texture, sub_rect, flip_horizontaly, alpha,
beta, destination_size, command_buffer,
@@ -205,13 +205,12 @@ class SubRectExtractorMetal {
return texture;
}
mediapipe::Status InternalExecute(id<MTLTexture> input_texture,
const RotatedRect& sub_rect,
bool flip_horizontaly, float alpha,
float beta,
const tflite::gpu::HW& destination_size,
id<MTLCommandBuffer> command_buffer,
id<MTLTexture> output_texture) {
absl::Status InternalExecute(id<MTLTexture> input_texture,
const RotatedRect& sub_rect,
bool flip_horizontaly, float alpha, float beta,
const tflite::gpu::HW& destination_size,
id<MTLCommandBuffer> command_buffer,
id<MTLTexture> output_texture) {
RET_CHECK(command_buffer != nil);
RET_CHECK(output_texture != nil);
@@ -254,10 +253,10 @@ class SubRectExtractorMetal {
vertexCount:6];
[command_encoder endEncoding];
return mediapipe::OkStatus();
return absl::OkStatus();
}
static mediapipe::Status MakePipelineState(
static absl::Status MakePipelineState(
id<MTLDevice> device, OutputFormat output_format, BorderMode border_mode,
id<MTLRenderPipelineState>* pipeline_state) {
RET_CHECK(pipeline_state != nil);
@@ -328,7 +327,7 @@ class SubRectExtractorMetal {
RET_CHECK(error == nil) << "Couldn't create a pipeline state"
<< [[error localizedDescription] UTF8String];
return mediapipe::OkStatus();
return absl::OkStatus();
}
id<MTLBuffer> positions_buffer_;
@@ -340,25 +339,19 @@ class SubRectExtractorMetal {
class MetalProcessor : public ImageToTensorConverter {
public:
mediapipe::Status Init(CalculatorContext* cc, BorderMode border_mode) {
absl::Status Init(CalculatorContext* cc, BorderMode border_mode) {
metal_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(metal_helper_);
ASSIGN_OR_RETURN(extractor_, SubRectExtractorMetal::Make(
metal_helper_.mtlDevice,
OutputFormat::kF32C4, border_mode));
return mediapipe::OkStatus();
return absl::OkStatus();
}
Size GetImageSize(const Packet& image_packet) override {
const auto& image = image_packet.Get<mediapipe::GpuBuffer>();
return {image.width(), image.height()};
}
mediapipe::StatusOr<Tensor> Convert(const Packet& image_packet,
const RotatedRect& roi,
const Size& output_dims, float range_min,
float range_max) override {
const auto& input = image_packet.Get<mediapipe::GpuBuffer>();
absl::StatusOr<Tensor> Convert(const mediapipe::Image& input,
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 "
@@ -367,7 +360,8 @@ class MetalProcessor : public ImageToTensorConverter {
}
@autoreleasepool {
id<MTLTexture> texture = [metal_helper_ metalTextureWithGpuBuffer:input];
id<MTLTexture> texture =
[metal_helper_ metalTextureWithGpuBuffer:input.GetGpuBuffer()];
constexpr int kNumChannels = 4;
Tensor tensor(Tensor::ElementType::kFloat32,
@@ -400,8 +394,8 @@ class MetalProcessor : public ImageToTensorConverter {
} // namespace
mediapipe::StatusOr<std::unique_ptr<ImageToTensorConverter>>
CreateMetalConverter(CalculatorContext* cc, BorderMode border_mode) {
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>> CreateMetalConverter(
CalculatorContext* cc, BorderMode border_mode) {
auto result = absl::make_unique<MetalProcessor>();
MP_RETURN_IF_ERROR(result->Init(cc, border_mode));
mediapipe/calculators/tensor/image_to_tensor_converter_metal.h
+2 -2
View File
@@ -30,8 +30,8 @@ namespace mediapipe {
// Creates Metal image-to-tensor converter.
// NOTE: [MPPMetalHelper updateContract:...] invocation must precede
// converter creation.
mediapipe::StatusOr<std::unique_ptr<ImageToTensorConverter>>
CreateMetalConverter(CalculatorContext* cc, BorderMode border_mode);
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>> CreateMetalConverter(
CalculatorContext* cc, BorderMode border_mode);
} // namespace mediapipe
mediapipe/calculators/tensor/image_to_tensor_converter_opencv.cc
+11 -17
View File
@@ -20,9 +20,9 @@
#include "mediapipe/calculators/tensor/image_to_tensor_converter.h"
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_frame_opencv.h"
#include "mediapipe/framework/formats/image_opencv.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/opencv_core_inc.h"
@@ -46,21 +46,15 @@ class OpenCvProcessor : public ImageToTensorConverter {
}
}
Size GetImageSize(const Packet& image_packet) override {
const auto& image = image_packet.Get<mediapipe::ImageFrame>();
return {image.Width(), image.Height()};
}
mediapipe::StatusOr<Tensor> Convert(const Packet& image_packet,
const RotatedRect& roi,
const Size& output_dims, float range_min,
float range_max) override {
const auto& input = image_packet.Get<mediapipe::ImageFrame>();
if (input.Format() != mediapipe::ImageFormat::SRGB &&
input.Format() != mediapipe::ImageFormat::SRGBA) {
absl::StatusOr<Tensor> Convert(const mediapipe::Image& input,
const RotatedRect& roi,
const Size& output_dims, float range_min,
float range_max) override {
if (input.image_format() != mediapipe::ImageFormat::SRGB &&
input.image_format() != mediapipe::ImageFormat::SRGBA) {
return InvalidArgumentError(
absl::StrCat("Only RGBA/RGB formats are supported, passed format: ",
static_cast<uint32_t>(input.Format())));
static_cast<uint32_t>(input.image_format())));
}
cv::Mat src = mediapipe::formats::MatView(&input);
@@ -118,8 +112,8 @@ class OpenCvProcessor : public ImageToTensorConverter {
} // namespace
mediapipe::StatusOr<std::unique_ptr<ImageToTensorConverter>>
CreateOpenCvConverter(CalculatorContext* cc, BorderMode border_mode) {
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>> CreateOpenCvConverter(
CalculatorContext* cc, BorderMode border_mode) {
// Simply "return absl::make_unique<OpenCvProcessor>()" failed to build on
// macOS with bazel.
return std::unique_ptr<ImageToTensorConverter>(
mediapipe/calculators/tensor/image_to_tensor_converter_opencv.h
+2 -2
View File
@@ -24,8 +24,8 @@
namespace mediapipe {
// Creates OpenCV image-to-tensor converter.
mediapipe::StatusOr<std::unique_ptr<ImageToTensorConverter>>
CreateOpenCvConverter(CalculatorContext* cc, BorderMode border_mode);
absl::StatusOr<std::unique_ptr<ImageToTensorConverter>> CreateOpenCvConverter(
CalculatorContext* cc, BorderMode border_mode);
} // namespace mediapipe
mediapipe/calculators/tensor/image_to_tensor_utils.cc
+46 -5
View File
@@ -38,10 +38,10 @@ RotatedRect GetRoi(int input_width, int input_height,
/*rotation =*/0};
}
mediapipe::StatusOr<std::array<float, 4>> PadRoi(int input_tensor_width,
int input_tensor_height,
bool keep_aspect_ratio,
RotatedRect* roi) {
absl::StatusOr<std::array<float, 4>> PadRoi(int input_tensor_width,
int input_tensor_height,
bool keep_aspect_ratio,
RotatedRect* roi) {
if (!keep_aspect_ratio) {
return std::array<float, 4>{0.0f, 0.0f, 0.0f, 0.0f};
}
@@ -76,7 +76,7 @@ mediapipe::StatusOr<std::array<float, 4>> PadRoi(int input_tensor_width,
horizontal_padding, vertical_padding};
}
mediapipe::StatusOr<ValueTransformation> GetValueRangeTransformation(
absl::StatusOr<ValueTransformation> GetValueRangeTransformation(
float from_range_min, float from_range_max, float to_range_min,
float to_range_max) {
RET_CHECK_LT(from_range_min, from_range_max)
@@ -173,4 +173,45 @@ void GetRotatedSubRectToRectTransformMatrix(const RotatedRect& sub_rect,
matrix[15] = 1.0f;
}
void GetTransposedRotatedSubRectToRectTransformMatrix(
const RotatedRect& sub_rect, int rect_width, int rect_height,
bool flip_horizontaly, std::array<float, 16>* matrix_ptr) {
std::array<float, 16>& matrix = *matrix_ptr;
// See comments in GetRotatedSubRectToRectTransformMatrix for detailed
// calculations.
const float a = sub_rect.width;
const float b = sub_rect.height;
const float flip = flip_horizontaly ? -1 : 1;
const float c = std::cos(sub_rect.rotation);
const float d = std::sin(sub_rect.rotation);
const float e = sub_rect.center_x;
const float f = sub_rect.center_y;
const float g = 1.0f / rect_width;
const float h = 1.0f / rect_height;
// row 1 (indices 0,4,8,12 from non-transposed fcn)
matrix[0] = a * c * flip * g;
matrix[1] = a * d * flip * h;
matrix[2] = 0.0f;
matrix[3] = 0.0f;
// row 2 (indices 1,5,9,13 from non-transposed fcn)
matrix[4] = -b * d * g;
matrix[5] = b * c * h;
matrix[6] = 0.0f;
matrix[7] = 0.0f;
// row 3 (indices 2,6,10,14 from non-transposed fcn)
matrix[8] = 0.0f;
matrix[9] = 0.0f;
matrix[10] = a * g;
matrix[11] = 0.0f;
// row 4 (indices 3,7,11,15 from non-transposed fcn)
matrix[12] = (-0.5f * a * c * flip + 0.5f * b * d + e) * g;
matrix[13] = (-0.5f * b * c - 0.5f * a * d * flip + f) * h;
matrix[14] = 0.0f;
matrix[15] = 1.0f;
}
} // namespace mediapipe
mediapipe/calculators/tensor/image_to_tensor_utils.h
+23 -5
View File
@@ -37,10 +37,10 @@ RotatedRect GetRoi(int input_width, int input_height,
// Pads ROI, so extraction happens correctly if aspect ratio is to be kept.
// Returns letterbox padding applied.
mediapipe::StatusOr<std::array<float, 4>> PadRoi(int input_tensor_width,
int input_tensor_height,
bool keep_aspect_ratio,
RotatedRect* roi);
absl::StatusOr<std::array<float, 4>> PadRoi(int input_tensor_width,
int input_tensor_height,
bool keep_aspect_ratio,
RotatedRect* roi);
// Represents a transformation of value which involves scaling and offsetting.
// To apply transformation:
@@ -55,7 +55,7 @@ struct ValueTransformation {
// [from_range_min, from_range_max] into [to_range_min, to_range_max] range.
// from_range_min must be less than from_range_max
// to_range_min must be less than to_range_max
mediapipe::StatusOr<ValueTransformation> GetValueRangeTransformation(
absl::StatusOr<ValueTransformation> GetValueRangeTransformation(
float from_range_min, float from_range_max, float to_range_min,
float to_range_max);
@@ -77,6 +77,24 @@ void GetRotatedSubRectToRectTransformMatrix(const RotatedRect& sub_rect,
bool flip_horizontaly,
std::array<float, 16>* matrix);
// Returns the transpose of the matrix found with
// "GetRotatedSubRectToRectTransformMatrix". That is to say, this populates a
// 4x4 "matrix" with col major order transformation matrix which maps (x, y) in
// range [0, 1] (describing points of @sub_rect) to (x', y') in range [0, 1]***
// (describing points of a rect: [0, @rect_width] x [0, @rect_height] = RECT).
//
// *** (x', y') will go out of the range for points from @sub_rect
// which are not contained by RECT and it's expected behavior
//
// @sub_rect - rotated sub rect in absolute coordinates
// @rect_width - rect width
// @rect_height - rect height
// @flip_horizontaly - we need to flip the output buffer.
// @matrix - 4x4 matrix (array of 16 elements) to populate
void GetTransposedRotatedSubRectToRectTransformMatrix(
const RotatedRect& sub_rect, int rect_width, int rect_height,
bool flip_horizontaly, std::array<float, 16>* matrix);
} // namespace mediapipe
#endif // MEDIAPIPE_CALCULATORS_TENSOR_IMAGE_TO_TENSOR_UTILS_H_
mediapipe/calculators/tensor/image_to_tensor_utils_test.cc
+8 -8
View File
@@ -70,7 +70,7 @@ TEST(PadRoi, NoPadding) {
.rotation = 5};
auto status_or_value = PadRoi(10, 10, /*keep_aspect_ratio=*/false, &roi);
MP_ASSERT_OK(status_or_value);
EXPECT_THAT(status_or_value.ValueOrDie(),
EXPECT_THAT(status_or_value.value(),
ElementsAreArray({0.0f, 0.0f, 0.0f, 0.0f}));
EXPECT_THAT(roi, EqRotatedRect(100, 200, 20, 10, 5));
}
@@ -83,7 +83,7 @@ TEST(PadRoi, HorizontalPadding) {
.rotation = 5};
auto status_or_value = PadRoi(10, 10, /*keep_aspect_ratio=*/true, &roi);
MP_ASSERT_OK(status_or_value);
EXPECT_THAT(status_or_value.ValueOrDie(),
EXPECT_THAT(status_or_value.value(),
ElementsAreArray({0.25f, 0.0f, 0.25f, 0.0f}));
EXPECT_THAT(roi, EqRotatedRect(200, 200, 20, 10, 5));
}
@@ -95,7 +95,7 @@ TEST(PadRoi, VerticalPadding) {
auto status_or_value = PadRoi(10, 10, /*keep_aspect_ratio=*/true, &roi);
MP_ASSERT_OK(status_or_value);
EXPECT_THAT(
status_or_value.ValueOrDie(),
status_or_value.value(),
ElementsAre(testing::FloatEq(0.0f),
testing::FloatNear(expected_horizontal_padding, 1e-6),
testing::FloatEq(0.0f),
@@ -115,7 +115,7 @@ TEST(GetValueRangeTransformation, PixelToFloatZeroCenter) {
/*from_range_min=*/0.0f, /*from_range_max=*/255.0f,
/*to_range_min=*/-1.0f, /*to_range_max=*/1.0f);
MP_ASSERT_OK(status_or_value);
EXPECT_THAT(status_or_value.ValueOrDie(),
EXPECT_THAT(status_or_value.value(),
EqValueTransformation(/*scale=*/2 / 255.0f,
/*offset=*/-1.0f));
}
@@ -125,7 +125,7 @@ TEST(GetValueRangeTransformation, PixelToFloat) {
/*from_range_min=*/0.0f, /*from_range_max=*/255.0f,
/*to_range_min=*/0.0f, /*to_range_max=*/1.0f);
MP_ASSERT_OK(status_or_value);
EXPECT_THAT(status_or_value.ValueOrDie(),
EXPECT_THAT(status_or_value.value(),
EqValueTransformation(/*scale=*/1 / 255.0f,
/*offset=*/0.0f));
}
@@ -135,7 +135,7 @@ TEST(GetValueRangeTransformation, FloatToFloatNoOp) {
/*from_range_min=*/0.0f, /*from_range_max=*/1.0f,
/*to_range_min=*/0.0f, /*to_range_max=*/1.0f);
MP_ASSERT_OK(status_or_value);
EXPECT_THAT(status_or_value.ValueOrDie(),
EXPECT_THAT(status_or_value.value(),
EqValueTransformation(/*scale=*/1.0f, /*offset=*/0.0f));
}
@@ -144,7 +144,7 @@ TEST(GetValueRangeTransformation, PixelToPixelNoOp) {
/*from_range_min=*/0.0f, /*from_range_max=*/255.0f,
/*to_range_min=*/0.0f, /*to_range_max=*/255.0f);
MP_ASSERT_OK(status_or_value);
EXPECT_THAT(status_or_value.ValueOrDie(),
EXPECT_THAT(status_or_value.value(),
EqValueTransformation(/*scale=*/1.0f, /*offset=*/0.0f));
}
@@ -153,7 +153,7 @@ TEST(GetValueRangeTransformation, FloatToPixel) {
/*from_range_min=*/0.0f, /*from_range_max=*/1.0f,
/*to_range_min=*/0.0f, /*to_range_max=*/255.0f);
MP_ASSERT_OK(status_or_value);
EXPECT_THAT(status_or_value.ValueOrDie(),
EXPECT_THAT(status_or_value.value(),
EqValueTransformation(/*scale=*/255.0f, /*offset=*/0.0f));
}
mediapipe/calculators/tensor/inference_calculator.cc
+36 -799
View File
@@ -12,822 +12,59 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include "mediapipe/calculators/tensor/inference_calculator.h"
#include <cstring>
#include <memory>
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "mediapipe/calculators/tensor/inference_calculator.pb.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/tflite/config.h"
#include "mediapipe/util/tflite/tflite_model_loader.h"
#if !defined(__EMSCRIPTEN__) || defined(__EMSCRIPTEN_PTHREADS__)
#include "mediapipe/util/cpu_util.h"
#endif // !__EMSCRIPTEN__ || __EMSCRIPTEN_PTHREADS__
#include "tensorflow/lite/error_reporter.h"
#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/model.h"
#if defined(MEDIAPIPE_ANDROID)
#include "mediapipe/util/android/file/base/file.h"
#include "mediapipe/util/android/file/base/filesystem.h"
#include "mediapipe/util/android/file/base/helpers.h"
#endif // ANDROID
#if MEDIAPIPE_TFLITE_GL_INFERENCE
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/util/tflite/tflite_gpu_runner.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
#if MEDIAPIPE_TFLITE_METAL_INFERENCE
#import <CoreVideo/CoreVideo.h>
#import <Metal/Metal.h>
#import <MetalKit/MetalKit.h>
#import "mediapipe/gpu/MPPMetalHelper.h"
#include "mediapipe/gpu/MPPMetalUtil.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/metal/buffer_convert.h"
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
#include "tensorflow/lite/delegates/gpu/metal_delegate_internal.h"
#endif // MEDIAPIPE_TFLITE_METAL_INFERENCE
#if !defined(MEDIAPIPE_EDGE_TPU)
#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
#endif // !EDGETPU
#if defined(MEDIAPIPE_ANDROID)
#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
#endif // ANDROID
namespace {
// Commonly used to compute the number of blocks to launch in a kernel.
int NumGroups(const int size, const int group_size) { // NOLINT
return (size + group_size - 1) / group_size;
}
// Round up n to next multiple of m.
template <typename T>
T RoundUp(T n, T m) {
return ((n + m - T{1}) / m) * m;
}
bool ShouldUseGpu(const mediapipe::InferenceCalculatorOptions& options) {
return (
!options.has_delegate() || // Use GPU delegate if delegate not specified
(options.has_delegate() && options.delegate().has_gpu()));
}
constexpr char kTensorsTag[] = "TENSORS";
#if defined(MEDIAPIPE_EDGE_TPU)
#include "edgetpu.h"
// Creates and returns an Edge TPU interpreter to run the given edgetpu model.
std::unique_ptr<tflite::Interpreter> BuildEdgeTpuInterpreter(
const tflite::FlatBufferModel& model,
tflite::ops::builtin::BuiltinOpResolver* resolver,
edgetpu::EdgeTpuContext* edgetpu_context) {
resolver->AddCustom(edgetpu::kCustomOp, edgetpu::RegisterCustomOp());
std::unique_ptr<tflite::Interpreter> interpreter;
if (tflite::InterpreterBuilder(model, *resolver)(&interpreter) != kTfLiteOk) {
std::cerr << "Failed to build edge TPU interpreter." << std::endl;
}
interpreter->SetExternalContext(kTfLiteEdgeTpuContext, edgetpu_context);
interpreter->SetNumThreads(1);
if (interpreter->AllocateTensors() != kTfLiteOk) {
std::cerr << "Failed to allocate edge TPU tensors." << std::endl;
}
return interpreter;
}
#endif // MEDIAPIPE_EDGE_TPU
} // namespace
#include "absl/strings/string_view.h"
#include "mediapipe/framework/tool/subgraph_expansion.h"
namespace mediapipe {
namespace api2 {
#if MEDIAPIPE_TFLITE_METAL_INFERENCE
namespace {
tflite::gpu::BHWC BhwcFromTensorShape(const Tensor::Shape& shape) {
tflite::gpu::BHWC result;
result.b = shape.dims[0];
switch (shape.dims.size()) {
case 1:
// result.b is already filled.
break;
case 2:
result.h = 1;
result.w = 1;
result.c = shape.dims[1];
break;
case 3:
result.h = 1;
result.w = shape.dims[1];
result.c = shape.dims[2];
break;
case 4:
result.h = shape.dims[1];
result.w = shape.dims[2];
result.c = shape.dims[3];
break;
default:
// Handles 0 and >4.
LOG(FATAL)
<< "Dimensions size must be in range [1,4] for GPU inference, but "
<< shape.dims.size() << " is provided";
}
return result;
}
} // namespace
#endif // MEDIAPIPE_TFLITE_METAL_INFERENCE
// Returns number of threads to configure XNNPACK delegate with.
// (Equal to user provided value if specified. Otherwise, it returns number of
// high cores (hard-coded to 1 for Emscripten without Threads extension))
int GetXnnpackNumThreads(const mediapipe::InferenceCalculatorOptions& opts) {
static constexpr int kDefaultNumThreads = -1;
if (opts.has_delegate() && opts.delegate().has_xnnpack() &&
opts.delegate().xnnpack().num_threads() != kDefaultNumThreads) {
return opts.delegate().xnnpack().num_threads();
}
#if !defined(__EMSCRIPTEN__) || defined(__EMSCRIPTEN_PTHREADS__)
return InferHigherCoreIds().size();
#else
return 1;
#endif // !__EMSCRIPTEN__ || __EMSCRIPTEN_PTHREADS__
}
// Calculator Header Section
// Runs inference on the provided input Tensors and TFLite model.
//
// Creates an interpreter with given model and calls invoke().
// Optionally run inference on CPU/GPU.
//
// This calculator can be used with TensorConverterCalculator to get the
// appropriate inputs.
//
// When the input tensors are on CPU, gpu inference is optional and can be
// specified in the calculator options.
// When the input tensors are on GPU, inference is GPU and output can be CPU or
// GPU.
//
// Input:
// TENSORS - Vector of Tensors
//
// Output:
// TENSORS - Vector of Tensors
//
// Input side packet:
// CUSTOM_OP_RESOLVER (optional) - Use a custom op resolver,
// instead of the builtin one.
// MODEL (optional) - Use to specify TfLite model
// (std::unique_ptr<tflite::FlatBufferModel,
// std::function<void(tflite::FlatBufferModel*)>>)
//
// Example use:
// node {
// calculator: "InferenceCalculator"
// input_stream: "TENSORS:tensor_image"
// output_stream: "TENSORS:tensors"
// options: {
// [mediapipe.InferenceCalculatorOptions.ext] {
// model_path: "modelname.tflite"
// }
// }
// }
//
// or
//
// node {
// calculator: "InferenceCalculator"
// input_stream: "TENSORS:tensor_image"
// input_side_packet: "MODEL:model"
// output_stream: "TENSORS:tensors"
// options: {
// [mediapipe.InferenceCalculatorOptions.ext] {
// model_path: "modelname.tflite"
// delegate { gpu {} }
// }
// }
// }
//
// IMPORTANT Notes:
// Tensors are assumed to be ordered correctly (sequentially added to model).
// Input tensors are assumed to be of the correct size and already normalized.
class InferenceCalculator : public Node {
class InferenceCalculatorSelectorImpl
: public SubgraphImpl<InferenceCalculatorSelector,
InferenceCalculatorSelectorImpl> {
public:
using TfLiteDelegatePtr =
std::unique_ptr<TfLiteDelegate, std::function<void(TfLiteDelegate*)>>;
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
static constexpr SideInput<tflite::ops::builtin::BuiltinOpResolver>::Optional
kSideInCustomOpResolver{"CUSTOM_OP_RESOLVER"};
static constexpr SideInput<TfLiteModelPtr>::Optional kSideInModel{"MODEL"};
static constexpr Output<std::vector<Tensor>> kOutTensors{"TENSORS"};
MEDIAPIPE_NODE_CONTRACT(kInTensors, kSideInCustomOpResolver, kSideInModel,
kOutTensors);
static mediapipe::Status UpdateContract(CalculatorContract* cc);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
mediapipe::Status Close(CalculatorContext* cc) override;
private:
mediapipe::Status ReadKernelsFromFile();
mediapipe::Status WriteKernelsToFile();
mediapipe::Status LoadModel(CalculatorContext* cc);
mediapipe::StatusOr<mediapipe::Packet> GetModelAsPacket(
const CalculatorContext& cc);
mediapipe::Status LoadDelegate(CalculatorContext* cc);
mediapipe::Status InitTFLiteGPURunner(CalculatorContext* cc);
mediapipe::Packet model_packet_;
std::unique_ptr<tflite::Interpreter> interpreter_;
TfLiteDelegatePtr delegate_;
#if MEDIAPIPE_TFLITE_GL_INFERENCE
mediapipe::GlCalculatorHelper gpu_helper_;
std::unique_ptr<tflite::gpu::TFLiteGPURunner> tflite_gpu_runner_;
bool allow_precision_loss_ = false;
mediapipe::InferenceCalculatorOptions::Delegate::Gpu::API
tflite_gpu_runner_api_;
#elif MEDIAPIPE_TFLITE_METAL_INFERENCE
MPPMetalHelper* gpu_helper_ = nullptr;
TFLBufferConvert* converter_to_BPHWC4_ = nil;
TFLBufferConvert* converter_from_BPHWC4_ = nil;
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
#if MEDIAPIPE_TFLITE_GPU_SUPPORTED
std::vector<Tensor::Shape> output_shapes_;
std::vector<std::unique_ptr<Tensor>> gpu_buffers_in_;
std::vector<std::unique_ptr<Tensor>> gpu_buffers_out_;
#endif // MEDIAPIPE_TFLITE_GPU_SUPPORTED
#if defined(MEDIAPIPE_EDGE_TPU)
std::shared_ptr<edgetpu::EdgeTpuContext> edgetpu_context_ =
edgetpu::EdgeTpuManager::GetSingleton()->OpenDevice();
#endif
bool use_advanced_gpu_api_ = false;
bool use_gpu_delegate_ = false;
bool use_kernel_caching_ = false;
std::string cached_kernel_filename_;
absl::StatusOr<CalculatorGraphConfig> GetConfig(
const CalculatorGraphConfig::Node& subgraph_node) {
const auto& options =
Subgraph::GetOptions<::mediapipe::InferenceCalculatorOptions>(
subgraph_node);
std::vector<absl::string_view> impls;
const bool should_use_gpu =
!options.has_delegate() || // Use GPU delegate if not specified
(options.has_delegate() && options.delegate().has_gpu());
if (should_use_gpu) {
impls.emplace_back("Metal");
impls.emplace_back("MlDrift");
impls.emplace_back("Gl");
}
impls.emplace_back("Cpu");
for (const auto& suffix : impls) {
const auto impl = absl::StrCat("InferenceCalculator", suffix);
if (!mediapipe::CalculatorBaseRegistry::IsRegistered(impl)) continue;
CalculatorGraphConfig::Node impl_node = subgraph_node;
impl_node.set_calculator(impl);
return tool::MakeSingleNodeGraph(std::move(impl_node));
}
return absl::UnimplementedError("no implementation available");
}
};
MEDIAPIPE_REGISTER_NODE(InferenceCalculator);
mediapipe::Status InferenceCalculator::UpdateContract(CalculatorContract* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected())
<< "Either model as side packet or model path in options is required.";
if (ShouldUseGpu(options)) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#elif MEDIAPIPE_TFLITE_METAL_INFERENCE
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#endif
}
return mediapipe::OkStatus();
}
mediapipe::Status InferenceCalculator::Open(CalculatorContext* cc) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE || MEDIAPIPE_TFLITE_METAL_INFERENCE
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
if (ShouldUseGpu(options)) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
use_advanced_gpu_api_ = options.has_delegate() &&
options.delegate().has_gpu() &&
options.delegate().gpu().use_advanced_gpu_api();
allow_precision_loss_ = options.delegate().gpu().allow_precision_loss();
tflite_gpu_runner_api_ = options.delegate().gpu().api();
use_kernel_caching_ =
use_advanced_gpu_api_ && options.delegate().gpu().use_kernel_caching();
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
use_gpu_delegate_ = !use_advanced_gpu_api_;
}
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE || MEDIAPIPE_TFLITE_METAL_INFERENCE
if (use_kernel_caching_) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE && defined(MEDIAPIPE_ANDROID)
cached_kernel_filename_ =
"/sdcard/" + mediapipe::File::Basename(options.model_path()) + ".ker";
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE && 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));
}
if (use_gpu_delegate_ || use_advanced_gpu_api_) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
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)
: LoadDelegate(cc);
}));
#elif MEDIAPIPE_TFLITE_METAL_INFERENCE
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(gpu_helper_);
MP_RETURN_IF_ERROR(LoadDelegate(cc));
#endif
} else {
MP_RETURN_IF_ERROR(LoadDelegate(cc));
}
return mediapipe::OkStatus();
}
mediapipe::Status InferenceCalculator::Process(CalculatorContext* cc) {
if (kInTensors(cc).IsEmpty()) {
return mediapipe::OkStatus();
}
const auto& input_tensors = *kInTensors(cc);
RET_CHECK(!input_tensors.empty());
auto output_tensors = absl::make_unique<std::vector<Tensor>>();
#if MEDIAPIPE_TFLITE_METAL_INFERENCE
id<MTLCommandBuffer> command_buffer;
id<MTLComputeCommandEncoder> compute_encoder;
#endif // MEDIAPIPE_TFLITE_METAL_INFERENCE
if (use_gpu_delegate_ || use_advanced_gpu_api_) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
if (use_advanced_gpu_api_) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors, &output_tensors]() -> ::mediapipe::Status {
for (int i = 0; i < input_tensors.size(); ++i) {
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToInputTensor(
input_tensors[i].GetOpenGlBufferReadView().name(), i));
}
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
output_shapes_[i]);
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToOutputTensor(
output_tensors->back().GetOpenGlBufferWriteView().name(), i));
}
return mediapipe::OkStatus();
}));
} else {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors]() -> ::mediapipe::Status {
// Explicitly copy input.
for (int i = 0; i < input_tensors.size(); ++i) {
glBindBuffer(GL_COPY_READ_BUFFER,
input_tensors[i].GetOpenGlBufferReadView().name());
glBindBuffer(
GL_COPY_WRITE_BUFFER,
gpu_buffers_in_[i]->GetOpenGlBufferWriteView().name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0,
0, input_tensors[i].bytes());
}
return mediapipe::OkStatus();
}));
}
#elif MEDIAPIPE_TFLITE_METAL_INFERENCE
command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"InferenceCalculator";
compute_encoder = [command_buffer computeCommandEncoder];
// Explicit copy input with conversion float 32 bits to 16 bits.
for (int i = 0; i < input_tensors.size(); ++i) {
auto input_view = input_tensors[i].GetMtlBufferReadView(command_buffer);
// Reshape tensor.
tflite::gpu::BHWC shape = BhwcFromTensorShape(input_tensors[i].shape());
auto gpu_buffer_view =
gpu_buffers_in_[i]->GetMtlBufferWriteView(command_buffer);
[converter_to_BPHWC4_ convertWithEncoder:compute_encoder
shape:shape
sourceBuffer:input_view.buffer()
convertedBuffer:gpu_buffer_view.buffer()];
}
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
} else {
// 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();
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());
}
}
// Run inference.
#if MEDIAPIPE_TFLITE_GL_INFERENCE
if (use_advanced_gpu_api_) {
RET_CHECK(tflite_gpu_runner_->Invoke().ok());
} else {
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
}
#else
#if MEDIAPIPE_TFLITE_METAL_INFERENCE
if (use_gpu_delegate_) {
RET_CHECK(
TFLGpuDelegateSetCommandEncoder(delegate_.get(), compute_encoder));
}
#endif // MEDIAPIPE_TFLITE_METAL_INFERENCE
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
if (use_gpu_delegate_ || use_advanced_gpu_api_) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
if (use_gpu_delegate_) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &output_tensors]() -> ::mediapipe::Status {
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
const auto& t = gpu_buffers_out_[i];
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
gpu_buffers_out_[i]->shape());
auto read_view = t->GetOpenGlBufferReadView();
glBindBuffer(GL_COPY_READ_BUFFER, read_view.name());
auto write_view =
output_tensors->back().GetOpenGlBufferWriteView();
glBindBuffer(GL_COPY_WRITE_BUFFER, write_view.name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0,
0, t->bytes());
}
return mediapipe::OkStatus();
}));
}
// Output tensors are already bound if use_advanced_gpu_api_ is true.
#elif MEDIAPIPE_TFLITE_METAL_INFERENCE
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
output_shapes_[i]);
// Reshape tensor.
tflite::gpu::BHWC shape = BhwcFromTensorShape(output_shapes_[i]);
auto read_view =
gpu_buffers_out_[i]->GetMtlBufferReadView(command_buffer);
auto write_view =
output_tensors->at(i).GetMtlBufferWriteView(command_buffer);
[converter_from_BPHWC4_ convertWithEncoder:compute_encoder
shape:shape
sourceBuffer:read_view.buffer()
convertedBuffer:write_view.buffer()];
}
[compute_encoder endEncoding];
[command_buffer commit];
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
} else {
// Output result tensors (CPU).
const auto& tensor_indexes = interpreter_->outputs();
output_tensors->reserve(tensor_indexes.size());
for (int i = 0; i < tensor_indexes.size(); ++i) {
TfLiteTensor* tensor = interpreter_->tensor(tensor_indexes[i]);
output_tensors->emplace_back(
Tensor::ElementType::kFloat32,
Tensor::Shape{std::vector<int>{
tensor->dims->data, tensor->dims->data + tensor->dims->size}});
auto cpu_view = output_tensors->back().GetCpuWriteView();
std::memcpy(cpu_view.buffer<float>(), tensor->data.f,
output_tensors->back().bytes());
}
}
kOutTensors(cc).Send(std::move(output_tensors));
return mediapipe::OkStatus();
}
mediapipe::Status InferenceCalculator::WriteKernelsToFile() {
#if MEDIAPIPE_TFLITE_GL_INFERENCE && defined(MEDIAPIPE_ANDROID)
if (use_kernel_caching_) {
// Save kernel file.
auto kernel_cache = absl::make_unique<std::vector<uint8_t>>(
tflite_gpu_runner_->GetSerializedBinaryCache());
std::string cache_str(kernel_cache->begin(), kernel_cache->end());
MP_RETURN_IF_ERROR(
mediapipe::file::SetContents(cached_kernel_filename_, cache_str));
}
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE && MEDIAPIPE_ANDROID
return mediapipe::OkStatus();
}
mediapipe::Status InferenceCalculator::Close(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(WriteKernelsToFile());
#if MEDIAPIPE_TFLITE_GL_INFERENCE
if (use_gpu_delegate_) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
gpu_buffers_in_.clear();
gpu_buffers_out_.clear();
return mediapipe::OkStatus();
}));
}
#elif MEDIAPIPE_TFLITE_METAL_INFERENCE
converter_to_BPHWC4_ = nil;
converter_from_BPHWC4_ = nil;
gpu_buffers_in_.clear();
gpu_buffers_out_.clear();
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
#if defined(MEDIAPIPE_EDGE_TPU)
edgetpu_context_.reset();
#endif
interpreter_ = nullptr;
delegate_ = nullptr;
return mediapipe::OkStatus();
}
mediapipe::Status InferenceCalculator::ReadKernelsFromFile() {
#if MEDIAPIPE_TFLITE_GL_INFERENCE && defined(MEDIAPIPE_ANDROID)
if (use_kernel_caching_) {
// Load pre-compiled kernel file.
if (mediapipe::File::Exists(cached_kernel_filename_)) {
std::string cache_str;
MP_RETURN_IF_ERROR(
mediapipe::file::GetContents(cached_kernel_filename_, &cache_str));
std::vector<uint8_t> cache_vec(cache_str.begin(), cache_str.end());
tflite_gpu_runner_->SetSerializedBinaryCache(std::move(cache_vec));
}
}
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE && MEDIAPIPE_ANDROID
return mediapipe::OkStatus();
}
mediapipe::Status InferenceCalculator::InitTFLiteGPURunner(
absl::StatusOr<Packet<TfLiteModelPtr>> InferenceCalculator::GetModelAsPacket(
CalculatorContext* cc) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(*cc));
const auto& model = *model_packet_.Get<TfLiteModelPtr>();
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolver());
// Create runner
tflite::gpu::InferenceOptions options;
options.priority1 = allow_precision_loss_
? tflite::gpu::InferencePriority::MIN_LATENCY
: tflite::gpu::InferencePriority::MAX_PRECISION;
options.priority2 = tflite::gpu::InferencePriority::AUTO;
options.priority3 = tflite::gpu::InferencePriority::AUTO;
options.usage = tflite::gpu::InferenceUsage::SUSTAINED_SPEED;
tflite_gpu_runner_ = std::make_unique<tflite::gpu::TFLiteGPURunner>(options);
switch (tflite_gpu_runner_api_) {
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::OPENGL: {
tflite_gpu_runner_->ForceOpenGL();
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::OPENCL: {
tflite_gpu_runner_->ForceOpenCL();
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::ANY: {
// Do not need to force any specific API.
break;
}
}
MP_RETURN_IF_ERROR(
tflite_gpu_runner_->InitializeWithModel(model, op_resolver));
// Create and bind OpenGL buffers for outputs.
// The buffers are created once and their ids are passed to calculator outputs
output_shapes_.resize(tflite_gpu_runner_->outputs_size());
for (int i = 0; i < tflite_gpu_runner_->outputs_size(); ++i) {
output_shapes_[i] = {tflite_gpu_runner_->GetOutputShapes()[i].b,
tflite_gpu_runner_->GetOutputShapes()[i].h,
tflite_gpu_runner_->GetOutputShapes()[i].w,
tflite_gpu_runner_->GetOutputShapes()[i].c};
}
MP_RETURN_IF_ERROR(ReadKernelsFromFile());
MP_RETURN_IF_ERROR(tflite_gpu_runner_->Build());
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
return mediapipe::OkStatus();
}
mediapipe::Status InferenceCalculator::LoadModel(CalculatorContext* cc) {
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(*cc));
const auto& model = *model_packet_.Get<TfLiteModelPtr>();
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolver());
#if defined(MEDIAPIPE_EDGE_TPU)
interpreter_ =
BuildEdgeTpuInterpreter(model, &op_resolver, edgetpu_context_.get());
#else
tflite::InterpreterBuilder(model, op_resolver)(&interpreter_);
#endif // MEDIAPIPE_EDGE_TPU
RET_CHECK(interpreter_);
#if defined(__EMSCRIPTEN__) || defined(MEDIAPIPE_EDGE_TPU)
interpreter_->SetNumThreads(1);
#else
interpreter_->SetNumThreads(
cc->Options<mediapipe::InferenceCalculatorOptions>().cpu_num_thread());
#endif // __EMSCRIPTEN__
RET_CHECK_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
// TODO: Support quantized tensors.
CHECK(interpreter_->tensor(interpreter_->inputs()[0])->quantization.type !=
kTfLiteAffineQuantization);
return mediapipe::OkStatus();
}
mediapipe::StatusOr<mediapipe::Packet> InferenceCalculator::GetModelAsPacket(
const CalculatorContext& cc) {
const auto& options = cc.Options<mediapipe::InferenceCalculatorOptions>();
const auto& options = cc->Options<mediapipe::InferenceCalculatorOptions>();
if (!options.model_path().empty()) {
return TfLiteModelLoader::LoadFromPath(options.model_path());
}
if (cc.InputSidePackets().HasTag("MODEL")) {
return cc.InputSidePackets().Tag("MODEL");
}
return mediapipe::Status(
mediapipe::StatusCode::kNotFound,
"Must specify TFLite model as path or loaded model.");
}
mediapipe::Status InferenceCalculator::LoadDelegate(CalculatorContext* cc) {
const auto& calculator_opts =
cc->Options<mediapipe::InferenceCalculatorOptions>();
if (calculator_opts.has_delegate() &&
calculator_opts.delegate().has_tflite()) {
// Default tflite inference requeqsted - no need to modify graph.
return mediapipe::OkStatus();
}
if (!use_gpu_delegate_) {
#if defined(MEDIAPIPE_ANDROID)
const bool nnapi_requested = calculator_opts.has_delegate()
? calculator_opts.delegate().has_nnapi()
: calculator_opts.use_nnapi();
if (nnapi_requested) {
// Attempt to use NNAPI.
// If not supported, the default CPU delegate will be created and used.
interpreter_->SetAllowFp16PrecisionForFp32(1);
delegate_ =
TfLiteDelegatePtr(tflite::NnApiDelegate(), [](TfLiteDelegate*) {
// No need to free according to tflite::NnApiDelegate()
// documentation.
});
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
return mediapipe::OkStatus();
}
#endif // MEDIAPIPE_ANDROID
#if defined(__EMSCRIPTEN__)
const bool xnnpack_requested = true;
#else
const bool xnnpack_requested = calculator_opts.has_delegate() &&
calculator_opts.delegate().has_xnnpack();
#endif // __EMSCRIPTEN__
#if !defined(MEDIAPIPE_EDGE_TPU)
if (xnnpack_requested) {
TfLiteXNNPackDelegateOptions xnnpack_opts{};
xnnpack_opts.num_threads = GetXnnpackNumThreads(calculator_opts);
delegate_ = TfLiteDelegatePtr(TfLiteXNNPackDelegateCreate(&xnnpack_opts),
&TfLiteXNNPackDelegateDelete);
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
return mediapipe::OkStatus();
}
#endif // !EDGETPU
// Return, no need for GPU delegate below.
return mediapipe::OkStatus();
} else {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
// Configure and create the delegate.
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
options.compile_options.precision_loss_allowed = 1;
options.compile_options.preferred_gl_object_type =
TFLITE_GL_OBJECT_TYPE_FASTEST;
options.compile_options.dynamic_batch_enabled = 0;
options.compile_options.inline_parameters = 1;
delegate_ = TfLiteDelegatePtr(TfLiteGpuDelegateCreate(&options),
&TfLiteGpuDelegateDelete);
// Get input image sizes.
const auto& input_indices = interpreter_->inputs();
for (int i = 0; i < input_indices.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
gpu_buffers_in_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat32,
Tensor::Shape{std::vector<int>{
tensor->dims->data, tensor->dims->data + tensor->dims->size}}));
RET_CHECK_EQ(
TfLiteGpuDelegateBindBufferToTensor(
delegate_.get(),
gpu_buffers_in_.back()->GetOpenGlBufferWriteView().name(),
interpreter_->inputs()[i]),
kTfLiteOk);
}
interpreter_->SetAllowBufferHandleOutput(true);
// Get output image sizes.
const auto& output_indices = interpreter_->outputs();
output_shapes_.resize(output_indices.size());
// Create and bind output buffers.
for (int i = 0; i < output_shapes_.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(output_indices[i]);
gpu_buffers_out_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat32,
Tensor::Shape{std::vector<int>{
tensor->dims->data, tensor->dims->data + tensor->dims->size}}));
RET_CHECK_EQ(
TfLiteGpuDelegateBindBufferToTensor(
delegate_.get(),
gpu_buffers_out_.back()->GetOpenGlBufferWriteView().name(),
output_indices[i]),
kTfLiteOk);
}
// Must call this last.
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
#elif MEDIAPIPE_TFLITE_METAL_INFERENCE
// Configure and create the delegate.
TFLGpuDelegateOptions options;
options.allow_precision_loss = true;
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypeDoNotWait;
delegate_ = TfLiteDelegatePtr(TFLGpuDelegateCreate(&options),
&TFLGpuDelegateDelete);
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
id<MTLDevice> device = gpu_helper_.mtlDevice;
// Get input image sizes.
const auto& input_indices = interpreter_->inputs();
for (int i = 0; i < input_indices.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
// Create and bind input buffer.
std::vector<int> dims{tensor->dims->data,
tensor->dims->data + tensor->dims->size};
dims.back() = RoundUp(dims.back(), 4);
gpu_buffers_in_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat16, Tensor::Shape{dims}));
auto buffer_view =
gpu_buffers_in_[i]->GetMtlBufferWriteView(gpu_helper_.mtlDevice);
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
delegate_.get(), input_indices[i], buffer_view.buffer()),
true);
}
interpreter_->SetAllowBufferHandleOutput(true);
// Get output image sizes.
const auto& output_indices = interpreter_->outputs();
output_shapes_.resize(output_indices.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(output_indices[i]);
RET_CHECK(tensor->dims->size <= 4);
// Create and bind output buffers.
// Channels are always padded to multiple of 4.
std::vector<int> dims{tensor->dims->data,
tensor->dims->data + tensor->dims->size};
output_shapes_[i] = {dims};
dims.back() = RoundUp(dims.back(), 4);
gpu_buffers_out_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat16, Tensor::Shape{dims}));
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
delegate_.get(), output_indices[i],
gpu_buffers_out_[i]
->GetMtlBufferWriteView(gpu_helper_.mtlDevice)
.buffer()),
true);
}
// Create converter for GPU input.
converter_to_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
isFloat16:true
convertToPBHWC4:true];
if (converter_to_BPHWC4_ == nil) {
return mediapipe::InternalError(
"Error initializating input buffer converter");
}
// Create converter for GPU output.
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
isFloat16:true
convertToPBHWC4:false];
if (converter_from_BPHWC4_ == nil) {
return mediapipe::InternalError(
"Error initializating output buffer converter");
}
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
}
return mediapipe::OkStatus();
if (!kSideInModel(cc).IsEmpty()) return kSideInModel(cc);
return absl::Status(mediapipe::StatusCode::kNotFound,
"Must specify TFLite model as path or loaded model.");
}
} // namespace api2
mediapipe/calculators/tensor/inference_calculator.h
+136
View File
@@ -0,0 +1,136 @@
// 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_TENSOR_INFERENCE_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_TENSOR_INFERENCE_CALCULATOR_H_
#include <cstring>
#include <memory>
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "mediapipe/calculators/tensor/inference_calculator.pb.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/tflite/tflite_model_loader.h"
#include "tensorflow/lite/error_reporter.h"
#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/model.h"
namespace mediapipe {
namespace api2 {
// Runs inference on the provided input Tensors and TFLite model.
//
// Creates an interpreter with given model and calls invoke().
// Optionally run inference on CPU/GPU.
//
// This calculator can be used with TensorConverterCalculator to get the
// appropriate inputs.
//
// When the input tensors are on CPU, gpu inference is optional and can be
// specified in the calculator options.
// When the input tensors are on GPU, inference is GPU and output can be CPU or
// GPU.
//
// Input:
// TENSORS - Vector of Tensors
//
// Output:
// TENSORS - Vector of Tensors
//
// Input side packet:
// CUSTOM_OP_RESOLVER (optional) - Use a custom op resolver,
// instead of the builtin one.
// MODEL (optional) - Use to specify TfLite model
// (std::unique_ptr<tflite::FlatBufferModel,
// std::function<void(tflite::FlatBufferModel*)>>)
//
// Example use:
// node {
// calculator: "InferenceCalculator"
// input_stream: "TENSORS:tensor_image"
// output_stream: "TENSORS:tensors"
// options: {
// [mediapipe.InferenceCalculatorOptions.ext] {
// model_path: "modelname.tflite"
// }
// }
// }
//
// or
//
// node {
// calculator: "InferenceCalculator"
// input_stream: "TENSORS:tensor_image"
// input_side_packet: "MODEL:model"
// output_stream: "TENSORS:tensors"
// options: {
// [mediapipe.InferenceCalculatorOptions.ext] {
// model_path: "modelname.tflite"
// delegate { gpu {} }
// }
// }
// }
//
// IMPORTANT Notes:
// Tensors are assumed to be ordered correctly (sequentially added to model).
// Input tensors are assumed to be of the correct size and already normalized.
class InferenceCalculator : public NodeIntf {
public:
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
static constexpr SideInput<tflite::ops::builtin::BuiltinOpResolver>::Optional
kSideInCustomOpResolver{"CUSTOM_OP_RESOLVER"};
static constexpr SideInput<TfLiteModelPtr>::Optional kSideInModel{"MODEL"};
static constexpr Output<std::vector<Tensor>> kOutTensors{"TENSORS"};
MEDIAPIPE_NODE_CONTRACT(kInTensors, kSideInCustomOpResolver, kSideInModel,
kOutTensors);
protected:
using TfLiteDelegatePtr =
std::unique_ptr<TfLiteDelegate, std::function<void(TfLiteDelegate*)>>;
absl::StatusOr<Packet<TfLiteModelPtr>> GetModelAsPacket(
CalculatorContext* cc);
};
struct InferenceCalculatorSelector : public InferenceCalculator {
static constexpr char kCalculatorName[] = "InferenceCalculator";
};
struct InferenceCalculatorGl : public InferenceCalculator {
static constexpr char kCalculatorName[] = "InferenceCalculatorGl";
};
struct InferenceCalculatorMlDrift : public InferenceCalculator {
static constexpr char kCalculatorName[] = "InferenceCalculatorMlDrift";
};
struct InferenceCalculatorMetal : public InferenceCalculator {
static constexpr char kCalculatorName[] = "InferenceCalculatorMetal";
};
struct InferenceCalculatorCpu : public InferenceCalculator {
static constexpr char kCalculatorName[] = "InferenceCalculatorCpu";
};
} // namespace api2
} // namespace mediapipe
#endif // MEDIAPIPE_CALCULATORS_TENSOR_INFERENCE_CALCULATOR_H_
mediapipe/calculators/tensor/inference_calculator.proto
+7 -2
View File
@@ -65,7 +65,8 @@ message InferenceCalculatorOptions {
// Load pre-compiled serialized binary cache to accelerate init process.
// Only available for OpenCL delegate on Android.
optional bool use_kernel_caching = 2 [default = false];
// Kernel caching will only be enabled if this path is set.
optional string cached_kernel_path = 2;
}
// Android only.
message Nnapi {}
@@ -104,7 +105,11 @@ message InferenceCalculatorOptions {
optional int32 cpu_num_thread = 4 [default = -1];
// TfLite delegate to run inference.
// NOTE: calculator is free to choose delegate if not specified explicitly.
// If not specified, TFLite GPU delegate is used by default (as if "gpu {}"
// is specified) unless GPU support is disabled in the build (i.e., with
// --define MEDIAPIPE_DISABLE_GPU=1), in which case regular TFLite on CPU is
// used (as if "tflite {}" is specified) except when building with emscripten
// where xnnpack is used.
// NOTE: use_gpu/use_nnapi are ignored if specified. (Delegate takes
// precedence over use_* deprecated options.)
optional Delegate delegate = 5;
mediapipe/calculators/tensor/inference_calculator_cpu.cc
+205
View File
@@ -0,0 +1,205 @@
// Copyright 2019 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cstring>
#include <memory>
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "mediapipe/calculators/tensor/inference_calculator.h"
#if defined(MEDIAPIPE_ANDROID)
#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
#endif // ANDROID
#if !defined(__EMSCRIPTEN__) || defined(__EMSCRIPTEN_PTHREADS__)
#include "mediapipe/util/cpu_util.h"
#endif // !__EMSCRIPTEN__ || __EMSCRIPTEN_PTHREADS__
#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
namespace mediapipe {
namespace api2 {
namespace {
// Returns number of threads to configure XNNPACK delegate with.
// (Equal to user provided value if specified. Otherwise, it returns number of
// high cores (hard-coded to 1 for Emscripten without Threads extension))
int GetXnnpackNumThreads(const mediapipe::InferenceCalculatorOptions& opts) {
static constexpr int kDefaultNumThreads = -1;
if (opts.has_delegate() && opts.delegate().has_xnnpack() &&
opts.delegate().xnnpack().num_threads() != kDefaultNumThreads) {
return opts.delegate().xnnpack().num_threads();
}
#if !defined(__EMSCRIPTEN__) || defined(__EMSCRIPTEN_PTHREADS__)
return InferHigherCoreIds().size();
#else
return 1;
#endif // !__EMSCRIPTEN__ || __EMSCRIPTEN_PTHREADS__
}
} // namespace
class InferenceCalculatorCpuImpl
: public NodeImpl<InferenceCalculatorCpu, InferenceCalculatorCpuImpl> {
public:
static absl::Status UpdateContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
absl::Status Close(CalculatorContext* cc) override;
private:
absl::Status LoadModel(CalculatorContext* cc);
absl::Status LoadDelegate(CalculatorContext* cc);
// 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_;
};
absl::Status InferenceCalculatorCpuImpl::UpdateContract(
CalculatorContract* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected())
<< "Either model as side packet or model path in options is required.";
return absl::OkStatus();
}
absl::Status InferenceCalculatorCpuImpl::Open(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(LoadModel(cc));
MP_RETURN_IF_ERROR(LoadDelegate(cc));
return absl::OkStatus();
}
absl::Status InferenceCalculatorCpuImpl::Process(CalculatorContext* cc) {
if (kInTensors(cc).IsEmpty()) {
return absl::OkStatus();
}
const auto& input_tensors = *kInTensors(cc);
RET_CHECK(!input_tensors.empty());
auto output_tensors = absl::make_unique<std::vector<Tensor>>();
// 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();
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());
}
// Run inference.
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
// Output result tensors (CPU).
const auto& tensor_indexes = interpreter_->outputs();
output_tensors->reserve(tensor_indexes.size());
for (int i = 0; i < tensor_indexes.size(); ++i) {
TfLiteTensor* tensor = interpreter_->tensor(tensor_indexes[i]);
output_tensors->emplace_back(
Tensor::ElementType::kFloat32,
Tensor::Shape{std::vector<int>{
tensor->dims->data, tensor->dims->data + tensor->dims->size}});
auto cpu_view = output_tensors->back().GetCpuWriteView();
std::memcpy(cpu_view.buffer<float>(), tensor->data.f,
output_tensors->back().bytes());
}
kOutTensors(cc).Send(std::move(output_tensors));
return absl::OkStatus();
}
absl::Status InferenceCalculatorCpuImpl::Close(CalculatorContext* cc) {
interpreter_ = nullptr;
delegate_ = nullptr;
return absl::OkStatus();
}
absl::Status InferenceCalculatorCpuImpl::LoadModel(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::BuiltinOpResolver());
tflite::InterpreterBuilder(model, op_resolver)(&interpreter_);
RET_CHECK(interpreter_);
#if defined(__EMSCRIPTEN__)
interpreter_->SetNumThreads(1);
#else
interpreter_->SetNumThreads(
cc->Options<mediapipe::InferenceCalculatorOptions>().cpu_num_thread());
#endif // __EMSCRIPTEN__
RET_CHECK_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
// TODO: Support quantized tensors.
CHECK(interpreter_->tensor(interpreter_->inputs()[0])->quantization.type !=
kTfLiteAffineQuantization);
return absl::OkStatus();
}
absl::Status InferenceCalculatorCpuImpl::LoadDelegate(CalculatorContext* cc) {
const auto& calculator_opts =
cc->Options<mediapipe::InferenceCalculatorOptions>();
if (calculator_opts.has_delegate() &&
calculator_opts.delegate().has_tflite()) {
// Default tflite inference requeqsted - no need to modify graph.
return absl::OkStatus();
}
#if defined(MEDIAPIPE_ANDROID)
const bool nnapi_requested = calculator_opts.has_delegate()
? calculator_opts.delegate().has_nnapi()
: calculator_opts.use_nnapi();
if (nnapi_requested) {
// Attempt to use NNAPI.
// If not supported, the default CPU delegate will be created and used.
interpreter_->SetAllowFp16PrecisionForFp32(1);
delegate_ = TfLiteDelegatePtr(tflite::NnApiDelegate(), [](TfLiteDelegate*) {
// No need to free according to tflite::NnApiDelegate() documentation.
});
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
return absl::OkStatus();
}
#endif // MEDIAPIPE_ANDROID
#if defined(__EMSCRIPTEN__)
const bool use_xnnpack = true;
#else
const bool use_xnnpack = calculator_opts.has_delegate() &&
calculator_opts.delegate().has_xnnpack();
#endif // defined(__EMSCRIPTEN__)
if (use_xnnpack) {
TfLiteXNNPackDelegateOptions xnnpack_opts{};
xnnpack_opts.num_threads = GetXnnpackNumThreads(calculator_opts);
delegate_ = TfLiteDelegatePtr(TfLiteXNNPackDelegateCreate(&xnnpack_opts),
&TfLiteXNNPackDelegateDelete);
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
}
return absl::OkStatus();
}
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/inference_calculator_face_detection_test.cc
+186
View File
@@ -0,0 +1,186 @@
// Copyright 2019 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "absl/memory/memory.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "mediapipe/calculators/tensor/image_to_tensor_calculator.pb.h"
#include "mediapipe/calculators/tensor/inference_calculator.pb.h"
#include "mediapipe/calculators/tensor/tensors_to_detections_calculator.pb.h"
#include "mediapipe/framework/api2/builder.h"
#include "mediapipe/framework/api2/packet.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/detection.pb.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/graph_test_base.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/status_matchers.h" // NOLINT
#include "mediapipe/framework/tool/subgraph_expansion.h"
#include "mediapipe/framework/tool/test_util.h"
namespace mediapipe {
namespace api2 {
namespace {
using mediapipe::Detection;
using mediapipe::InferenceCalculatorOptions_Delegate;
using testing::ElementsAre;
using testing::EqualsProto;
using testing::proto::Approximately;
struct Param {
std::string name; // Appended to the test name.
std::string impl_suffix; // Expected InferenceCalculator backend.
InferenceCalculatorOptions_Delegate delegate;
};
const std::vector<Param>& GetParams() {
static auto all_params = [] {
static std::vector<Param> p;
p.push_back({"TfLite", "Cpu"});
p.back().delegate.mutable_tflite();
#if TARGET_OS_IPHONE && !TARGET_IPHONE_SIMULATOR
// Metal is not available on the iOS simulator.
p.push_back({"Metal", "Metal"});
p.back().delegate.mutable_gpu();
#endif // TARGET_IPHONE_SIMULATOR
#if __EMSCRIPTEN__
p.push_back({"MlDrift", "MlDrift"});
p.back().delegate.mutable_gpu();
#endif // __EMSCRIPTEN__
#if __ANDROID__ && 0 // Disabled for now since emulator can't go GLESv3
p.push_back({"Gl", "Gl"});
p.back().delegate.mutable_gpu();
// This requires API level 27
p.push_back({"NnApi", "Cpu"});
p.back().delegate.mutable_nnapi();
#endif // __ANDROID__
p.push_back({"XnnPack", "Cpu"});
p.back().delegate.mutable_xnnpack();
return p;
}();
return all_params;
}
class InferenceCalculatorTest : public testing::TestWithParam<Param> {
protected:
#if __EMSCRIPTEN__
// TODO: fix Tensor locking.
// The MlDrift backend currently fails in debug mode without this,
// because of Tensor locking issues. I am adding this temporarily since
// the calculator is already being used and it's better to have test
// coverage for it. Also, the issue doesn't apply to our Emscripten
// build in practice since it's single-threaded.
void SetUp(void) override {
absl::SetMutexDeadlockDetectionMode(absl::OnDeadlockCycle::kIgnore);
}
#endif // __EMSCRIPTEN__
void SetDelegateForParam(mediapipe::CalculatorGraphConfig_Node* node) {
*node->mutable_options()
->MutableExtension(mediapipe::InferenceCalculatorOptions::ext)
->mutable_delegate() = GetParam().delegate;
}
};
TEST_P(InferenceCalculatorTest, TestBackendSelection) {
CalculatorGraphConfig config;
auto node = config.add_node();
node->set_calculator("InferenceCalculator");
SetDelegateForParam(node);
MP_ASSERT_OK(tool::ExpandSubgraphs(&config));
EXPECT_EQ(config.node(0).calculator(),
absl::StrCat("InferenceCalculator", GetParam().impl_suffix));
}
TEST_P(InferenceCalculatorTest, TestFaceDetection) {
CalculatorGraphConfig config;
ASSERT_TRUE(LoadTestGraph(
&config, file::JoinPath(GetTestRootDir(),
"mediapipe/calculators/tensor/"
"testdata/face_detection_test.binarypb")));
// Expand subgraphs to find any nested instances of InferenceCalculator.
MP_ASSERT_OK(tool::ExpandSubgraphs(&config));
int found = 0;
for (auto& node : *config.mutable_node()) {
// The InferenceCalculator subgraph itself will have expanded to a specific
// implementation. Replace it.
// TODO: make it possible to exclude it from expansion above.
if (absl::StartsWith(node.calculator(), "InferenceCalculator")) {
++found;
node.set_calculator("InferenceCalculator");
SetDelegateForParam(&node);
}
}
ASSERT_EQ(found, 1);
std::vector<mediapipe::Packet> detection_packets;
tool::AddVectorSink("detections", &config, &detection_packets);
std::vector<mediapipe::Packet> rendering_packets;
tool::AddVectorSink("rendering", &config, &rendering_packets);
// Load test image.
std::unique_ptr<ImageFrame> input_image = LoadTestPng(
file::JoinPath(GetTestRootDir(), "mediapipe/objc/testdata/sergey.png"));
ASSERT_THAT(input_image, testing::NotNull());
std::unique_ptr<ImageFrame> expected_image =
LoadTestPng(file::JoinPath(GetTestRootDir(),
"mediapipe/calculators/tensor/"
"testdata/face_detection_expected.png"));
ASSERT_THAT(expected_image, testing::NotNull());
std::string binary;
Detection expected_detection;
MP_ASSERT_OK(
file::GetContents(file::JoinPath(GetTestRootDir(),
"mediapipe/calculators/tensor/"
"testdata/expected_detection.binarypb"),
&binary));
expected_detection.ParseFromArray(binary.data(), binary.size());
// Prepare test inputs.
std::unordered_map<std::string, std::unique_ptr<ImageFrame>> input_streams;
input_streams.insert(std::make_pair("image", std::move(input_image)));
std::string output_stream = "rendering";
// Test graph with relaxed color difference tolerance.
// Compare with CPU generated image.
Timestamp ts0 = Timestamp(0);
TestGraphConfig(config, input_streams, output_stream, expected_image, {}, ts0,
2.0, 2.0, 1.0);
ASSERT_EQ(detection_packets.size(), 1);
std::vector<Detection> dets =
detection_packets[0].Get<std::vector<Detection>>();
#if !defined(MEDIAPIPE_PROTO_LITE)
// Approximately is not available with lite protos (b/178137094).
EXPECT_THAT(dets,
ElementsAre(Approximately(EqualsProto(expected_detection))));
#endif
}
INSTANTIATE_TEST_SUITE_P(Implementation, InferenceCalculatorTest,
testing::ValuesIn(GetParams()),
[](const testing::TestParamInfo<Param>& info) {
return info.param.name;
});
} // namespace
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/inference_calculator_gl.cc
+368
View File
@@ -0,0 +1,368 @@
// Copyright 2019 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cstring>
#include <memory>
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "mediapipe/calculators/tensor/inference_calculator.h"
#include "mediapipe/util/tflite/config.h"
#if MEDIAPIPE_TFLITE_GL_INFERENCE
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/util/tflite/tflite_gpu_runner.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
#if defined(MEDIAPIPE_ANDROID)
#include "mediapipe/util/android/file/base/file.h"
#include "mediapipe/util/android/file/base/filesystem.h"
#include "mediapipe/util/android/file/base/helpers.h"
#endif // ANDROID
namespace mediapipe {
namespace api2 {
class InferenceCalculatorGlImpl
: public NodeImpl<InferenceCalculatorGl, InferenceCalculatorGlImpl> {
public:
static absl::Status UpdateContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
absl::Status Close(CalculatorContext* cc) override;
private:
absl::Status ReadKernelsFromFile();
absl::Status WriteKernelsToFile();
absl::Status LoadModel(CalculatorContext* cc);
absl::Status LoadDelegate(CalculatorContext* cc);
absl::Status InitTFLiteGPURunner(CalculatorContext* cc);
// 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_;
#if MEDIAPIPE_TFLITE_GL_INFERENCE
mediapipe::GlCalculatorHelper gpu_helper_;
std::unique_ptr<tflite::gpu::TFLiteGPURunner> tflite_gpu_runner_;
bool allow_precision_loss_ = false;
mediapipe::InferenceCalculatorOptions::Delegate::Gpu::API
tflite_gpu_runner_api_;
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
#if MEDIAPIPE_TFLITE_GPU_SUPPORTED
std::vector<Tensor::Shape> output_shapes_;
std::vector<std::unique_ptr<Tensor>> gpu_buffers_in_;
std::vector<std::unique_ptr<Tensor>> gpu_buffers_out_;
#endif // MEDIAPIPE_TFLITE_GPU_SUPPORTED
bool use_advanced_gpu_api_ = false;
bool use_gpu_delegate_ = false;
bool use_kernel_caching_ = false;
std::string cached_kernel_filename_;
};
absl::Status InferenceCalculatorGlImpl::UpdateContract(CalculatorContract* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected())
<< "Either model as side packet or model path in options is required.";
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::Open(CalculatorContext* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
use_advanced_gpu_api_ = options.has_delegate() &&
options.delegate().has_gpu() &&
options.delegate().gpu().use_advanced_gpu_api();
allow_precision_loss_ = options.delegate().gpu().allow_precision_loss();
tflite_gpu_runner_api_ = options.delegate().gpu().api();
use_kernel_caching_ = use_advanced_gpu_api_ &&
options.delegate().gpu().has_cached_kernel_path();
use_gpu_delegate_ = !use_advanced_gpu_api_;
if (use_kernel_caching_) {
#ifdef MEDIAPIPE_ANDROID
cached_kernel_filename_ = options.delegate().gpu().cached_kernel_path() +
mediapipe::File::Basename(options.model_path()) +
".ker";
#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) : LoadDelegate(cc);
}));
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::Process(CalculatorContext* cc) {
if (kInTensors(cc).IsEmpty()) {
return absl::OkStatus();
}
const auto& input_tensors = *kInTensors(cc);
RET_CHECK(!input_tensors.empty());
auto output_tensors = absl::make_unique<std::vector<Tensor>>();
if (use_advanced_gpu_api_) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors, &output_tensors]() -> ::mediapipe::Status {
for (int i = 0; i < input_tensors.size(); ++i) {
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToInputTensor(
input_tensors[i].GetOpenGlBufferReadView().name(), i));
}
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
output_shapes_[i]);
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToOutputTensor(
output_tensors->back().GetOpenGlBufferWriteView().name(), i));
}
return absl::OkStatus();
}));
} else {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors]() -> ::mediapipe::Status {
// Explicitly copy input.
for (int i = 0; i < input_tensors.size(); ++i) {
glBindBuffer(GL_COPY_READ_BUFFER,
input_tensors[i].GetOpenGlBufferReadView().name());
glBindBuffer(GL_COPY_WRITE_BUFFER,
gpu_buffers_in_[i]->GetOpenGlBufferWriteView().name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0,
input_tensors[i].bytes());
}
return absl::OkStatus();
}));
}
// Run inference.
if (use_advanced_gpu_api_) {
RET_CHECK(tflite_gpu_runner_->Invoke().ok());
} else {
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
}
if (use_gpu_delegate_) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &output_tensors]() -> ::mediapipe::Status {
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
const auto& t = gpu_buffers_out_[i];
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
gpu_buffers_out_[i]->shape());
auto read_view = t->GetOpenGlBufferReadView();
glBindBuffer(GL_COPY_READ_BUFFER, read_view.name());
auto write_view = output_tensors->back().GetOpenGlBufferWriteView();
glBindBuffer(GL_COPY_WRITE_BUFFER, write_view.name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0,
t->bytes());
}
return absl::OkStatus();
}));
}
// Output tensors are already bound if use_advanced_gpu_api_ is true.
kOutTensors(cc).Send(std::move(output_tensors));
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::WriteKernelsToFile() {
#ifdef MEDIAPIPE_ANDROID
if (use_kernel_caching_) {
// Save kernel file.
auto kernel_cache = absl::make_unique<std::vector<uint8_t>>(
tflite_gpu_runner_->GetSerializedBinaryCache());
std::string cache_str(kernel_cache->begin(), kernel_cache->end());
MP_RETURN_IF_ERROR(
mediapipe::file::SetContents(cached_kernel_filename_, cache_str));
}
#endif // MEDIAPIPE_ANDROID
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::Close(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(WriteKernelsToFile());
if (use_gpu_delegate_) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
gpu_buffers_in_.clear();
gpu_buffers_out_.clear();
return absl::OkStatus();
}));
}
interpreter_ = nullptr;
delegate_ = nullptr;
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::ReadKernelsFromFile() {
#ifdef MEDIAPIPE_ANDROID
if (use_kernel_caching_) {
// Load pre-compiled kernel file.
if (mediapipe::File::Exists(cached_kernel_filename_)) {
std::string cache_str;
MP_RETURN_IF_ERROR(
mediapipe::file::GetContents(cached_kernel_filename_, &cache_str));
std::vector<uint8_t> cache_vec(cache_str.begin(), cache_str.end());
tflite_gpu_runner_->SetSerializedBinaryCache(std::move(cache_vec));
}
}
#endif // MEDIAPIPE_ANDROID
return absl::OkStatus();
}
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::BuiltinOpResolver());
// Create runner
tflite::gpu::InferenceOptions options;
options.priority1 = allow_precision_loss_
? tflite::gpu::InferencePriority::MIN_LATENCY
: tflite::gpu::InferencePriority::MAX_PRECISION;
options.priority2 = tflite::gpu::InferencePriority::AUTO;
options.priority3 = tflite::gpu::InferencePriority::AUTO;
options.usage = tflite::gpu::InferenceUsage::SUSTAINED_SPEED;
tflite_gpu_runner_ = std::make_unique<tflite::gpu::TFLiteGPURunner>(options);
switch (tflite_gpu_runner_api_) {
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::OPENGL: {
tflite_gpu_runner_->ForceOpenGL();
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::OPENCL: {
tflite_gpu_runner_->ForceOpenCL();
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::ANY: {
// Do not need to force any specific API.
break;
}
}
MP_RETURN_IF_ERROR(
tflite_gpu_runner_->InitializeWithModel(model, op_resolver));
// Create and bind OpenGL buffers for outputs.
// The buffers are created once and their ids are passed to calculator outputs
output_shapes_.resize(tflite_gpu_runner_->outputs_size());
for (int i = 0; i < tflite_gpu_runner_->outputs_size(); ++i) {
output_shapes_[i] = {tflite_gpu_runner_->GetOutputShapes()[i].b,
tflite_gpu_runner_->GetOutputShapes()[i].h,
tflite_gpu_runner_->GetOutputShapes()[i].w,
tflite_gpu_runner_->GetOutputShapes()[i].c};
}
MP_RETURN_IF_ERROR(ReadKernelsFromFile());
MP_RETURN_IF_ERROR(tflite_gpu_runner_->Build());
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::LoadModel(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::BuiltinOpResolver());
tflite::InterpreterBuilder(model, op_resolver)(&interpreter_);
RET_CHECK(interpreter_);
#if defined(__EMSCRIPTEN__)
interpreter_->SetNumThreads(1);
#else
interpreter_->SetNumThreads(
cc->Options<mediapipe::InferenceCalculatorOptions>().cpu_num_thread());
#endif // __EMSCRIPTEN__
RET_CHECK_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
// TODO: Support quantized tensors.
CHECK(interpreter_->tensor(interpreter_->inputs()[0])->quantization.type !=
kTfLiteAffineQuantization);
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::LoadDelegate(CalculatorContext* cc) {
// Configure and create the delegate.
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
options.compile_options.precision_loss_allowed = 1;
options.compile_options.preferred_gl_object_type =
TFLITE_GL_OBJECT_TYPE_FASTEST;
options.compile_options.dynamic_batch_enabled = 0;
options.compile_options.inline_parameters = 1;
delegate_ = TfLiteDelegatePtr(TfLiteGpuDelegateCreate(&options),
&TfLiteGpuDelegateDelete);
// Get input image sizes.
const auto& input_indices = interpreter_->inputs();
for (int i = 0; i < input_indices.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
gpu_buffers_in_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat32,
Tensor::Shape{std::vector<int>{
tensor->dims->data, tensor->dims->data + tensor->dims->size}}));
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
delegate_.get(),
gpu_buffers_in_.back()->GetOpenGlBufferWriteView().name(),
interpreter_->inputs()[i]),
kTfLiteOk);
}
interpreter_->SetAllowBufferHandleOutput(true);
// Get output image sizes.
const auto& output_indices = interpreter_->outputs();
output_shapes_.resize(output_indices.size());
// Create and bind output buffers.
for (int i = 0; i < output_shapes_.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(output_indices[i]);
gpu_buffers_out_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat32,
Tensor::Shape{std::vector<int>{
tensor->dims->data, tensor->dims->data + tensor->dims->size}}));
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
delegate_.get(),
gpu_buffers_out_.back()->GetOpenGlBufferWriteView().name(),
output_indices[i]),
kTfLiteOk);
}
// Must call this last.
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
return absl::OkStatus();
}
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/inference_calculator_metal.cc
+293
View File
@@ -0,0 +1,293 @@
// 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.
#import <CoreVideo/CoreVideo.h>
#import <Metal/Metal.h>
#import <MetalKit/MetalKit.h>
#include <cstring>
#include <memory>
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "mediapipe/calculators/tensor/inference_calculator.h"
#import "mediapipe/gpu/MPPMetalHelper.h"
#include "mediapipe/gpu/MPPMetalUtil.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/util/tflite/config.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/metal/buffer_convert.h"
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
#include "tensorflow/lite/delegates/gpu/metal_delegate_internal.h"
namespace {
// Round up n to next multiple of m.
template <typename T>
T RoundUp(T n, T m) {
return ((n + m - T{1}) / m) * m;
}
} // namespace
namespace mediapipe {
namespace api2 {
#if MEDIAPIPE_TFLITE_METAL_INFERENCE
namespace {
tflite::gpu::BHWC BhwcFromTensorShape(const Tensor::Shape& shape) {
tflite::gpu::BHWC result;
result.b = shape.dims[0];
switch (shape.dims.size()) {
case 1:
// result.b is already filled.
break;
case 2:
result.h = 1;
result.w = 1;
result.c = shape.dims[1];
break;
case 3:
result.h = 1;
result.w = shape.dims[1];
result.c = shape.dims[2];
break;
case 4:
result.h = shape.dims[1];
result.w = shape.dims[2];
result.c = shape.dims[3];
break;
default:
// Handles 0 and >4.
LOG(FATAL)
<< "Dimensions size must be in range [1,4] for GPU inference, but "
<< shape.dims.size() << " is provided";
}
return result;
}
} // namespace
#endif // MEDIAPIPE_TFLITE_METAL_INFERENCE
class InferenceCalculatorMetalImpl
: public NodeImpl<InferenceCalculatorMetal, InferenceCalculatorMetalImpl> {
public:
static absl::Status UpdateContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
absl::Status Close(CalculatorContext* cc) override;
private:
absl::Status LoadModel(CalculatorContext* cc);
absl::Status LoadDelegate(CalculatorContext* cc);
// 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_;
#if MEDIAPIPE_TFLITE_METAL_INFERENCE
MPPMetalHelper* gpu_helper_ = nullptr;
TFLBufferConvert* converter_to_BPHWC4_ = nil;
TFLBufferConvert* converter_from_BPHWC4_ = nil;
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
#if MEDIAPIPE_TFLITE_GPU_SUPPORTED
std::vector<Tensor::Shape> output_shapes_;
std::vector<std::unique_ptr<Tensor>> gpu_buffers_in_;
std::vector<std::unique_ptr<Tensor>> gpu_buffers_out_;
#endif // MEDIAPIPE_TFLITE_GPU_SUPPORTED
};
absl::Status InferenceCalculatorMetalImpl::UpdateContract(
CalculatorContract* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected())
<< "Either model as side packet or model path in options is required.";
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
return absl::OkStatus();
}
absl::Status InferenceCalculatorMetalImpl::Open(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(LoadModel(cc));
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(gpu_helper_);
MP_RETURN_IF_ERROR(LoadDelegate(cc));
return absl::OkStatus();
}
absl::Status InferenceCalculatorMetalImpl::Process(CalculatorContext* cc) {
if (kInTensors(cc).IsEmpty()) {
return absl::OkStatus();
}
const auto& input_tensors = *kInTensors(cc);
RET_CHECK(!input_tensors.empty());
auto output_tensors = absl::make_unique<std::vector<Tensor>>();
id<MTLCommandBuffer> command_buffer;
command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"InferenceCalculator";
// Explicit copy input with conversion float 32 bits to 16 bits.
for (int i = 0; i < input_tensors.size(); ++i) {
auto input_view = input_tensors[i].GetMtlBufferReadView(command_buffer);
// Reshape tensor.
tflite::gpu::BHWC shape = BhwcFromTensorShape(input_tensors[i].shape());
auto gpu_buffer_view =
gpu_buffers_in_[i]->GetMtlBufferWriteView(command_buffer);
id<MTLComputeCommandEncoder> input_encoder =
[command_buffer computeCommandEncoder];
[converter_to_BPHWC4_ convertWithEncoder:input_encoder
shape:shape
sourceBuffer:input_view.buffer()
convertedBuffer:gpu_buffer_view.buffer()];
[input_encoder endEncoding];
}
// Run inference.
RET_CHECK(TFLGpuDelegateSetCommandBuffer(delegate_.get(), command_buffer));
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
output_shapes_[i]);
// Reshape tensor.
tflite::gpu::BHWC shape = BhwcFromTensorShape(output_shapes_[i]);
auto read_view = gpu_buffers_out_[i]->GetMtlBufferReadView(command_buffer);
auto write_view =
output_tensors->at(i).GetMtlBufferWriteView(command_buffer);
id<MTLComputeCommandEncoder> output_encoder =
[command_buffer computeCommandEncoder];
[converter_from_BPHWC4_ convertWithEncoder:output_encoder
shape:shape
sourceBuffer:read_view.buffer()
convertedBuffer:write_view.buffer()];
[output_encoder endEncoding];
}
[command_buffer commit];
kOutTensors(cc).Send(std::move(output_tensors));
return absl::OkStatus();
}
absl::Status InferenceCalculatorMetalImpl::Close(CalculatorContext* cc) {
converter_to_BPHWC4_ = nil;
converter_from_BPHWC4_ = nil;
gpu_buffers_in_.clear();
gpu_buffers_out_.clear();
interpreter_ = nullptr;
delegate_ = nullptr;
return absl::OkStatus();
}
absl::Status InferenceCalculatorMetalImpl::LoadModel(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::BuiltinOpResolver());
tflite::InterpreterBuilder(model, op_resolver)(&interpreter_);
RET_CHECK(interpreter_);
interpreter_->SetNumThreads(
cc->Options<mediapipe::InferenceCalculatorOptions>().cpu_num_thread());
RET_CHECK_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
// TODO: Support quantized tensors.
CHECK(interpreter_->tensor(interpreter_->inputs()[0])->quantization.type !=
kTfLiteAffineQuantization);
return absl::OkStatus();
}
absl::Status InferenceCalculatorMetalImpl::LoadDelegate(CalculatorContext* cc) {
const auto& calculator_opts =
cc->Options<mediapipe::InferenceCalculatorOptions>();
// Configure and create the delegate.
TFLGpuDelegateOptions options;
options.allow_precision_loss = true;
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypeDoNotWait;
delegate_ =
TfLiteDelegatePtr(TFLGpuDelegateCreate(&options), &TFLGpuDelegateDelete);
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
kTfLiteOk);
id<MTLDevice> device = gpu_helper_.mtlDevice;
// Get input image sizes.
const auto& input_indices = interpreter_->inputs();
for (int i = 0; i < input_indices.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
// Create and bind input buffer.
std::vector<int> dims{tensor->dims->data,
tensor->dims->data + tensor->dims->size};
dims.back() = RoundUp(dims.back(), 4);
gpu_buffers_in_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat16, Tensor::Shape{dims}));
auto buffer_view =
gpu_buffers_in_[i]->GetMtlBufferWriteView(gpu_helper_.mtlDevice);
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
delegate_.get(), input_indices[i], buffer_view.buffer()),
true);
}
interpreter_->SetAllowBufferHandleOutput(true);
// Get output image sizes.
const auto& output_indices = interpreter_->outputs();
output_shapes_.resize(output_indices.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(output_indices[i]);
RET_CHECK(tensor->dims->size <= 4);
// Create and bind output buffers.
// Channels are always padded to multiple of 4.
std::vector<int> dims{tensor->dims->data,
tensor->dims->data + tensor->dims->size};
output_shapes_[i] = {dims};
dims.back() = RoundUp(dims.back(), 4);
gpu_buffers_out_.emplace_back(absl::make_unique<Tensor>(
Tensor::ElementType::kFloat16, Tensor::Shape{dims}));
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
delegate_.get(), output_indices[i],
gpu_buffers_out_[i]
->GetMtlBufferWriteView(gpu_helper_.mtlDevice)
.buffer()),
true);
}
// Create converter for GPU input.
converter_to_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
isFloat16:true
convertToPBHWC4:true];
if (converter_to_BPHWC4_ == nil) {
return mediapipe::InternalError(
"Error initializating input buffer converter");
}
// Create converter for GPU output.
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
isFloat16:true
convertToPBHWC4:false];
if (converter_from_BPHWC4_ == nil) {
return absl::InternalError("Error initializating output buffer converter");
}
return absl::OkStatus();
}
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/inference_calculator_test.cc
+3 -3
View File
@@ -111,9 +111,9 @@ TEST(InferenceCalculatorTest, SmokeTest) {
// Test CPU inference only.
DoSmokeTest(/*graph_proto=*/absl::StrReplaceAll(
graph_proto, {{"$delegate", "delegate { tflite {} }"}}));
DoSmokeTest(/*graph_proto=*/absl::StrReplaceAll(
graph_proto, {{"$delegate", "delegate { xnnpack {} }"}}));
DoSmokeTest(/*graph_proto=*/absl::StrReplaceAll(
DoSmokeTest(absl::StrReplaceAll(graph_proto,
{{"$delegate", "delegate { xnnpack {} }"}}));
DoSmokeTest(absl::StrReplaceAll(
graph_proto,
{{"$delegate", "delegate { xnnpack { num_threads: 10 } }"}}));
}
mediapipe/calculators/tensor/tensor_converter_calculator.cc
+41 -43
View File
@@ -100,21 +100,21 @@ namespace mediapipe {
class TensorConverterCalculator : public CalculatorBase {
public:
static mediapipe::Status GetContract(CalculatorContract* cc);
static absl::Status GetContract(CalculatorContract* cc);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
mediapipe::Status Close(CalculatorContext* cc) override;
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
absl::Status Close(CalculatorContext* cc) override;
private:
mediapipe::Status InitGpu(CalculatorContext* cc);
mediapipe::Status LoadOptions(CalculatorContext* cc);
absl::Status InitGpu(CalculatorContext* cc);
absl::Status LoadOptions(CalculatorContext* cc);
template <class T>
mediapipe::Status NormalizeImage(const ImageFrame& image_frame,
bool flip_vertically, float* tensor_ptr);
mediapipe::Status CopyMatrixToTensor(const Matrix& matrix, float* tensor_ptr);
mediapipe::Status ProcessCPU(CalculatorContext* cc);
mediapipe::Status ProcessGPU(CalculatorContext* cc);
absl::Status NormalizeImage(const ImageFrame& image_frame,
bool flip_vertically, float* tensor_ptr);
absl::Status CopyMatrixToTensor(const Matrix& matrix, float* tensor_ptr);
absl::Status ProcessCPU(CalculatorContext* cc);
absl::Status ProcessGPU(CalculatorContext* cc);
#if MEDIAPIPE_METAL_ENABLED
MPPMetalHelper* gpu_helper_ = nullptr;
@@ -139,8 +139,7 @@ class TensorConverterCalculator : public CalculatorBase {
};
REGISTER_CALCULATOR(TensorConverterCalculator);
mediapipe::Status TensorConverterCalculator::GetContract(
CalculatorContract* cc) {
absl::Status TensorConverterCalculator::GetContract(CalculatorContract* cc) {
// Confirm only one of the input streams is present.
RET_CHECK(static_cast<int>(cc->Inputs().HasTag(kImageFrameTag)) +
static_cast<int>(cc->Inputs().HasTag(kGpuBufferTag)) +
@@ -167,10 +166,10 @@ mediapipe::Status TensorConverterCalculator::GetContract(
RET_CHECK(cc->Outputs().HasTag(kTensorsTag));
cc->Outputs().Tag(kTensorsTag).Set<std::vector<Tensor>>();
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorConverterCalculator::Open(CalculatorContext* cc) {
absl::Status TensorConverterCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
MP_RETURN_IF_ERROR(LoadOptions(cc));
@@ -187,13 +186,13 @@ mediapipe::Status TensorConverterCalculator::Open(CalculatorContext* cc) {
}
#endif // !MEDIAPIPE_DISABLE_GPU
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorConverterCalculator::Process(CalculatorContext* cc) {
absl::Status TensorConverterCalculator::Process(CalculatorContext* cc) {
if (use_gpu_) {
if (cc->Inputs().Tag(kGpuBufferTag).IsEmpty()) {
return mediapipe::OkStatus();
return absl::OkStatus();
}
// Convert to GPU tensors type.
MP_RETURN_IF_ERROR(ProcessGPU(cc));
@@ -201,10 +200,10 @@ mediapipe::Status TensorConverterCalculator::Process(CalculatorContext* cc) {
// Convert to CPU tensors or Matrix type.
MP_RETURN_IF_ERROR(ProcessCPU(cc));
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorConverterCalculator::Close(CalculatorContext* cc) {
absl::Status TensorConverterCalculator::Close(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
if (use_gpu_) {
#if MEDIAPIPE_METAL_ENABLED
@@ -221,14 +220,14 @@ mediapipe::Status TensorConverterCalculator::Close(CalculatorContext* cc) {
#endif // MEDIAPIPE_METAL_ENABLED
}
#endif // !MEDIAPIPE_DISABLE_GPU
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorConverterCalculator::ProcessCPU(CalculatorContext* cc) {
absl::Status TensorConverterCalculator::ProcessCPU(CalculatorContext* cc) {
auto output_tensors = absl::make_unique<std::vector<Tensor>>();
if (cc->Inputs().HasTag(kImageFrameTag)) {
if (cc->Inputs().Tag(kImageFrameTag).IsEmpty()) {
return mediapipe::OkStatus();
return absl::OkStatus();
}
const auto& image_frame =
cc->Inputs().Tag(kImageFrameTag).Get<ImageFrame>();
@@ -257,12 +256,12 @@ mediapipe::Status TensorConverterCalculator::ProcessCPU(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(NormalizeImage<float>(image_frame, flip_vertically_,
cpu_view.buffer<float>()));
} else {
return mediapipe::InternalError(
return absl::InternalError(
"Only byte-based (8 bit) and float (32 bit) images supported.");
}
} else if (cc->Inputs().HasTag(kMatrixTag)) {
if (cc->Inputs().Tag(kMatrixTag).IsEmpty()) {
return mediapipe::OkStatus();
return absl::OkStatus();
}
const auto& matrix = cc->Inputs().Tag(kMatrixTag).Get<Matrix>();
const int height = matrix.rows();
@@ -273,16 +272,16 @@ mediapipe::Status TensorConverterCalculator::ProcessCPU(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(CopyMatrixToTensor(
matrix, output_tensors->back().GetCpuWriteView().buffer<float>()));
} else {
return mediapipe::OkStatus();
return absl::OkStatus();
}
cc->Outputs()
.Tag(kTensorsTag)
.Add(output_tensors.release(), cc->InputTimestamp());
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorConverterCalculator::ProcessGPU(CalculatorContext* cc) {
absl::Status TensorConverterCalculator::ProcessGPU(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
if (!initialized_) {
MP_RETURN_IF_ERROR(InitGpu(cc));
@@ -318,7 +317,7 @@ mediapipe::Status TensorConverterCalculator::ProcessGPU(CalculatorContext* cc) {
[command_buffer commit];
#elif MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &output_tensors, &input]() -> mediapipe::Status {
[this, &output_tensors, &input]() -> absl::Status {
auto src = gpu_helper_.CreateSourceTexture(input);
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
// Convert GL texture into SSBO.
@@ -361,7 +360,7 @@ mediapipe::Status TensorConverterCalculator::ProcessGPU(CalculatorContext* cc) {
glBindTexture(GL_TEXTURE_2D, 0);
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
src.Release();
return mediapipe::OkStatus();
return absl::OkStatus();
}));
#endif // MEDIAPIPE_METAL_ENABLED
cc->Outputs()
@@ -371,10 +370,10 @@ mediapipe::Status TensorConverterCalculator::ProcessGPU(CalculatorContext* cc) {
RET_CHECK_FAIL() << "GPU processing is not enabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
absl::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
// Get input image sizes.
const auto& input =
@@ -448,7 +447,7 @@ mediapipe::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
&input,
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
&single_channel]()
-> mediapipe::Status {
-> absl::Status {
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
// Shader to convert GL Texture to Shader Storage Buffer Object (SSBO),
// with normalization to either: [0,1] or [-1,1].
@@ -558,15 +557,14 @@ mediapipe::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
glGenFramebuffers(1, &framebuffer_);
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
return mediapipe::OkStatus();
return absl::OkStatus();
}));
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#endif // !MEDIAPIPE_DISABLE_GPU
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorConverterCalculator::LoadOptions(
CalculatorContext* cc) {
absl::Status TensorConverterCalculator::LoadOptions(CalculatorContext* cc) {
// Get calculator options specified in the graph.
const auto& options =
cc->Options<::mediapipe::TensorConverterCalculatorOptions>();
@@ -604,11 +602,11 @@ mediapipe::Status TensorConverterCalculator::LoadOptions(
CHECK_GE(max_num_channels_, 1);
CHECK_LE(max_num_channels_, 4);
CHECK_NE(max_num_channels_, 2);
return mediapipe::OkStatus();
return absl::OkStatus();
}
template <class T>
mediapipe::Status TensorConverterCalculator::NormalizeImage(
absl::Status TensorConverterCalculator::NormalizeImage(
const ImageFrame& image_frame, bool flip_vertically, float* tensor_ptr) {
const int height = image_frame.Height();
const int width = image_frame.Width();
@@ -652,11 +650,11 @@ mediapipe::Status TensorConverterCalculator::NormalizeImage(
}
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorConverterCalculator::CopyMatrixToTensor(
const Matrix& matrix, float* tensor_ptr) {
absl::Status TensorConverterCalculator::CopyMatrixToTensor(const Matrix& matrix,
float* tensor_ptr) {
if (row_major_matrix_) {
auto matrix_map =
Eigen::Map<RowMajorMatrixXf>(tensor_ptr, matrix.rows(), matrix.cols());
@@ -667,7 +665,7 @@ mediapipe::Status TensorConverterCalculator::CopyMatrixToTensor(
matrix_map = matrix;
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
} // namespace mediapipe
mediapipe/calculators/tensor/tensors_to_classification_calculator.cc
+11 -13
View File
@@ -16,6 +16,7 @@
#include <unordered_map>
#include <vector>
#include "absl/container/node_hash_map.h"
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "mediapipe/calculators/tensor/tensors_to_classification_calculator.pb.h"
@@ -66,20 +67,19 @@ class TensorsToClassificationCalculator : public Node {
"CLASSIFICATIONS"};
MEDIAPIPE_NODE_CONTRACT(kInTensors, kOutClassificationList);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
mediapipe::Status Close(CalculatorContext* cc) override;
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
absl::Status Close(CalculatorContext* cc) override;
private:
::mediapipe::TensorsToClassificationCalculatorOptions options_;
int top_k_ = 0;
std::unordered_map<int, std::string> label_map_;
absl::node_hash_map<int, std::string> label_map_;
bool label_map_loaded_ = false;
};
MEDIAPIPE_REGISTER_NODE(TensorsToClassificationCalculator);
mediapipe::Status TensorsToClassificationCalculator::Open(
CalculatorContext* cc) {
absl::Status TensorsToClassificationCalculator::Open(CalculatorContext* cc) {
options_ =
cc->Options<::mediapipe::TensorsToClassificationCalculatorOptions>();
@@ -100,11 +100,10 @@ mediapipe::Status TensorsToClassificationCalculator::Open(
label_map_loaded_ = true;
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToClassificationCalculator::Process(
CalculatorContext* cc) {
absl::Status TensorsToClassificationCalculator::Process(CalculatorContext* cc) {
const auto& input_tensors = *kInTensors(cc);
RET_CHECK_EQ(input_tensors.size(), 1);
@@ -168,12 +167,11 @@ mediapipe::Status TensorsToClassificationCalculator::Process(
top_k_, raw_classification_list->size() - top_k_);
}
kOutClassificationList(cc).Send(std::move(classification_list));
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToClassificationCalculator::Close(
CalculatorContext* cc) {
return mediapipe::OkStatus();
absl::Status TensorsToClassificationCalculator::Close(CalculatorContext* cc) {
return absl::OkStatus();
}
} // namespace api2
mediapipe/calculators/tensor/tensors_to_detections_calculator.cc
+57 -48
View File
@@ -134,26 +134,27 @@ class TensorsToDetectionsCalculator : public Node {
"ANCHORS"};
static constexpr Output<std::vector<Detection>> kOutDetections{"DETECTIONS"};
MEDIAPIPE_NODE_CONTRACT(kInTensors, kInAnchors, kOutDetections);
static mediapipe::Status UpdateContract(CalculatorContract* cc);
static absl::Status UpdateContract(CalculatorContract* cc);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
mediapipe::Status Close(CalculatorContext* cc) override;
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
absl::Status Close(CalculatorContext* cc) override;
private:
mediapipe::Status ProcessCPU(CalculatorContext* cc,
std::vector<Detection>* output_detections);
mediapipe::Status ProcessGPU(CalculatorContext* cc,
std::vector<Detection>* output_detections);
absl::Status ProcessCPU(CalculatorContext* cc,
std::vector<Detection>* output_detections);
absl::Status ProcessGPU(CalculatorContext* cc,
std::vector<Detection>* output_detections);
mediapipe::Status LoadOptions(CalculatorContext* cc);
mediapipe::Status GpuInit(CalculatorContext* cc);
mediapipe::Status DecodeBoxes(const float* raw_boxes,
const std::vector<Anchor>& anchors,
std::vector<float>* boxes);
mediapipe::Status ConvertToDetections(
const float* detection_boxes, const float* detection_scores,
const int* detection_classes, std::vector<Detection>* output_detections);
absl::Status LoadOptions(CalculatorContext* cc);
absl::Status GpuInit(CalculatorContext* cc);
absl::Status DecodeBoxes(const float* raw_boxes,
const std::vector<Anchor>& anchors,
std::vector<float>* boxes);
absl::Status ConvertToDetections(const float* detection_boxes,
const float* detection_scores,
const int* detection_classes,
std::vector<Detection>* output_detections);
Detection ConvertToDetection(float box_ymin, float box_xmin, float box_ymax,
float box_xmax, float score, int class_id,
bool flip_vertically);
@@ -179,12 +180,13 @@ class TensorsToDetectionsCalculator : public Node {
std::unique_ptr<Tensor> decoded_boxes_buffer_;
std::unique_ptr<Tensor> scored_boxes_buffer_;
bool gpu_inited_ = false;
bool gpu_input_ = false;
bool anchors_init_ = false;
};
MEDIAPIPE_REGISTER_NODE(TensorsToDetectionsCalculator);
mediapipe::Status TensorsToDetectionsCalculator::UpdateContract(
absl::Status TensorsToDetectionsCalculator::UpdateContract(
CalculatorContract* cc) {
if (CanUseGpu()) {
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
@@ -194,10 +196,10 @@ mediapipe::Status TensorsToDetectionsCalculator::UpdateContract(
#endif // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToDetectionsCalculator::Open(CalculatorContext* cc) {
absl::Status TensorsToDetectionsCalculator::Open(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(LoadOptions(cc));
if (CanUseGpu()) {
@@ -207,14 +209,12 @@ mediapipe::Status TensorsToDetectionsCalculator::Open(CalculatorContext* cc) {
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(gpu_helper_);
#endif // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
MP_RETURN_IF_ERROR(GpuInit(cc));
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToDetectionsCalculator::Process(
CalculatorContext* cc) {
absl::Status TensorsToDetectionsCalculator::Process(CalculatorContext* cc) {
auto output_detections = absl::make_unique<std::vector<Detection>>();
bool gpu_processing = false;
if (CanUseGpu()) {
@@ -229,16 +229,20 @@ mediapipe::Status TensorsToDetectionsCalculator::Process(
}
if (gpu_processing) {
if (!gpu_inited_) {
MP_RETURN_IF_ERROR(GpuInit(cc));
gpu_inited_ = true;
}
MP_RETURN_IF_ERROR(ProcessGPU(cc, output_detections.get()));
} else {
MP_RETURN_IF_ERROR(ProcessCPU(cc, output_detections.get()));
}
kOutDetections(cc).Send(std::move(output_detections));
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToDetectionsCalculator::ProcessCPU(
absl::Status TensorsToDetectionsCalculator::ProcessCPU(
CalculatorContext* cc, std::vector<Detection>* output_detections) {
const auto& input_tensors = *kInTensors(cc);
@@ -275,7 +279,7 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessCPU(
} else if (!kInAnchors(cc).IsEmpty()) {
anchors_ = *kInAnchors(cc);
} else {
return mediapipe::UnavailableError("No anchor data available.");
return absl::UnavailableError("No anchor data available.");
}
anchors_init_ = true;
}
@@ -362,10 +366,10 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessCPU(
detection_classes.data(),
output_detections));
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
absl::Status TensorsToDetectionsCalculator::ProcessGPU(
CalculatorContext* cc, std::vector<Detection>* output_detections) {
const auto& input_tensors = *kInTensors(cc);
RET_CHECK_GE(input_tensors.size(), 2);
@@ -373,7 +377,7 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, &input_tensors, &cc,
&output_detections]()
-> mediapipe::Status {
-> absl::Status {
if (!anchors_init_) {
if (input_tensors.size() == kNumInputTensorsWithAnchors) {
auto read_view = input_tensors[2].GetOpenGlBufferReadView();
@@ -388,7 +392,7 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
auto raw_anchors = anchors_view.buffer<float>();
ConvertAnchorsToRawValues(anchors, num_boxes_, raw_anchors);
} else {
return mediapipe::UnavailableError("No anchor data available.");
return absl::UnavailableError("No anchor data available.");
}
anchors_init_ = true;
}
@@ -414,7 +418,7 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
glUseProgram(score_program_);
glDispatchCompute(num_boxes_, 1, 1);
}
return mediapipe::OkStatus();
return absl::OkStatus();
}));
// TODO: b/138851969. Is it possible to output a float vector
@@ -459,7 +463,7 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
ConvertAnchorsToRawValues(anchors, num_boxes_,
raw_anchors_view.buffer<float>());
} else {
return mediapipe::UnavailableError("No anchor data available.");
return absl::UnavailableError("No anchor data available.");
}
anchors_init_ = true;
}
@@ -520,10 +524,10 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
#else
LOG(ERROR) << "GPU input on non-Android not supported yet.";
#endif // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToDetectionsCalculator::Close(CalculatorContext* cc) {
absl::Status TensorsToDetectionsCalculator::Close(CalculatorContext* cc) {
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
gpu_helper_.RunInGlContext([this] {
decoded_boxes_buffer_ = nullptr;
@@ -540,11 +544,10 @@ mediapipe::Status TensorsToDetectionsCalculator::Close(CalculatorContext* cc) {
score_program_ = nil;
#endif // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToDetectionsCalculator::LoadOptions(
CalculatorContext* cc) {
absl::Status TensorsToDetectionsCalculator::LoadOptions(CalculatorContext* cc) {
// Get calculator options specified in the graph.
options_ = cc->Options<::mediapipe::TensorsToDetectionsCalculatorOptions>();
RET_CHECK(options_.has_num_classes());
@@ -567,10 +570,10 @@ mediapipe::Status TensorsToDetectionsCalculator::LoadOptions(
ignore_classes_.insert(options_.ignore_classes(i));
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToDetectionsCalculator::DecodeBoxes(
absl::Status TensorsToDetectionsCalculator::DecodeBoxes(
const float* raw_boxes, const std::vector<Anchor>& anchors,
std::vector<float>* boxes) {
for (int i = 0; i < num_boxes_; ++i) {
@@ -631,10 +634,10 @@ mediapipe::Status TensorsToDetectionsCalculator::DecodeBoxes(
}
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToDetectionsCalculator::ConvertToDetections(
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) {
@@ -671,7 +674,7 @@ mediapipe::Status TensorsToDetectionsCalculator::ConvertToDetections(
}
output_detections->emplace_back(detection);
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
Detection TensorsToDetectionsCalculator::ConvertToDetection(
@@ -694,10 +697,9 @@ Detection TensorsToDetectionsCalculator::ConvertToDetection(
return detection;
}
mediapipe::Status TensorsToDetectionsCalculator::GpuInit(
CalculatorContext* cc) {
absl::Status TensorsToDetectionsCalculator::GpuInit(CalculatorContext* cc) {
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> mediapipe::Status {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> absl::Status {
// A shader to decode detection boxes.
const std::string decode_src = absl::Substitute(
R"( #version 310 es
@@ -801,7 +803,14 @@ void main() {
glCompileShader(shader);
GLint compiled = GL_FALSE;
glGetShaderiv(shader, GL_COMPILE_STATUS, &compiled);
RET_CHECK(compiled == GL_TRUE);
RET_CHECK(compiled == GL_TRUE) << "Shader compilation error: " << [shader] {
GLint length;
glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &length);
std::string str;
str.reserve(length);
glGetShaderInfoLog(shader, length, nullptr, str.data());
return str;
}();
decode_program_ = glCreateProgram();
glAttachShader(decode_program_, shader);
glDeleteShader(shader);
@@ -910,7 +919,7 @@ void main() {
scored_boxes_buffer_ = absl::make_unique<Tensor>(
Tensor::ElementType::kFloat32, Tensor::Shape{1, num_boxes_ * 2});
return mediapipe::OkStatus();
return absl::OkStatus();
}));
#elif MEDIAPIPE_METAL_ENABLED
@@ -1128,7 +1137,7 @@ kernel void scoreKernel(
#endif // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
return mediapipe::OkStatus();
return absl::OkStatus();
}
} // namespace api2
mediapipe/calculators/tensor/tensors_to_floats_calculator.cc
+9 -10
View File
@@ -53,28 +53,27 @@ class TensorsToFloatsCalculator : public Node {
MEDIAPIPE_NODE_INTERFACE(TensorsToFloatsCalculator, kInTensors, kOutFloat,
kOutFloats);
static mediapipe::Status UpdateContract(CalculatorContract* cc);
mediapipe::Status Open(CalculatorContext* cc) final;
mediapipe::Status Process(CalculatorContext* cc) final;
static absl::Status UpdateContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc) final;
absl::Status Process(CalculatorContext* cc) final;
private:
::mediapipe::TensorsToFloatsCalculatorOptions options_;
};
MEDIAPIPE_REGISTER_NODE(TensorsToFloatsCalculator);
mediapipe::Status TensorsToFloatsCalculator::UpdateContract(
CalculatorContract* cc) {
absl::Status TensorsToFloatsCalculator::UpdateContract(CalculatorContract* cc) {
// Only exactly a single output allowed.
RET_CHECK(kOutFloat(cc).IsConnected() ^ kOutFloats(cc).IsConnected());
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToFloatsCalculator::Open(CalculatorContext* cc) {
absl::Status TensorsToFloatsCalculator::Open(CalculatorContext* cc) {
options_ = cc->Options<::mediapipe::TensorsToFloatsCalculatorOptions>();
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToFloatsCalculator::Process(CalculatorContext* cc) {
absl::Status TensorsToFloatsCalculator::Process(CalculatorContext* cc) {
const auto& input_tensors = *kInTensors(cc);
RET_CHECK(!input_tensors.empty());
// TODO: Add option to specify which tensor to take from.
@@ -101,7 +100,7 @@ mediapipe::Status TensorsToFloatsCalculator::Process(CalculatorContext* cc) {
} else {
kOutFloats(cc).Send(std::move(output_floats));
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/tensors_to_landmarks_calculator.cc
+12 -13
View File
@@ -100,23 +100,23 @@ class TensorsToLandmarksCalculator : public Node {
MEDIAPIPE_NODE_CONTRACT(kInTensors, kFlipHorizontally, kFlipVertically,
kOutLandmarkList, kOutNormalizedLandmarkList);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
private:
mediapipe::Status LoadOptions(CalculatorContext* cc);
absl::Status LoadOptions(CalculatorContext* cc);
int num_landmarks_ = 0;
::mediapipe::TensorsToLandmarksCalculatorOptions options_;
};
MEDIAPIPE_REGISTER_NODE(TensorsToLandmarksCalculator);
mediapipe::Status TensorsToLandmarksCalculator::Open(CalculatorContext* cc) {
absl::Status TensorsToLandmarksCalculator::Open(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(LoadOptions(cc));
if (kOutNormalizedLandmarkList(cc).IsConnected()) {
RET_CHECK(options_.has_input_image_height() &&
options_.has_input_image_width())
<< "Must provide input with/height for getting normalized landmarks.";
<< "Must provide input width/height for getting normalized landmarks.";
}
if (kOutLandmarkList(cc).IsConnected() &&
(options_.flip_horizontally() || options_.flip_vertically() ||
@@ -124,15 +124,15 @@ mediapipe::Status TensorsToLandmarksCalculator::Open(CalculatorContext* cc) {
kFlipVertically(cc).IsConnected())) {
RET_CHECK(options_.has_input_image_height() &&
options_.has_input_image_width())
<< "Must provide input with/height for using flipping when outputing "
<< "Must provide input width/height for using flipping when outputing "
"landmarks in absolute coordinates.";
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToLandmarksCalculator::Process(CalculatorContext* cc) {
absl::Status TensorsToLandmarksCalculator::Process(CalculatorContext* cc) {
if (kInTensors(cc).IsEmpty()) {
return mediapipe::OkStatus();
return absl::OkStatus();
}
bool flip_horizontally =
kFlipHorizontally(cc).GetOr(options_.flip_horizontally());
@@ -204,17 +204,16 @@ mediapipe::Status TensorsToLandmarksCalculator::Process(CalculatorContext* cc) {
kOutLandmarkList(cc).Send(std::move(output_landmarks));
}
return mediapipe::OkStatus();
return absl::OkStatus();
}
mediapipe::Status TensorsToLandmarksCalculator::LoadOptions(
CalculatorContext* cc) {
absl::Status TensorsToLandmarksCalculator::LoadOptions(CalculatorContext* cc) {
// Get calculator options specified in the graph.
options_ = cc->Options<::mediapipe::TensorsToLandmarksCalculatorOptions>();
RET_CHECK(options_.has_num_landmarks());
num_landmarks_ = options_.num_landmarks();
return mediapipe::OkStatus();
return absl::OkStatus();
}
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/testdata/expected_detection.pbtxt
Vendored
+35
View File
@@ -0,0 +1,35 @@
label_id: 0
score: 0.92843366
location_data {
format: RELATIVE_BOUNDING_BOX
relative_bounding_box {
xmin: 0.21061149
ymin: 0.29150677
width: 0.5657704
height: 0.5657307
}
relative_keypoints {
x: 0.37730268
y: 0.44038114
}
relative_keypoints {
x: 0.6250565
y: 0.44425336
}
relative_keypoints {
x: 0.50687385
y: 0.5767085
}
relative_keypoints {
x: 0.50173956
y: 0.6991459
}
relative_keypoints {
x: 0.2383742
y: 0.49879026
}
relative_keypoints {
x: 0.7404449
y: 0.50361776
}
}
mediapipe/calculators/tensor/testdata/face_detection_expected.png
Vendored
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 475 KiB

mediapipe/calculators/tensor/testdata/face_detection_test.pbtxt
Vendored
+31
View File
@@ -0,0 +1,31 @@
input_stream: "image"
output_stream: "rendering"
output_stream: "detections"
# Subgraph that detects faces.
node {
calculator: "FaceDetectionFrontCpu"
input_stream: "IMAGE:image"
output_stream: "DETECTIONS:detections"
}
# Converts the detections to drawing primitives for annotation overlay.
node {
calculator: "DetectionsToRenderDataCalculator"
input_stream: "DETECTIONS:detections"
output_stream: "RENDER_DATA:render_data"
options: {
[mediapipe.DetectionsToRenderDataCalculatorOptions.ext] {
thickness: 4.0
color { r: 255 g: 0 b: 0 }
}
}
}
# Draws annotations and overlays them on top of the input images.
node {
calculator: "AnnotationOverlayCalculator"
input_stream: "IMAGE:image"
input_stream: "render_data"
output_stream: "IMAGE:rendering"
}