Realtime Compositor: Implement despeckle node

This patch implements the despeckle node for the realtime compositor.

Differential Revision: https://developer.blender.org/D15673

Reviewed By: Clement Foucault

4 files changed, 127 insertions, 1 deletions

diff --git a/source/blender/gpu/CMakeLists.txt b/source/blender/gpu/CMakeLists.txt
index fd732a09e08..40ede29eab6 100644
--- a/source/blender/gpu/CMakeLists.txt
+++ b/source/blender/gpu/CMakeLists.txt
@@ -318,6 +318,7 @@ set(GLSL_SRC
   shaders/compositor/compositor_bokeh_image.glsl
   shaders/compositor/compositor_box_mask.glsl
   shaders/compositor/compositor_convert.glsl
+  shaders/compositor/compositor_despeckle.glsl
   shaders/compositor/compositor_directional_blur.glsl
   shaders/compositor/compositor_edge_filter.glsl
   shaders/compositor/compositor_ellipse_mask.glsl
@@ -568,6 +569,7 @@ set(SRC_SHADER_CREATE_INFOS
   shaders/compositor/infos/compositor_bokeh_image_info.hh
   shaders/compositor/infos/compositor_box_mask_info.hh
   shaders/compositor/infos/compositor_convert_info.hh
+  shaders/compositor/infos/compositor_despeckle_info.hh
   shaders/compositor/infos/compositor_directional_blur_info.hh
   shaders/compositor/infos/compositor_edge_filter_info.hh
   shaders/compositor/infos/compositor_ellipse_mask_info.hh
diff --git a/source/blender/gpu/shaders/compositor/compositor_despeckle.glsl b/source/blender/gpu/shaders/compositor/compositor_despeckle.glsl
new file mode 100644
index 00000000000..e4743d69d17
--- /dev/null
+++ b/source/blender/gpu/shaders/compositor/compositor_despeckle.glsl
@@ -0,0 +1,70 @@
+#pragma BLENDER_REQUIRE(gpu_shader_compositor_texture_utilities.glsl)
+
+/* Returns true if the given color is close enough to the given reference color within the
+ * threshold supplied by the user, and returns false otherwise. */
+bool is_close(vec4 reference_color, vec4 color)
+{
+  return all(lessThan(abs(reference_color - color).rgb, vec3(threshold)));
+}
+
+void main()
+{
+  ivec2 texel = ivec2(gl_GlobalInvocationID.xy);
+
+  /* A 3x3 weights kernel whose weights are the inverse of the distance to the center of the
+   * kernel. So the center weight is zero, the corners weights are (1 / sqrt(2)), and the rest
+   * of the weights are 1. The total sum of weights is 4 plus quadruple the corner weight. */
+  float corner_weight = 1.0 / sqrt(2.0);
+  float sum_of_weights = 4.0 + corner_weight * 4.0;
+  mat3 weights = mat3(vec3(corner_weight, 1.0, corner_weight),
+                      vec3(1.0, 0.0, 1.0),
+                      vec3(corner_weight, 1.0, corner_weight));
+
+  vec4 center_color = texture_load(input_tx, texel);
+
+  /* Go over the pixels in the 3x3 window around the center pixel and compute the total sum of
+   * their colors multiplied by their weights. Additionally, for pixels whose colors are not close
+   * enough to the color of the center pixel, accumulate their color as well as their weights. */
+  vec4 sum_of_colors = vec4(0);
+  float accumulated_weight = 0.0;
+  vec4 accumulated_color = vec4(0);
+  for (int j = 0; j < 3; j++) {
+    for (int i = 0; i < 3; i++) {
+      float weight = weights[j][i];
+      vec4 color = texture_load(input_tx, texel + ivec2(i - 1, j - 1)) * weight;
+      sum_of_colors += color;
+      if (!is_close(center_color, color)) {
+        accumulated_color += color;
+        accumulated_weight += weight;
+      }
+    }
+  }
+
+  /* If the accumulated weight is zero, that means all pixels in the 3x3 window are similar and no
+   * need to despeckle anything, so write the original center color and return. */
+  if (accumulated_weight == 0.0) {
+    imageStore(output_img, texel, center_color);
+    return;
+  }
+
+  /* If the ratio between the accumulated weights and the total sum of weights is not larger than
+   * the user specified neighbor threshold, then the number of pixels in the neighborhood that are
+   * not close enough to the center pixel is low, and no need to despeckle anything, so write the
+   * original center color and return. */
+  if (accumulated_weight / sum_of_weights < neighbor_threshold) {
+    imageStore(output_img, texel, center_color);
+    return;
+  }
+
+  /* If the weighted average color of the neighborhood is close enough to the center pixel, then no
+   * need to despeckle anything, so write the original center color and return. */
+  if (is_close(center_color, sum_of_colors / sum_of_weights)) {
+    imageStore(output_img, texel, center_color);
+    return;
+  }
+
+  /* We need to despeckle, so write the mean accumulated color. */
+  float factor = texture_load(factor_tx, texel).x;
+  vec4 mean_color = accumulated_color / accumulated_weight;
+  imageStore(output_img, texel, mix(center_color, mean_color, factor));
+}
diff --git a/source/blender/gpu/shaders/compositor/infos/compositor_despeckle_info.hh b/source/blender/gpu/shaders/compositor/infos/compositor_despeckle_info.hh
new file mode 100644
index 00000000000..df86c3a8258
--- /dev/null
+++ b/source/blender/gpu/shaders/compositor/infos/compositor_despeckle_info.hh
@@ -0,0 +1,13 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include "gpu_shader_create_info.hh"
+
+GPU_SHADER_CREATE_INFO(compositor_despeckle)
+    .local_group_size(16, 16)
+    .push_constant(Type::FLOAT, "threshold")
+    .push_constant(Type::FLOAT, "neighbor_threshold")
+    .sampler(0, ImageType::FLOAT_2D, "input_tx")
+    .sampler(1, ImageType::FLOAT_2D, "factor_tx")
+    .image(0, GPU_RGBA16F, Qualifier::WRITE, ImageType::FLOAT_2D, "output_img")
+    .compute_source("compositor_despeckle.glsl")
+    .do_static_compilation(true);
diff --git a/source/blender/nodes/composite/nodes/node_composite_despeckle.cc b/source/blender/nodes/composite/nodes/node_composite_despeckle.cc
index 0b9f9c8f76d..51a98812f0d 100644
--- a/source/blender/nodes/composite/nodes/node_composite_despeckle.cc
+++ b/source/blender/nodes/composite/nodes/node_composite_despeckle.cc
@@ -8,7 +8,10 @@
 #include "UI_interface.h"
 #include "UI_resources.h"
 
+#include "GPU_shader.h"
+
 #include "COM_node_operation.hh"
+#include "COM_utilities.hh"
 
 #include "node_composite_util.hh"
 
@@ -46,7 +49,45 @@ class DespeckleOperation : public NodeOperation {
 
   void execute() override
   {
-    get_input("Image").pass_through(get_result("Image"));
+    const Result &input_image = get_input("Image");
+    /* Single value inputs can't be despeckled and are returned as is. */
+    if (input_image.is_single_value()) {
+      get_input("Image").pass_through(get_result("Image"));
+      return;
+    }
+
+    GPUShader *shader = shader_manager().get("compositor_despeckle");
+    GPU_shader_bind(shader);
+
+    GPU_shader_uniform_1f(shader, "threshold", get_threshold());
+    GPU_shader_uniform_1f(shader, "neighbor_threshold", get_neighbor_threshold());
+
+    input_image.bind_as_texture(shader, "input_tx");
+
+    const Result &factor_image = get_input("Fac");
+    factor_image.bind_as_texture(shader, "factor_tx");
+
+    const Domain domain = compute_domain();
+    Result &output_image = get_result("Image");
+    output_image.allocate_texture(domain);
+    output_image.bind_as_image(shader, "output_img");
+
+    compute_dispatch_threads_at_least(shader, domain.size);
+
+    GPU_shader_unbind();
+    output_image.unbind_as_image();
+    input_image.unbind_as_texture();
+    factor_image.unbind_as_texture();
+  }
+
+  float get_threshold()
+  {
+    return bnode().custom3;
+  }
+
+  float get_neighbor_threshold()
+  {
+    return bnode().custom4;
   }
 };