Skip to content

Commit

Permalink
Clean up TensorConverterCalculator flipping behavior
Browse files Browse the repository at this point in the history 
Returns an error if
- gpu_origin is specified for a CPU image, and
- gpu_origin and flip_vertically are both specified.
Adds a test for an IMAGE_GPU input to validate flipping.

PiperOrigin-RevId: 565311456
  • Loading branch information
@copybara-github
MediaPipe Team authored and copybara-github committed Sep 14, 2023
1 parent 2164600 commit 124a4de
Show file tree
Hide file tree
Showing 4 changed files with 166 additions and 95 deletions.
6 changes: 6 additions & 0 deletions mediapipe/calculators/tensor/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -660,7 +660,12 @@ cc_library(
"//mediapipe/gpu:gpu_buffer_format",
"//mediapipe/gpu:gpu_origin_cc_proto",
"//mediapipe/util:resource_util",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:absl_check",
"@com_google_absl//absl/log:absl_log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
] + select({
"//mediapipe/gpu:disable_gpu": [],
Expand Down Expand Up @@ -715,6 +720,7 @@ cc_test(
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/tool:validate_type",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
Expand Down
190 changes: 105 additions & 85 deletions mediapipe/calculators/tensor/tensor_converter_calculator.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#include <vector>

#include "absl/log/absl_check.h"
#include "absl/log/absl_log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
Expand Down Expand Up @@ -57,11 +58,25 @@ int NumGroups(const int size, const int group_size) { // NOLINT
}

absl::StatusOr<bool> ShouldFlipVertically(
const mediapipe::TensorConverterCalculatorOptions& options) {
const mediapipe::TensorConverterCalculatorOptions& options, bool use_gpu) {
if (options.has_flip_vertically() && options.has_gpu_origin()) {
return absl::FailedPreconditionError(absl::StrFormat(
"Cannot specify both flip_vertically and gpu_origin options"));
}

if (!options.has_gpu_origin()) {
// Fall back to flip_vertically.
return options.flip_vertically();
}

// Warn if gpu_origin is specified with a CPU input image.
// Those are always TOP_LEFT, so no flipping is necessary.
if (!use_gpu) {
ABSL_LOG(WARNING)
<< "Ignoring gpu_origin option since IMAGE_GPU input is not specified";
return false;
}

switch (options.gpu_origin()) {
case mediapipe::GpuOrigin::TOP_LEFT:
return false;
Expand Down Expand Up @@ -140,7 +155,7 @@ class TensorConverterCalculator : public CalculatorBase {

private:
absl::Status InitGpu(CalculatorContext* cc);
absl::Status LoadOptions(CalculatorContext* cc);
absl::Status LoadOptions(CalculatorContext* cc, bool use_gpu);
template <class T>
absl::Status NormalizeImage(const ImageFrame& image_frame,
bool flip_vertically, float* tensor_ptr);
Expand Down Expand Up @@ -176,7 +191,8 @@ absl::Status TensorConverterCalculator::GetContract(CalculatorContract* cc) {
RET_CHECK(static_cast<int>(cc->Inputs().HasTag(kImageFrameTag)) +
static_cast<int>(cc->Inputs().HasTag(kGpuBufferTag)) +
static_cast<int>(cc->Inputs().HasTag(kMatrixTag)) ==
1);
1)
<< "Only one input tag of {IMAGE, IMAGE_GPU, MATRIX} may be specified";

if (cc->Inputs().HasTag(kImageFrameTag)) {
cc->Inputs().Tag(kImageFrameTag).Set<ImageFrame>();
Expand Down Expand Up @@ -204,8 +220,6 @@ absl::Status TensorConverterCalculator::GetContract(CalculatorContract* cc) {
absl::Status TensorConverterCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));

MP_RETURN_IF_ERROR(LoadOptions(cc));

#if !MEDIAPIPE_DISABLE_GPU
if (cc->Inputs().HasTag(kGpuBufferTag)) {
use_gpu_ = true;
Expand All @@ -218,6 +232,8 @@ absl::Status TensorConverterCalculator::Open(CalculatorContext* cc) {
}
#endif // !MEDIAPIPE_DISABLE_GPU

MP_RETURN_IF_ERROR(LoadOptions(cc, use_gpu_));

return absl::OkStatus();
}

Expand Down Expand Up @@ -436,7 +452,7 @@ absl::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
// Shader to convert GL Texture to Metal Buffer,
// with normalization to either: [0,1] or [-1,1].
const std::string shader_source = absl::Substitute(
R"(
R"glsl(
#include <metal_stdlib>
using namespace metal;
Expand All @@ -455,7 +471,7 @@ absl::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
$3 // g & b channels
$4 // alpha channel
}
)",
)glsl",
/*$0=*/
output_range_.has_value()
? absl::Substitute("pixel = pixel * half($0) + half($1);",
Expand All @@ -465,8 +481,8 @@ absl::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
/*$1=*/max_num_channels_,
/*$2=*/flip_vertically_ ? "(in_tex.get_height() - 1 - gid.y)" : "gid.y",
/*$3=*/
single_channel ? "" : R"(out_buf[linear_index + 1] = pixel.y;
out_buf[linear_index + 2] = pixel.z;)",
single_channel ? "" : R"glsl(out_buf[linear_index + 1] = pixel.y;
out_buf[linear_index + 2] = pixel.z;)glsl",
/*$4=*/include_alpha ? "out_buf[linear_index + 3] = pixel.w;" : "");

NSString* library_source =
Expand All @@ -484,17 +500,17 @@ absl::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
RET_CHECK(to_buffer_program_ != nil) << "Couldn't create pipeline state " <<
[[error localizedDescription] UTF8String];
#elif MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, &include_alpha,
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, &include_alpha,
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
&input,
&input,
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
&single_channel]()
-> absl::Status {
&single_channel]() -> 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].
const std::string shader_source = absl::Substitute(
R"( #version 310 es
// Shader to convert GL Texture to Shader Storage Buffer Object (SSBO),
// with normalization to either: [0,1] or [-1,1].
const std::string shader_source = absl::Substitute(
R"glsl( #version 310 es
layout(local_size_x = $0, local_size_y = $0) in;
layout(binding = 0) uniform sampler2D input_texture;
layout(std430, binding = 1) buffer Output {float elements[];} output_data;
Expand All @@ -508,38 +524,40 @@ absl::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
output_data.elements[linear_index + 0] = pixel.x; // r channel
$5 // g & b channels
$6 // alpha channel
})",
/*$0=*/kWorkgroupSize, /*$1=*/input.width(), /*$2=*/input.height(),
/*$3=*/
output_range_.has_value()
? absl::Substitute("pixel = pixel * float($0) + float($1);",
(output_range_->second - output_range_->first),
output_range_->first)
: "",
/*$4=*/flip_vertically_ ? "(width_height.y - 1 - gid.y)" : "gid.y",
/*$5=*/
single_channel ? ""
: R"(output_data.elements[linear_index + 1] = pixel.y;
output_data.elements[linear_index + 2] = pixel.z;)",
/*$6=*/
include_alpha ? "output_data.elements[linear_index + 3] = pixel.w;"
: "",
/*$7=*/max_num_channels_);
GLuint shader = glCreateShader(GL_COMPUTE_SHADER);
const GLchar* sources[] = {shader_source.c_str()};
glShaderSource(shader, 1, sources, NULL);
glCompileShader(shader);
GLint compiled = GL_FALSE;
glGetShaderiv(shader, GL_COMPILE_STATUS, &compiled);
RET_CHECK(compiled == GL_TRUE);
to_buffer_program_ = glCreateProgram();
glAttachShader(to_buffer_program_, shader);
glDeleteShader(shader);
glLinkProgram(to_buffer_program_);
})glsl",
/*$0=*/kWorkgroupSize, /*$1=*/input.width(), /*$2=*/input.height(),
/*$3=*/
output_range_.has_value()
? absl::Substitute(
"pixel = pixel * float($0) + float($1);",
(output_range_->second - output_range_->first),
output_range_->first)
: "",
/*$4=*/flip_vertically_ ? "(width_height.y - 1 - gid.y)" : "gid.y",
/*$5=*/
single_channel
? ""
: R"glsl(output_data.elements[linear_index + 1] = pixel.y;
output_data.elements[linear_index + 2] = pixel.z;)glsl",
/*$6=*/
include_alpha ? "output_data.elements[linear_index + 3] = pixel.w;"
: "",
/*$7=*/max_num_channels_);
GLuint shader = glCreateShader(GL_COMPUTE_SHADER);
const GLchar* sources[] = {shader_source.c_str()};
glShaderSource(shader, 1, sources, NULL);
glCompileShader(shader);
GLint compiled = GL_FALSE;
glGetShaderiv(shader, GL_COMPILE_STATUS, &compiled);
RET_CHECK(compiled == GL_TRUE);
to_buffer_program_ = glCreateProgram();
glAttachShader(to_buffer_program_, shader);
glDeleteShader(shader);
glLinkProgram(to_buffer_program_);
#else
// OpenGL ES 3.0 fragment shader Texture2d -> Texture2d conversion.
const std::string shader_source = absl::Substitute(
R"(
// OpenGL ES 3.0 fragment shader Texture2d -> Texture2d conversion.
const std::string shader_source = absl::Substitute(
R"glsl(
#if __VERSION__ < 130
#define in varying
#endif // __VERSION__ < 130
Expand All @@ -565,49 +583,51 @@ absl::Status TensorConverterCalculator::InitGpu(CalculatorContext* cc) {
fragColor.r = pixel.r; // r channel
$3 // g & b channels
$4 // alpha channel
})",
/*$0=*/single_channel ? "vec1" : "vec4",
/*$1=*/
flip_vertically_
? "vec2(sample_coordinate.x, 1.0 - sample_coordinate.y);"
: "sample_coordinate;",
/*$2=*/output_range_.has_value()
? absl::Substitute("pixel = pixel * float($0) + float($1);",
(output_range_->second - output_range_->first),
output_range_->first)
: "",
/*$3=*/single_channel ? "" : R"(fragColor.g = pixel.g;
fragColor.b = pixel.b;)",
/*$4=*/
include_alpha ? "fragColor.a = pixel.a;"
: (single_channel ? "" : "fragColor.a = 1.0;"));

const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
};
const GLchar* attr_name[NUM_ATTRIBUTES] = {
"position",
"texture_coordinate",
};
// shader program and params
mediapipe::GlhCreateProgram(
mediapipe::kBasicVertexShader, shader_source.c_str(), NUM_ATTRIBUTES,
&attr_name[0], attr_location, &to_tex2d_program_);
RET_CHECK(to_tex2d_program_) << "Problem initializing the program.";
glUseProgram(to_tex2d_program_);
glUniform1i(glGetUniformLocation(to_tex2d_program_, "frame"), 1);
glGenFramebuffers(1, &framebuffer_);
})glsl",
/*$0=*/single_channel ? "vec1" : "vec4",
/*$1=*/
flip_vertically_
? "vec2(sample_coordinate.x, 1.0 - sample_coordinate.y);"
: "sample_coordinate;",
/*$2=*/output_range_.has_value()
? absl::Substitute(
"pixel = pixel * float($0) + float($1);",
(output_range_->second - output_range_->first),
output_range_->first)
: "",
/*$3=*/single_channel ? "" : R"glsl(fragColor.g = pixel.g;
fragColor.b = pixel.b;)glsl",
/*$4=*/
include_alpha ? "fragColor.a = pixel.a;"
: (single_channel ? "" : "fragColor.a = 1.0;"));

const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
};
const GLchar* attr_name[NUM_ATTRIBUTES] = {
"position",
"texture_coordinate",
};
// shader program and params
mediapipe::GlhCreateProgram(
mediapipe::kBasicVertexShader, shader_source.c_str(),
NUM_ATTRIBUTES, &attr_name[0], attr_location, &to_tex2d_program_);
RET_CHECK(to_tex2d_program_) << "Problem initializing the program.";
glUseProgram(to_tex2d_program_);
glUniform1i(glGetUniformLocation(to_tex2d_program_, "frame"), 1);
glGenFramebuffers(1, &framebuffer_);

#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
return absl::OkStatus();
}));
return absl::OkStatus();
}));
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#endif // !MEDIAPIPE_DISABLE_GPU
return absl::OkStatus();
}

absl::Status TensorConverterCalculator::LoadOptions(CalculatorContext* cc) {
absl::Status TensorConverterCalculator::LoadOptions(CalculatorContext* cc,
bool use_gpu) {
// Get calculator options specified in the graph.
const auto& options =
cc->Options<::mediapipe::TensorConverterCalculatorOptions>();
Expand Down Expand Up @@ -635,7 +655,7 @@ absl::Status TensorConverterCalculator::LoadOptions(CalculatorContext* cc) {
}

// Get y-flip mode.
ASSIGN_OR_RETURN(flip_vertically_, ShouldFlipVertically(options));
ASSIGN_OR_RETURN(flip_vertically_, ShouldFlipVertically(options, use_gpu));

// Get row_major_matrix mode.
row_major_matrix_ = options.row_major_matrix();
Expand Down
6 changes: 4 additions & 2 deletions mediapipe/calculators/tensor/tensor_converter_calculator.proto
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,12 +44,14 @@ message TensorConverterCalculatorOptions {
// with a coordinate system where the origin is at the bottom-left corner
// (e.g., in OpenGL) whereas the ML model expects an image with a top-left
// origin.
// Prefer gpu_origin over this field.
// Prefer gpu_origin over this field when using GPU input images.
optional bool flip_vertically = 2 [default = false];

// Determines when the input image should be flipped vertically.
// Determines when the input GPU image should be flipped vertically.
// See GpuOrigin.Mode for more information.
// Affects only IMAGE_GPU inputs.
// If unset, falls back to flip_vertically for backwards compatibility.
// Cannot set both gpu_origin and flip_vertically.
optional GpuOrigin.Mode gpu_origin = 10;

// Controls how many channels of the input image get passed through to the
Expand Down
Loading

0 comments on commit 124a4de

Please sign in to comment.