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

No public description

Browse Source 
PiperOrigin-RevId: 572266397
This commit is contained in:
Sebastian Schmidt 2023-10-10 09:00:31 -07:00 committed by Copybara-Service
parent ef6e712a88
commit dc63a5401c
9 changed files with 83 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mediapipe/calculators/image/BUILD
View File

@ -786,6 +786,7 @@ cc_library(
":affine_transformation_runner_gl",
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:gpu_service",
],
}) + select({
"//mediapipe/framework/port:disable_opencv": [],

11
mediapipe/calculators/image/image_clone_calculator.cc
View File

@ -64,7 +64,8 @@ class ImageCloneCalculator : public Node {
"GPU processing is disabled in build flags");
}
#else
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
cc, /*requesst_gpu_as_optional=*/true));
#endif // MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
@ -72,9 +73,6 @@ class ImageCloneCalculator : public Node {
absl::Status Open(CalculatorContext* cc) override {
const auto& options = cc->Options<mediapipe::ImageCloneCalculatorOptions>();
output_on_gpu_ = options.output_on_gpu();
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
@ -104,6 +102,10 @@ class ImageCloneCalculator : public Node {
if (output_on_gpu_ && !input_on_gpu) {
#if !MEDIAPIPE_DISABLE_GPU
if (!gpu_initialized_) {
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
gpu_initialized_ = true;
}
gpu_helper_.RunInGlContext([&output]() { output->ConvertToGpu(); });
#endif // !MEDIAPIPE_DISABLE_GPU
} else if (!output_on_gpu_ && input_on_gpu) {
@ -118,6 +120,7 @@ class ImageCloneCalculator : public Node {
bool output_on_gpu_;
#if !MEDIAPIPE_DISABLE_GPU
mediapipe::GlCalculatorHelper gpu_helper_;
bool gpu_initialized_ = false;
#endif // !MEDIAPIPE_DISABLE_GPU
};
MEDIAPIPE_REGISTER_NODE(ImageCloneCalculator);

20
mediapipe/calculators/image/segmentation_smoothing_calculator.cc
View File

@ -117,7 +117,8 @@ absl::Status SegmentationSmoothingCalculator::GetContract(
cc->Outputs().Tag(kOutputMaskTag).Set<Image>();
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
cc, /*requesst_gpu_as_optional=*/true));
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
@ -130,10 +131,6 @@ absl::Status SegmentationSmoothingCalculator::Open(CalculatorContext* cc) {
cc->Options<mediapipe::SegmentationSmoothingCalculatorOptions>();
combine_with_previous_ratio_ = options.combine_with_previous_ratio();
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
@ -154,6 +151,9 @@ absl::Status SegmentationSmoothingCalculator::Process(CalculatorContext* cc) {
if (use_gpu) {
#if !MEDIAPIPE_DISABLE_GPU
if (!gpu_initialized_) {
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
}
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, cc]() -> absl::Status {
if (!gpu_initialized_) {
MP_RETURN_IF_ERROR(GlSetup(cc));
@ -178,10 +178,12 @@ absl::Status SegmentationSmoothingCalculator::Process(CalculatorContext* cc) {
absl::Status SegmentationSmoothingCalculator::Close(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
gpu_helper_.RunInGlContext([this] {
if (program_) glDeleteProgram(program_);
program_ = 0;
});
if (gpu_initialized_) {
gpu_helper_.RunInGlContext([this] {
if (program_) glDeleteProgram(program_);
program_ = 0;
});
}
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();

24
mediapipe/calculators/image/warp_affine_calculator.cc
View File

@ -36,6 +36,7 @@
#if !MEDIAPIPE_DISABLE_GPU
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/gpu_service.h"
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
@ -106,6 +107,7 @@ class WarpAffineRunnerHolder<mediapipe::GpuBuffer> {
cc->Options<mediapipe::WarpAffineCalculatorOptions>().interpolation());
return gl_helper_->Open(cc);
}
absl::StatusOr<RunnerType*> GetRunner() {
if (!runner_) {
MP_ASSIGN_OR_RETURN(
@ -142,7 +144,10 @@ class WarpAffineRunnerHolder<mediapipe::Image> {
MP_RETURN_IF_ERROR(cpu_holder_.Open(cc));
#endif // !MEDIAPIPE_DISABLE_OPENCV
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(gpu_holder_.Open(cc));
if (cc->Service(kGpuService).IsAvailable()) {
MP_RETURN_IF_ERROR(gpu_holder_.Open(cc));
gpu_holder_initialized_ = true;
}
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}
@ -151,6 +156,9 @@ class WarpAffineRunnerHolder<mediapipe::Image> {
const AffineTransformation::Size& size,
AffineTransformation::BorderMode border_mode) override {
if (input.UsesGpu()) {
if (!gpu_holder_initialized_) {
return absl::UnavailableError("GPU support is not available");
}
#if !MEDIAPIPE_DISABLE_GPU
MP_ASSIGN_OR_RETURN(auto* runner, gpu_holder_.GetRunner());
MP_ASSIGN_OR_RETURN(
@ -183,6 +191,7 @@ class WarpAffineRunnerHolder<mediapipe::Image> {
#endif // !MEDIAPIPE_DISABLE_OPENCV
#if !MEDIAPIPE_DISABLE_GPU
WarpAffineRunnerHolder<mediapipe::GpuBuffer> gpu_holder_;
bool gpu_holder_initialized_ = false;
#endif // !MEDIAPIPE_DISABLE_GPU
};
@ -196,20 +205,24 @@ class WarpAffineCalculatorImpl : public mediapipe::api2::NodeImpl<InterfaceT> {
static absl::Status UpdateContract(CalculatorContract* cc) {
if constexpr (std::is_same_v<InterfaceT, WarpAffineCalculatorGpu> ||
std::is_same_v<InterfaceT, WarpAffineCalculator>) {
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
cc, /*requesst_gpu_as_optional=*/true));
}
return absl::OkStatus();
}
#endif // !MEDIAPIPE_DISABLE_GPU
absl::Status Open(CalculatorContext* cc) override { return holder_.Open(cc); }
absl::Status Process(CalculatorContext* cc) override {
if (InterfaceT::kInImage(cc).IsEmpty() ||
InterfaceT::kMatrix(cc).IsEmpty() ||
InterfaceT::kOutputSize(cc).IsEmpty()) {
return absl::OkStatus();
}
if (!holder_initialized_) {
MP_RETURN_IF_ERROR(holder_.Open(cc));
holder_initialized_ = true;
}
const std::array<float, 16>& transform = *InterfaceT::kMatrix(cc);
auto [out_width, out_height] = *InterfaceT::kOutputSize(cc);
AffineTransformation::Size output_size;
@ -230,6 +243,7 @@ class WarpAffineCalculatorImpl : public mediapipe::api2::NodeImpl<InterfaceT> {
private:
WarpAffineRunnerHolder<typename decltype(InterfaceT::kInImage)::PayloadT>
holder_;
bool holder_initialized_ = false;
};
} // namespace

22
mediapipe/calculators/tensor/tensors_to_detections_calculator.cc
View File

@ -266,7 +266,8 @@ absl::Status TensorsToDetectionsCalculator::UpdateContract(
CalculatorContract* cc) {
if (CanUseGpu()) {
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
cc, /*requesst_gpu_as_optional=*/true));
#elif MEDIAPIPE_METAL_ENABLED
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#endif // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
@ -280,7 +281,6 @@ absl::Status TensorsToDetectionsCalculator::Open(CalculatorContext* cc) {
if (CanUseGpu()) {
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#elif MEDIAPIPE_METAL_ENABLED
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(gpu_helper_);
@ -676,13 +676,15 @@ absl::Status TensorsToDetectionsCalculator::ProcessGPU(
absl::Status TensorsToDetectionsCalculator::Close(CalculatorContext* cc) {
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
gpu_helper_.RunInGlContext([this] {
decoded_boxes_buffer_ = nullptr;
scored_boxes_buffer_ = nullptr;
raw_anchors_buffer_ = nullptr;
glDeleteProgram(decode_program_);
glDeleteProgram(score_program_);
});
if (gpu_inited_) {
gpu_helper_.RunInGlContext([this] {
decoded_boxes_buffer_ = nullptr;
scored_boxes_buffer_ = nullptr;
raw_anchors_buffer_ = nullptr;
glDeleteProgram(decode_program_);
glDeleteProgram(score_program_);
});
}
#elif MEDIAPIPE_METAL_ENABLED
decoded_boxes_buffer_ = nullptr;
scored_boxes_buffer_ = nullptr;
@ -942,6 +944,7 @@ absl::Status TensorsToDetectionsCalculator::GpuInit(CalculatorContext* cc) {
break;
}
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, output_format_flag]()
-> absl::Status {
// A shader to decode detection boxes.
@ -1420,7 +1423,6 @@ kernel void scoreKernel(
num_classes_, max_wg_size));
}
}
#endif // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
return absl::OkStatus();

31
mediapipe/calculators/tensor/tensors_to_segmentation_calculator.cc
View File

@ -173,6 +173,7 @@ class TensorsToSegmentationCalculator : public CalculatorBase {
#if !MEDIAPIPE_DISABLE_GPU
mediapipe::GlCalculatorHelper gpu_helper_;
GLuint upsample_program_;
bool gpu_initialized_ = false;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
int cached_width_ = 0;
int cached_height_ = 0;
@ -206,7 +207,8 @@ absl::Status TensorsToSegmentationCalculator::GetContract(
if (CanUseGpu()) {
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(
cc, /*requesst_gpu_as_optional=*/true));
#if MEDIAPIPE_METAL_ENABLED
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#endif // MEDIAPIPE_METAL_ENABLED
@ -218,12 +220,9 @@ absl::Status TensorsToSegmentationCalculator::GetContract(
absl::Status TensorsToSegmentationCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
bool use_gpu = false;
if (CanUseGpu()) {
#if !MEDIAPIPE_DISABLE_GPU
use_gpu = true;
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#if MEDIAPIPE_METAL_ENABLED
metal_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(metal_helper_);
@ -233,14 +232,6 @@ absl::Status TensorsToSegmentationCalculator::Open(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(LoadOptions(cc));
if (use_gpu) {
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(InitGpu(cc));
#else
RET_CHECK_FAIL() << "GPU processing disabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
}
return absl::OkStatus();
}
@ -285,6 +276,15 @@ absl::Status TensorsToSegmentationCalculator::Process(CalculatorContext* cc) {
}
if (use_gpu) {
#if !MEDIAPIPE_DISABLE_GPU
if (!gpu_initialized_) {
MP_RETURN_IF_ERROR(InitGpu(cc));
gpu_initialized_ = true;
}
#else
RET_CHECK_FAIL() << "GPU processing disabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, cc]() -> absl::Status {
MP_RETURN_IF_ERROR(ProcessGpu(cc));
@ -306,6 +306,10 @@ absl::Status TensorsToSegmentationCalculator::Process(CalculatorContext* cc) {
absl::Status TensorsToSegmentationCalculator::Close(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
if (!gpu_initialized_) {
return absl::OkStatus();
}
gpu_helper_.RunInGlContext([this] {
if (upsample_program_) glDeleteProgram(upsample_program_);
upsample_program_ = 0;
@ -635,6 +639,7 @@ absl::Status TensorsToSegmentationCalculator::LoadOptions(
absl::Status TensorsToSegmentationCalculator::InitGpu(CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> absl::Status {
// A shader to process a segmentation tensor into an output mask.
// Currently uses 4 channels for output, and sets R+A channels as mask value.
@ -899,6 +904,8 @@ void main() {
return absl::OkStatus();
}));
gpu_initialized_ = true;
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();

9
mediapipe/gpu/gl_calculator_helper.cc
View File

@ -56,8 +56,13 @@ void GlCalculatorHelper::InitializeForTest(GpuResources* gpu_resources) {
}
// static
absl::Status GlCalculatorHelper::UpdateContract(CalculatorContract* cc) {
cc->UseService(kGpuService);
absl::Status GlCalculatorHelper::UpdateContract(CalculatorContract* cc,
bool requesst_gpu_as_optional) {
if (requesst_gpu_as_optional) {
cc->UseService(kGpuService).Optional();
} else {
cc->UseService(kGpuService);
}
// Allow the legacy side packet to be provided, too, for backwards
// compatibility with existing graphs. It will just be ignored.
auto& input_side_packets = cc->InputSidePackets();

3
mediapipe/gpu/gl_calculator_helper.h
View File

@ -67,7 +67,8 @@ class GlCalculatorHelper {
// This method can be called from GetContract to set up the needed GPU
// resources.
static absl::Status UpdateContract(CalculatorContract* cc);
static absl::Status UpdateContract(CalculatorContract* cc,
bool requesst_gpu_as_optional = false);
// This method can be called from FillExpectations to set the correct types
// for the shared GL input side packet(s).

5
mediapipe/python/pybind/calculator_graph.cc
View File

@ -112,6 +112,11 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
"graph with a ValidatedGraphConfig object.");
}
auto calculator_graph = absl::make_unique<CalculatorGraph>();
// Disable default service initialization. This allows us to use
// the CPU versions of calculators that only optionally request
// kGpuService.
RaisePyErrorIfNotOk(
calculator_graph->DisallowServiceDefaultInitialization());
RaisePyErrorIfNotOk(calculator_graph->Initialize(graph_config_proto));
return calculator_graph.release();
}),