Upgrades and fixes for image segmentation category mask on GPU

PiperOrigin-RevId: 523204584

mediapipe/tasks/cc/vision/image_segmenter/calculators/segmentation_postprocessor_gl.cc

@@ -100,58 +100,89 @@ void main() {
   gl_FragColor = vec4(out_value, out_value, out_value, out_value);
 })";
 
-// Hard-coded for max of 3 textures for now, so num classes must be <= 12, and
+// For our argmax shader, we use a simple iterative approach to avoid the extra
-// the cost of this shader will be higher than necessary for smaller numbers of
+// hassle that accompanies usage of depth buffer for this, since we're not as
-// classes.
+// concerned with performance. Since we run the shader chunk-by-chunk, we can
-// TODO: Improve this.
+// simply hard-code our different max comparisons.
 static constexpr char kArgmaxShader[] = R"(
-DEFAULT_PRECISION(mediump, float)
+DEFAULT_PRECISION(highp, float)
 in vec2 sample_coordinate;
-uniform sampler2D input_texture0;
+uniform sampler2D prev_max_texture;  // prev_max_value, prev_max_arg, 0, 1
-uniform sampler2D input_texture1;
+uniform sampler2D current_chunk;
-uniform sampler2D input_texture2;
+uniform int num_channels;  // how many channels from current chunk to use (1-4)
+uniform int argmax_offset;  // index of first confidence mask in current chunk
 
-int argmax4(vec4 vec) {
+float max4(vec4 vec, out int argmax) {
   float aMax = max(vec.x, vec.y);
   float bMax = max(vec.z, vec.w);
   if (aMax >= bMax) {
-    if (vec.x >= vec.y) return 0;
+    if (vec.x >= vec.y) {
-    return 1;
+        argmax = 0;
-  } else if (vec.z >= vec.w) return 2;
+        return vec.x;
-  return 3;
+    }
+    argmax = 1;
+    return vec.y;
+  } else if (vec.z >= vec.w) {
+    argmax = 2;
+    return vec.z;
+  }
+  argmax = 3;
+  return vec.w;
 }
 
-float max4(vec4 vec) {
+float max3(vec4 vec, out int argmax) {
-  return max(max(vec.x, vec.y), max(vec.z, vec.w));
+    if (vec.x >= vec.y) {
+        if (vec.x >= vec.z) {
+            argmax = 0;
+            return vec.x;
+        }
+        argmax = 2;
+        return vec.z;
+    } else if (vec.y >= vec.z) {
+        argmax = 1;
+        return vec.y;
+    }
+    argmax = 2;
+    return vec.z;
+}
+
+float max2(vec4 vec, out int argmax) {
+    if (vec.x >= vec.y) {
+        argmax = 0;
+        return vec.x;
+    }
+    argmax = 1;
+    return vec.y;
 }
 
 void main() {
-  // Grab all vecs
+    vec2 prev_pixel = texture2D(prev_max_texture, sample_coordinate).xy;
-  vec4 pixel0 = texture2D(input_texture0, sample_coordinate);
+    float max_value = prev_pixel.x;
-  vec4 pixel1 = texture2D(input_texture1, sample_coordinate);
+    vec4 chunk_pixel = texture2D(current_chunk, sample_coordinate);
-  vec4 pixel2 = texture2D(input_texture2, sample_coordinate);
 
-  // Find vector which contains maximum value, and return its argmax
+    int chunk_argmax;
-  float max0 = max4(pixel0);
+    float chunk_max_value;
-  float max1 = max4(pixel1);
+    if (num_channels == 1) {
-  float max2 = max4(pixel2);
+      chunk_max_value = chunk_pixel.x;
-
+      chunk_argmax = 0;
-  int argmax;
+    } else if (num_channels == 2) {
-  float out_value;
+      chunk_max_value = max2(chunk_pixel, chunk_argmax);
-  if (max0 >= max1) {
+    } else if (num_channels == 3) {
-    if (max0 >= max2) {
+      chunk_max_value = max3(chunk_pixel, chunk_argmax);
-      argmax = argmax4(pixel0);
     } else {
-      argmax = argmax4(pixel2) + 8;
+      chunk_max_value = max4(chunk_pixel, chunk_argmax);
-    }
-  } else if (max1 >= max2) {
-    argmax = argmax4(pixel1) + 4;
-  } else {
-    argmax = argmax4(pixel2) + 8;
     }
 
-  out_value = float(argmax) / 255.0;
+    // Now compare against previous max_value
-  gl_FragColor = vec4(out_value, out_value, out_value, out_value);
+    if (chunk_max_value > max_value) {
+      // For now we convert our final integral argmax
+      // (chunk_argmax + argmax_offset) to a float from 0.0 to 1.0 in steps of
+      // 1/255.0.
+      float final_argmax = float(chunk_argmax + argmax_offset) / 255.0;
+      gl_FragColor = vec4(chunk_max_value, final_argmax, 0.0, 1.0);
+    } else {
+      gl_FragColor = vec4(max_value, prev_pixel.y, 0.0, 1.0);
+    }
 })";
 
 // Softmax is in 3 steps:
@@ -316,8 +347,7 @@ absl::Status SegmentationPostprocessorGl::GlInit() {
     }
 
     const std::string activation_shader_source =
-        absl::StrCat(std::string(mediapipe::kMediaPipeFragmentShaderPreamble),
+        absl::StrFormat(kActivationFragmentShader, activation_fn);
-                     absl::StrFormat(kActivationFragmentShader, activation_fn));
 
     const std::string split_fragment_shader_source =
         absl::StrCat(std::string(mediapipe::kMediaPipeFragmentShaderPreamble),
@@ -326,13 +356,14 @@ absl::Status SegmentationPostprocessorGl::GlInit() {
         absl::StrCat(std::string(mediapipe::kMediaPipeVertexShaderPreamble),
                      std::string(kSplitVertexShader));
 
-    const std::string channel_select_shader_source =
+    // Compile all our shader programs and grab uniforms.
-        absl::StrCat(std::string(mediapipe::kMediaPipeFragmentShaderPreamble),
+    // Simple shaders (Activation and Channel-select)
-                     std::string(kChannelSelectShader));
+    MP_RETURN_IF_ERROR(CreateBasicFragmentShaderProgram(
-
+        "activation", activation_shader_source, {"input_texture"},
-    const std::string argmax_shader_source =
+        &activation_shader_));
-        absl::StrCat(std::string(mediapipe::kMediaPipeFragmentShaderPreamble),
+    MP_RETURN_IF_ERROR(CreateBasicFragmentShaderProgram(
-                     std::string(kArgmaxShader));
+        "channel select", kChannelSelectShader,
+        {"input_texture", "channel_select"}, &channel_select_shader_));
 
     // Softmax shaders (Max, Transform+Sum, and Normalization)
     MP_RETURN_IF_ERROR(CreateBasicFragmentShaderProgram(
@@ -346,18 +377,14 @@ absl::Status SegmentationPostprocessorGl::GlInit() {
         "softmax normalization", kNormalizationShader,
         {"sum_texture", "current_chunk"}, &softmax_normalization_shader_));
 
-    // Compile all our shader programs.
+    // Category mask shaders (Argmax)
-    // Note: we enable `force_log_errors` so that we get full debugging error
+    MP_RETURN_IF_ERROR(CreateBasicFragmentShaderProgram(
-    //   messages when compiling shaders on web, where normally such errors are
+        "argmax", kArgmaxShader,
-    //   suppressed. See //mediapipe/gpu/shader_util.cc for more
+        {"prev_max_texture", "current_chunk", "num_channels", "argmax_offset"},
-    //   info.
+        &argmax_shader_));
-    mediapipe::GlhCreateProgram(
-        kBasicVertexShader, activation_shader_source.c_str(), NUM_ATTRIBUTES,
-        &attr_name[0], attr_location, &activation_program_,
-        /* force_log_errors */ true);
-    RET_CHECK(activation_program_)
-        << "Problem initializing the activation program.";
 
+    // Split shader. This is created separately since it uses a custom vertex
+    // shader. TODO: Refactor so this shares common init code as well.
     mediapipe::GlhCreateProgram(split_vertex_shader_source.c_str(),
                                 split_fragment_shader_source.c_str(),
                                 NUM_ATTRIBUTES, &attr_name[0], attr_location,
@@ -365,25 +392,7 @@ absl::Status SegmentationPostprocessorGl::GlInit() {
                                 /* force_log_errors */ true);
     RET_CHECK(split_program_) << "Problem initializing the split program.";
 
-    mediapipe::GlhCreateProgram(
+    // Get split program uniform locations.
-        kBasicVertexShader, channel_select_shader_source.c_str(),
-        NUM_ATTRIBUTES, &attr_name[0], attr_location, &channel_select_program_,
-        /* force_log_errors */ true);
-    RET_CHECK(channel_select_program_)
-        << "Problem initializing the channel select program.";
-
-    mediapipe::GlhCreateProgram(kBasicVertexShader,
-                                argmax_shader_source.c_str(), NUM_ATTRIBUTES,
-                                &attr_name[0], attr_location, &argmax_program_,
-                                /* force_log_errors */ true);
-    RET_CHECK(argmax_program_) << "Problem initializing the argmax program.";
-
-    // Get uniform locations.
-    activation_texture_uniform_ =
-        glGetUniformLocation(activation_program_, "input_texture");
-    RET_CHECK(activation_texture_uniform_ > 0)
-        << "activation input_texture uniform not found.";
-
     split_texture_uniform_ =
         glGetUniformLocation(split_program_, "input_texture");
     RET_CHECK(split_texture_uniform_ > 0)
@@ -392,28 +401,6 @@ absl::Status SegmentationPostprocessorGl::GlInit() {
     RET_CHECK(split_x_offset_uniform_ > 0)
         << "split x_offset uniform not found.";
 
-    channel_select_texture_uniform_ =
-        glGetUniformLocation(channel_select_program_, "input_texture");
-    RET_CHECK(channel_select_texture_uniform_ > 0)
-        << "channel select input_texture uniform not found.";
-    channel_select_index_uniform_ =
-        glGetUniformLocation(channel_select_program_, "channel_select");
-    RET_CHECK(channel_select_index_uniform_ > 0)
-        << "channel select indexing uniform not found.";
-
-    argmax_texture0_uniform_ =
-        glGetUniformLocation(argmax_program_, "input_texture0");
-    RET_CHECK(argmax_texture0_uniform_ > 0)
-        << "argmax input_texture0 uniform not found.";
-    argmax_texture1_uniform_ =
-        glGetUniformLocation(argmax_program_, "input_texture1");
-    RET_CHECK(argmax_texture1_uniform_ > 0)
-        << "argmax input_texture1 uniform not found.";
-    argmax_texture2_uniform_ =
-        glGetUniformLocation(argmax_program_, "input_texture2");
-    RET_CHECK(argmax_texture2_uniform_ > 0)
-        << "argmax input_texture2 uniform not found.";
-
     // TODO: If ES3.0+ only, switch to VAO for handling attributes.
     glGenBuffers(1, &square_vertices_);
     glBindBuffer(GL_ARRAY_BUFFER, square_vertices_);
@@ -426,7 +413,6 @@ absl::Status SegmentationPostprocessorGl::GlInit() {
                  kBasicTextureVertices, GL_STATIC_DRAW);
 
     glBindBuffer(GL_ARRAY_BUFFER, 0);
-
     return absl::OkStatus();
   });
 }
@@ -492,8 +478,8 @@ SegmentationPostprocessorGl::GetSegmentationResultGpu(const Shape& input_shape,
     glEnableVertexAttribArray(ATTRIB_TEXTURE_POSITION);
 
     // Step 1: apply activation pass
-    glUseProgram(activation_program_);
+    glUseProgram(activation_shader_.program);
-    glUniform1i(activation_texture_uniform_, 1);
+    glUniform1i(activation_shader_.uniforms["input_texture"], 1);
     GlTexture activated_texture = helper_.CreateDestinationTexture(
         input_width, input_height, activation_output_format);
     helper_.BindFramebuffer(activated_texture);
@@ -660,48 +646,69 @@ SegmentationPostprocessorGl::GetSegmentationResultGpu(const Shape& input_shape,
 
     std::vector<GlTexture> outputs;
     if (is_category_mask) {
-      // Step 3: For CATEGORY, apply argmax shader with up to 3 textures to
+      // Step 3: For CATEGORY, apply argmax shader iteratively with each chunk
-      // extract final index mask.
+      // to get a 2-channel texture representing "combined maxval" and "argmax",
-      RET_CHECK(num_chunks <= 3)
+      // and then slice off the second channel for the category mask output,
-          << "Cannot handle more than 12 classes in argmax shader.";
+      // using our usual channel_select program.
+      glUseProgram(argmax_shader_.program);
+      glUniform1i(argmax_shader_.uniforms["current_chunk"], 1);
+      glUniform1i(argmax_shader_.uniforms["prev_max_texture"], 2);
 
-      glUseProgram(argmax_program_);
+      GlTexture max_texture = helper_.CreateDestinationTexture(
-      glUniform1i(argmax_texture0_uniform_, 1);
+          output_width, output_height, chunk_output_format);
-      glUniform1i(argmax_texture1_uniform_, 2);
+      GlTexture next_max_texture = helper_.CreateDestinationTexture(
-      glUniform1i(argmax_texture2_uniform_, 3);
+          output_width, output_height, chunk_output_format);
-      outputs.push_back(helper_.CreateDestinationTexture(
-          output_width, output_height, final_output_format));
-      helper_.BindFramebuffer(outputs.back());
 
-      // Bind however many chunks we have
+      // GLSL uses IEEE 754 single-precision floating-point for encoding its
+      // floats (at least for number representation, although not necessarily
+      // for operations). So we can clear to a reasonable minimum float value
+      // accordingly.
+      const float kFloatMin32 = -3.402823466e+38;
+      glClearColor(kFloatMin32, -1.0, 0.0, 1.0);
+      helper_.BindFramebuffer(max_texture);
+      glClear(GL_COLOR_BUFFER_BIT);
+      // Set our clear color back to a "normal" default.
+      glClearColor(0.0, 0.0, 0.0, 0.0);
       for (int i = 0; i < num_chunks; ++i) {
-        glActiveTexture(GL_TEXTURE1 + i);
+        int num_channels = 4;
+        if ((i + 1) * 4 > num_outputs) num_channels = num_outputs % 4;
+        glUniform1i(argmax_shader_.uniforms["num_channels"], num_channels);
+        glUniform1i(argmax_shader_.uniforms["argmax_offset"], i * 4);
+        helper_.BindFramebuffer(next_max_texture);
+        glActiveTexture(GL_TEXTURE2);
+        glBindTexture(GL_TEXTURE_2D, max_texture.name());
+        glActiveTexture(GL_TEXTURE1);
         glBindTexture(GL_TEXTURE_2D, chunks[i].name());
-      }
+        // TODO: We probably don't actually need all these clears.
-
-      for (int i = num_chunks; i < 3; ++i) {  // 3 is hard-coded max chunks
-        glActiveTexture(GL_TEXTURE1 + i);
-        // If texture is unbound, sampling from it should always give zeros.
-        // This is not ideal, but is ok for now for not polluting the argmax
-        // shader results too much.
-        glBindTexture(GL_TEXTURE_2D, 0);
-      }
-
         glClear(GL_COLOR_BUFFER_BIT);
         glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
 
-      // Unbind the extra textures here.
+        // Put results into max_texture, so we can repeat the process easily.
-      for (int i = 0; i < num_chunks; ++i) {
+        std::swap(max_texture, next_max_texture);
-        glActiveTexture(GL_TEXTURE1 + i);
-        glBindTexture(GL_TEXTURE_2D, 0);
       }
+
+      // Do final channel-select on max_texture below, selecting for argmax
+      outputs.push_back(helper_.CreateDestinationTexture(
+          output_width, output_height, final_output_format));
+      helper_.BindFramebuffer(outputs.back());
+      glUseProgram(channel_select_shader_.program);
+      glUniform1i(channel_select_shader_.uniforms["input_texture"], 1);
+      // 0:max_val, 1:argmax
+      glUniform1i(channel_select_shader_.uniforms["channel_select"], 1);
+      glBindTexture(GL_TEXTURE_2D, max_texture.name());
+      // We can't interpolate across argmax values, so we disable linear
+      // interpolation there for this upsampling step.
+      glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
+      glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
+      glClear(GL_COLOR_BUFFER_BIT);
+      glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
     } else {
       // Step 3: For CONFIDENCE, apply channel-select repeatedly to extract
       // final textures.
-      glUseProgram(channel_select_program_);
+      glUseProgram(channel_select_shader_.program);
-      glUniform1i(channel_select_texture_uniform_, 1);
+      glUniform1i(channel_select_shader_.uniforms["input_texture"], 1);
       for (int i = 0; i < num_outputs; i++) {
-        glUniform1i(channel_select_index_uniform_, (i % 4));
+        glUniform1i(channel_select_shader_.uniforms["channel_select"], (i % 4));
         outputs.push_back(helper_.CreateDestinationTexture(
             output_width, output_height, final_output_format));
         helper_.BindFramebuffer(outputs.back());
@@ -744,19 +751,16 @@ SegmentationPostprocessorGl::GetSegmentationResultGpu(const Shape& input_shape,
 // Cleanup OpenGL resources on destruction
 SegmentationPostprocessorGl::~SegmentationPostprocessorGl() {
   helper_.RunInGlContext([this] {
-    glDeleteProgram(activation_program_);
-    glDeleteProgram(argmax_program_);
-    glDeleteProgram(channel_select_program_);
     glDeleteProgram(split_program_);
     glDeleteBuffers(1, &square_vertices_);
     glDeleteBuffers(1, &texture_vertices_);
-    activation_program_ = 0;
-    argmax_program_ = 0;
-    channel_select_program_ = 0;
     split_program_ = 0;
     square_vertices_ = 0;
     texture_vertices_ = 0;
 
+    glDeleteProgram(activation_shader_.program);
+    glDeleteProgram(argmax_shader_.program);
+    glDeleteProgram(channel_select_shader_.program);
     glDeleteProgram(softmax_max_shader_.program);
     glDeleteProgram(softmax_transform_and_sum_shader_.program);
     glDeleteProgram(softmax_normalization_shader_.program);

mediapipe/tasks/cc/vision/image_segmenter/calculators/segmentation_postprocessor_gl.h

@@ -54,21 +54,16 @@ class SegmentationPostprocessorGl {
   GlCalculatorHelper helper_;
 
   // GL references (programs, buffers, uniforms)
-  GLuint activation_program_ = 0;
+  // Split program is special because it uses a custom vertex shader.
-  GLuint argmax_program_ = 0;
-  GLuint channel_select_program_ = 0;
   GLuint split_program_ = 0;
   GLuint square_vertices_ = 0;
   GLuint texture_vertices_ = 0;
-  GLint activation_texture_uniform_;
-  GLint argmax_texture0_uniform_;
-  GLint argmax_texture1_uniform_;
-  GLint argmax_texture2_uniform_;
-  GLint channel_select_texture_uniform_;
-  GLint channel_select_index_uniform_;
   GLint split_texture_uniform_;
   GLint split_x_offset_uniform_;
 
+  GlShader activation_shader_;
+  GlShader argmax_shader_;
+  GlShader channel_select_shader_;
   GlShader softmax_max_shader_;
   GlShader softmax_transform_and_sum_shader_;
   GlShader softmax_normalization_shader_;