cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Implement TensorsToImageCalculator.

Browse Source 
PiperOrigin-RevId: 512687103
This commit is contained in:
Jiuqiang Tang 2023-02-27 11:41:25 -08:00 committed by Copybara-Service
parent 5913e5dd1d
commit aa61abe386
3 changed files with 578 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

108
mediapipe/tasks/cc/vision/face_stylizer/calculators/BUILD Normal file
View File

@ -0,0 +1,108 @@
# Copyright 2023 The MediaPipe Authors. All Rights Reserved.
#
# 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.
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_proto_library")
licenses(["notice"])
package(default_visibility = ["//mediapipe/tasks:internal"])
mediapipe_proto_library(
name = "tensors_to_image_calculator_proto",
srcs = ["tensors_to_image_calculator.proto"],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
"//mediapipe/gpu:gpu_origin_proto",
],
)
cc_library(
name = "tensors_to_image_calculator",
srcs = ["tensors_to_image_calculator.cc"],
copts = select({
"//mediapipe:apple": [
"-x objective-c++",
"-fobjc-arc", # enable reference-counting
],
"//conditions:default": [],
}),
features = ["-layering_check"], # allow depending on tensor_to_image_calculator_gpu_deps
linkopts = select({
"//mediapipe:apple": [
"-framework CoreVideo",
"-framework MetalKit",
],
"//conditions:default": [],
}),
deps = [
":tensors_to_image_calculator_cc_proto",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_options_cc_proto",
"//mediapipe/framework/api2:builder",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/api2:packet",
"//mediapipe/framework/api2:port",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:vector",
"//mediapipe/gpu:gpu_origin_cc_proto",
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": ["tensor_to_image_calculator_gpu_deps"],
}),
alwayslink = 1,
)
cc_library(
name = "tensor_to_image_calculator_gpu_deps",
visibility = ["//visibility:private"],
deps = select({
"//mediapipe:android": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_quad_renderer",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:util",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_texture",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl/converters:util",
],
"//mediapipe:ios": [
"//mediapipe/gpu:MPPMetalHelper",
"//mediapipe/gpu:MPPMetalUtil",
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gpu_buffer",
],
"//mediapipe:macos": [],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_quad_renderer",
"//mediapipe/gpu:gpu_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:util",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_texture",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl/converters:util",
],
}),
)

439
mediapipe/tasks/cc/vision/face_stylizer/calculators/tensors_to_image_calculator.cc Normal file
View File

@ -0,0 +1,439 @@
// Copyright 2023 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/api2/packet.h"
#include "mediapipe/framework/api2/port.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_options.pb.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/gpu/gpu_origin.pb.h"
#include "mediapipe/tasks/cc/vision/face_stylizer/calculators/tensors_to_image_calculator.pb.h"
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/gpu/gpu_buffer.h"
#if MEDIAPIPE_METAL_ENABLED
#import <CoreVideo/CoreVideo.h>
#import <Metal/Metal.h>
#import <MetalKit/MetalKit.h>
#include "mediapipe/framework/formats/tensor_mtl_buffer_view.h"
#import "mediapipe/gpu/MPPMetalHelper.h"
#else
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_quad_renderer.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
#include "tensorflow/lite/delegates/gpu/common/util.h"
#include "tensorflow/lite/delegates/gpu/gl/converters/util.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_texture.h"
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
#endif // MEDIAPIPE_METAL_ENABLED
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
namespace tasks {
namespace {
using ::mediapipe::api2::Input;
using ::mediapipe::api2::Node;
using ::mediapipe::api2::Output;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
using ::tflite::gpu::gl::GlProgram;
using ::tflite::gpu::gl::GlShader;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
// Commonly used to compute the number of blocks to launch in a kernel.
static int NumGroups(const int size, const int group_size) { // NOLINT
return (size + group_size - 1) / group_size;
}
} // namespace
// Converts a MediaPipe tensor to a MediaPipe Image.
//
// Input streams:
// TENSORS - std::vector<mediapipe::Tensor> that only contains one element.
//
// Output streams:
// OUTPUT - mediapipe::Image.
//
// TODO: Enable TensorsToImageCalculator to run on CPU.
class TensorsToImageCalculator : public Node {
public:
static constexpr Input<std::vector<Tensor>> kInputTensors{"TENSORS"};
static constexpr Output<Image> kOutputImage{"IMAGE"};
MEDIAPIPE_NODE_CONTRACT(kInputTensors, kOutputImage);
static absl::Status GetContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc);
absl::Status Process(CalculatorContext* cc);
absl::Status Close(CalculatorContext* cc);
private:
#if !MEDIAPIPE_DISABLE_GPU
#if MEDIAPIPE_METAL_ENABLED
bool metal_initialized_ = false;
MPPMetalHelper* gpu_helper_ = nullptr;
id<MTLComputePipelineState> to_buffer_program_;
absl::Status MetalSetup(CalculatorContext* cc);
absl::Status MetalProcess(CalculatorContext* cc);
#else
absl::Status GlSetup(CalculatorContext* cc);
GlCalculatorHelper gl_helper_;
bool gl_initialized_ = false;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
std::unique_ptr<tflite::gpu::gl::GlProgram> gl_compute_program_;
const tflite::gpu::uint3 workgroup_size_ = {8, 8, 1};
#else
GLuint program_ = 0;
std::unique_ptr<mediapipe::QuadRenderer> gl_renderer_;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
#endif // MEDIAPIPE_METAL_ENABLED
#endif // !MEDIAPIPE_DISABLE_GPU
};
MEDIAPIPE_REGISTER_NODE(::mediapipe::tasks::TensorsToImageCalculator);
absl::Status TensorsToImageCalculator::GetContract(CalculatorContract* cc) {
#if !MEDIAPIPE_DISABLE_GPU
#if MEDIAPIPE_METAL_ENABLED
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#else
return GlCalculatorHelper::UpdateContract(cc);
#endif // MEDIAPIPE_METAL_ENABLED
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
absl::Status TensorsToImageCalculator::Open(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
#if MEDIAPIPE_METAL_ENABLED
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(gpu_helper_);
#else
MP_RETURN_IF_ERROR(gl_helper_.Open(cc));
#endif // MEDIAPIPE_METAL_ENABLED
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
absl::Status TensorsToImageCalculator::Process(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
#if MEDIAPIPE_METAL_ENABLED
return MetalProcess(cc);
#else
return gl_helper_.RunInGlContext([this, cc]() -> absl::Status {
if (!gl_initialized_) {
MP_RETURN_IF_ERROR(GlSetup(cc));
gl_initialized_ = true;
}
if (kInputTensors(cc).IsEmpty()) {
return absl::OkStatus();
}
const auto& input_tensors = kInputTensors(cc).Get();
RET_CHECK_EQ(input_tensors.size(), 1)
<< "Expect 1 input tensor, but have " << input_tensors.size();
const int tensor_width = input_tensors[0].shape().dims[2];
const int tensor_height = input_tensors[0].shape().dims[1];
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
auto out_texture = std::make_unique<tflite::gpu::gl::GlTexture>();
MP_RETURN_IF_ERROR(CreateReadWriteRgbaImageTexture(
tflite::gpu::DataType::UINT8, // GL_RGBA8
{tensor_width, tensor_height}, out_texture.get()));
const int output_index = 0;
glBindImageTexture(output_index, out_texture->id(), 0, GL_FALSE, 0,
GL_WRITE_ONLY, GL_RGBA8);
auto read_view = input_tensors[0].GetOpenGlBufferReadView();
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, read_view.name());
const tflite::gpu::uint3 workload = {tensor_width, tensor_height, 1};
const tflite::gpu::uint3 workgroups =
tflite::gpu::DivideRoundUp(workload, workgroup_size_);
glUseProgram(gl_compute_program_->id());
glUniform2i(glGetUniformLocation(gl_compute_program_->id(), "out_size"),
tensor_width, tensor_height);
MP_RETURN_IF_ERROR(gl_compute_program_->Dispatch(workgroups));
auto texture_buffer = mediapipe::GlTextureBuffer::Wrap(
out_texture->target(), out_texture->id(), tensor_width, tensor_height,
mediapipe::GpuBufferFormat::kBGRA32,
[ptr = out_texture.release()](
std::shared_ptr<mediapipe::GlSyncPoint> sync_token) mutable {
delete ptr;
});
auto output =
std::make_unique<mediapipe::GpuBuffer>(std::move(texture_buffer));
kOutputImage(cc).Send(Image(*output));
;
#else
if (!input_tensors[0].ready_as_opengl_texture_2d()) {
(void)input_tensors[0].GetCpuReadView();
}
auto output_texture =
gl_helper_.CreateDestinationTexture(tensor_width, tensor_height);
gl_helper_.BindFramebuffer(output_texture); // GL_TEXTURE0
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D,
input_tensors[0].GetOpenGlTexture2dReadView().name());
MP_RETURN_IF_ERROR(gl_renderer_->GlRender(
tensor_width, tensor_height, output_texture.width(),
output_texture.height(), mediapipe::FrameScaleMode::kStretch,
mediapipe::FrameRotation::kNone,
/*flip_horizontal=*/false, /*flip_vertical=*/false,
/*flip_texture=*/false));
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, 0);
auto output = output_texture.GetFrame<GpuBuffer>();
kOutputImage(cc).Send(Image(*output));
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
return mediapipe::OkStatus();
});
#endif // MEDIAPIPE_METAL_ENABLED
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
absl::Status TensorsToImageCalculator::Close(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU && !MEDIAPIPE_METAL_ENABLED
gl_helper_.RunInGlContext([this] {
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
gl_compute_program_.reset();
#else
if (program_) glDeleteProgram(program_);
program_ = 0;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
});
#endif // !MEDIAPIPE_DISABLE_GPU && !MEDIAPIPE_METAL_ENABLED
return absl::OkStatus();
}
#if MEDIAPIPE_METAL_ENABLED
absl::Status TensorsToImageCalculator::MetalProcess(CalculatorContext* cc) {
if (!metal_initialized_) {
MP_RETURN_IF_ERROR(MetalSetup(cc));
metal_initialized_ = true;
}
if (kInputTensors(cc).IsEmpty()) {
return absl::OkStatus();
}
const auto& input_tensors = kInputTensors(cc).Get();
RET_CHECK_EQ(input_tensors.size(), 1)
<< "Expect 1 input tensor, but have " << input_tensors.size();
const int tensor_width = input_tensors[0].shape().dims[2];
const int tensor_height = input_tensors[0].shape().dims[1];
// TODO: Fix unused variable
[[maybe_unused]] id<MTLDevice> device = gpu_helper_.mtlDevice;
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"TensorsToImageCalculatorConvert";
id<MTLComputeCommandEncoder> compute_encoder =
[command_buffer computeCommandEncoder];
[compute_encoder setComputePipelineState:to_buffer_program_];
auto input_view =
mediapipe::MtlBufferView::GetReadView(input_tensors[0], command_buffer);
[compute_encoder setBuffer:input_view.buffer() offset:0 atIndex:0];
mediapipe::GpuBuffer output =
[gpu_helper_ mediapipeGpuBufferWithWidth:tensor_width
height:tensor_height];
id<MTLTexture> dst_texture = [gpu_helper_ metalTextureWithGpuBuffer:output];
[compute_encoder setTexture:dst_texture atIndex:1];
MTLSize threads_per_group = MTLSizeMake(8, 8, 1);
MTLSize threadgroups =
MTLSizeMake(NumGroups(tensor_width, 8), NumGroups(tensor_height, 8), 1);
[compute_encoder dispatchThreadgroups:threadgroups
threadsPerThreadgroup:threads_per_group];
[compute_encoder endEncoding];
[command_buffer commit];
kOutputImage(cc).Send(Image(output));
return absl::OkStatus();
}
absl::Status TensorsToImageCalculator::MetalSetup(CalculatorContext* cc) {
id<MTLDevice> device = gpu_helper_.mtlDevice;
const std::string shader_source =
R"(
#include <metal_stdlib>
using namespace metal;
kernel void convertKernel(
device float* in_buf [[ buffer(0) ]],
texture2d<float, access::read_write> out_tex [[ texture(1) ]],
uint2 gid [[ thread_position_in_grid ]]) {
if (gid.x >= out_tex.get_width() || gid.y >= out_tex.get_height()) return;
uint linear_index = 3 * (gid.y * out_tex.get_width() + gid.x);
float4 out_value = float4(in_buf[linear_index], in_buf[linear_index + 1], in_buf[linear_index + 2], 1.0);
out_tex.write(out_value, gid);
}
)";
NSString* library_source =
[NSString stringWithUTF8String:shader_source.c_str()];
NSError* error = nil;
id<MTLLibrary> library =
[device newLibraryWithSource:library_source options:nullptr error:&error];
RET_CHECK(library != nil) << "Couldn't create shader library "
<< [[error localizedDescription] UTF8String];
id<MTLFunction> kernel_func = nil;
kernel_func = [library newFunctionWithName:@"convertKernel"];
RET_CHECK(kernel_func != nil) << "Couldn't create kernel function.";
to_buffer_program_ =
[device newComputePipelineStateWithFunction:kernel_func error:&error];
RET_CHECK(to_buffer_program_ != nil) << "Couldn't create pipeline state " <<
[[error localizedDescription] UTF8String];
return mediapipe::OkStatus();
}
#endif // MEDIAPIPE_METAL_ENABLED
#if !MEDIAPIPE_DISABLE_GPU && !MEDIAPIPE_METAL_ENABLED
absl::Status TensorsToImageCalculator::GlSetup(CalculatorContext* cc) {
std::string maybe_flip_y_define;
#if !defined(__APPLE__)
const auto& options = cc->Options<TensorsToImageCalculatorOptions>();
if (options.gpu_origin() != mediapipe::GpuOrigin::TOP_LEFT) {
maybe_flip_y_define = R"(
#define FLIP_Y_COORD
)";
}
#endif // !defined(__APPLE__)
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
const std::string shader_header =
absl::StrCat(tflite::gpu::gl::GetShaderHeader(workgroup_size_), R"(
precision highp float;
layout(rgba8, binding = 0) writeonly uniform highp image2D output_texture;
uniform ivec2 out_size;
)");
const std::string shader_body = R"(
layout(std430, binding = 2) readonly buffer B0 {
float elements[];
} input_data; // data tensor
void main() {
int out_width = out_size.x;
int out_height = out_size.y;
ivec2 gid = ivec2(gl_GlobalInvocationID.xy);
if (gid.x >= out_width || gid.y >= out_height) { return; }
int linear_index = 3 * (gid.y * out_width + gid.x);
#ifdef FLIP_Y_COORD
int y_coord = out_height - gid.y - 1;
#else
int y_coord = gid.y;
#endif // defined(FLIP_Y_COORD)
ivec2 out_coordinate = ivec2(gid.x, y_coord);
vec4 out_value = vec4(input_data.elements[linear_index], input_data.elements[linear_index + 1], input_data.elements[linear_index + 2], 1.0);
imageStore(output_texture, out_coordinate, out_value);
})";
const std::string shader_full =
absl::StrCat(shader_header, maybe_flip_y_define, shader_body);
GlShader shader;
MP_RETURN_IF_ERROR(
GlShader::CompileShader(GL_COMPUTE_SHADER, shader_full, &shader));
gl_compute_program_ = std::make_unique<GlProgram>();
MP_RETURN_IF_ERROR(
GlProgram::CreateWithShader(shader, gl_compute_program_.get()));
#else
constexpr GLchar kFragColorOutputDeclaration[] = R"(
#ifdef GL_ES
#define fragColor gl_FragColor
#else
out vec4 fragColor;
#endif // defined(GL_ES);
)";
constexpr GLchar kBody[] = R"(
DEFAULT_PRECISION(mediump, float)
in vec2 sample_coordinate;
uniform sampler2D tensor;
void main() {
#ifdef FLIP_Y_COORD
float y_coord = 1.0 - sample_coordinate.y;
#else
float y_coord = sample_coordinate.y;
#endif // defined(FLIP_Y_COORD)
vec3 color = texture2D(tensor, vec2(sample_coordinate.x, y_coord)).rgb;
fragColor = vec4(color, 1.0);
}
)";
const std::string src =
absl::StrCat(mediapipe::kMediaPipeFragmentShaderPreamble,
kFragColorOutputDeclaration, maybe_flip_y_define, kBody);
gl_renderer_ = std::make_unique<mediapipe::QuadRenderer>();
MP_RETURN_IF_ERROR(gl_renderer_->GlSetup(src.c_str(), {"tensor"}));
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
return mediapipe::OkStatus();
}
#endif // !MEDIAPIPE_DISABLE_GPU && !MEDIAPIPE_METAL_ENABLED
} // namespace tasks
} // namespace mediapipe

31
mediapipe/tasks/cc/vision/face_stylizer/calculators/tensors_to_image_calculator.proto Normal file
View File

@ -0,0 +1,31 @@
// Copyright 2023 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
syntax = "proto2";
package mediapipe.tasks;
import "mediapipe/framework/calculator.proto";
import "mediapipe/gpu/gpu_origin.proto";
message TensorsToImageCalculatorOptions {
extend mediapipe.CalculatorOptions {
optional TensorsToImageCalculatorOptions ext = 511831156;
}
// For CONVENTIONAL mode for OpenGL, input image starts at bottom and needs
// to be flipped vertically as tensors are expected to start at top.
// (DEFAULT or unset interpreted as CONVENTIONAL.)
optional mediapipe.GpuOrigin.Mode gpu_origin = 1;
}