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

No public description

Browse Source 
PiperOrigin-RevId: 586813896
This commit is contained in:
Sebastian Schmidt 2023-11-30 16:01:16 -08:00 committed by Copybara-Service
parent 3433ba083a
commit 7013b23785
1 changed files with 21 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
mediapipe/tasks/cc/vision/face_stylizer/calculators/tensors_to_image_calculator.cc
View File

@ -33,6 +33,7 @@
#include "mediapipe/framework/port/opencv_imgproc_inc.h" #include "mediapipe/framework/port/opencv_imgproc_inc.h"
#include "mediapipe/framework/port/status.h" #include "mediapipe/framework/port/status.h"
#include "mediapipe/gpu/gpu_origin.pb.h" #include "mediapipe/gpu/gpu_origin.pb.h"
#include "mediapipe/gpu/gpu_service.h"
#include "mediapipe/tasks/cc/vision/face_stylizer/calculators/tensors_to_image_calculator.pb.h" #include "mediapipe/tasks/cc/vision/face_stylizer/calculators/tensors_to_image_calculator.pb.h"
#if !MEDIAPIPE_DISABLE_GPU #if !MEDIAPIPE_DISABLE_GPU
@ -145,7 +146,8 @@ absl::Status TensorsToImageCalculator::UpdateContract(CalculatorContract* cc) {
#if MEDIAPIPE_METAL_ENABLED #if MEDIAPIPE_METAL_ENABLED
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]); MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#else #else
return GlCalculatorHelper::UpdateContract(cc); return GlCalculatorHelper::UpdateContract(cc,
/*requesst_gpu_as_optional=*/true);
#endif // MEDIAPIPE_METAL_ENABLED #endif // MEDIAPIPE_METAL_ENABLED
#endif // !MEDIAPIPE_DISABLE_GPU #endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus(); return absl::OkStatus();
@ -153,16 +155,7 @@ absl::Status TensorsToImageCalculator::UpdateContract(CalculatorContract* cc) {
absl::Status TensorsToImageCalculator::Open(CalculatorContext* cc) { absl::Status TensorsToImageCalculator::Open(CalculatorContext* cc) {
options_ = cc->Options<TensorsToImageCalculatorOptions>(); options_ = cc->Options<TensorsToImageCalculatorOptions>();
if (CanUseGpu()) { if (!CanUseGpu()) {
#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
} else {
ABSL_CHECK(options_.has_input_tensor_float_range() ^ ABSL_CHECK(options_.has_input_tensor_float_range() ^
options_.has_input_tensor_uint_range()) options_.has_input_tensor_uint_range())
<< "Must specify either `input_tensor_float_range` or " << "Must specify either `input_tensor_float_range` or "
@ -179,7 +172,9 @@ absl::Status TensorsToImageCalculator::Process(CalculatorContext* cc) {
#if MEDIAPIPE_METAL_ENABLED #if MEDIAPIPE_METAL_ENABLED
return MetalProcess(cc); return MetalProcess(cc);
#else #else
return GlProcess(cc); if (cc->Service(kGpuService).IsAvailable()) {
return GlProcess(cc);
}
#endif // MEDIAPIPE_METAL_ENABLED #endif // MEDIAPIPE_METAL_ENABLED
#endif // !MEDIAPIPE_DISABLE_GPU #endif // !MEDIAPIPE_DISABLE_GPU
} }
@ -188,14 +183,16 @@ absl::Status TensorsToImageCalculator::Process(CalculatorContext* cc) {
absl::Status TensorsToImageCalculator::Close(CalculatorContext* cc) { absl::Status TensorsToImageCalculator::Close(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU && !MEDIAPIPE_METAL_ENABLED #if !MEDIAPIPE_DISABLE_GPU && !MEDIAPIPE_METAL_ENABLED
gl_helper_.RunInGlContext([this] { if (gl_initialized_) {
gl_helper_.RunInGlContext([this] {
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31 #if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
gl_compute_program_.reset(); gl_compute_program_.reset();
#else #else
if (program_) glDeleteProgram(program_); if (program_) glDeleteProgram(program_);
program_ = 0; program_ = 0;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31 #endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
}); });
}
#endif // !MEDIAPIPE_DISABLE_GPU && !MEDIAPIPE_METAL_ENABLED #endif // !MEDIAPIPE_DISABLE_GPU && !MEDIAPIPE_METAL_ENABLED
return absl::OkStatus(); return absl::OkStatus();
} }
@ -315,6 +312,9 @@ absl::Status TensorsToImageCalculator::MetalProcess(CalculatorContext* cc) {
} }
absl::Status TensorsToImageCalculator::MetalSetup(CalculatorContext* cc) { absl::Status TensorsToImageCalculator::MetalSetup(CalculatorContext* cc) {
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(gpu_helper_);
id<MTLDevice> device = gpu_helper_.mtlDevice; id<MTLDevice> device = gpu_helper_.mtlDevice;
const std::string shader_source = const std::string shader_source =
R"( R"(
@ -450,6 +450,10 @@ absl::Status TensorsToImageCalculator::GlSetup(CalculatorContext* cc) {
} }
absl::Status TensorsToImageCalculator::GlProcess(CalculatorContext* cc) { absl::Status TensorsToImageCalculator::GlProcess(CalculatorContext* cc) {
if (!gl_initialized_) {
MP_RETURN_IF_ERROR(gl_helper_.Open(cc));
}
return gl_helper_.RunInGlContext([this, cc]() -> absl::Status { return gl_helper_.RunInGlContext([this, cc]() -> absl::Status {
if (!gl_initialized_) { if (!gl_initialized_) {
MP_RETURN_IF_ERROR(GlSetup(cc)); MP_RETURN_IF_ERROR(GlSetup(cc));