Compositor: Implement Keying node for new CPU compositor

Reference #125968.

source/blender/nodes/composite/nodes/node_composite_keying.cc

@@ -7,6 +7,8 @@
  */
 
 #include "BLI_math_base.h"
+#include "BLI_math_color.h"
+#include "BLI_math_vector.h"
 #include "BLI_math_vector_types.hh"
 
 #include "DNA_movieclip_types.h"
@@ -87,17 +89,6 @@ class KeyingOperation : public NodeOperation {
 
   void execute() override
   {
-    /* Not yet supported on CPU. */
-    if (!context().use_gpu()) {
-      for (const bNodeSocket *output : this->node()->output_sockets()) {
-        Result &output_result = get_result(output->identifier);
-        if (output_result.should_compute()) {
-          output_result.allocate_invalid();
-        }
-      }
-      return;
-    }
-
     Result blurred_input = compute_blurred_input();
 
     Result matte = compute_matte(blurred_input);
@@ -158,6 +149,14 @@ class KeyingOperation : public NodeOperation {
   }
 
   Result extract_input_chroma()
+  {
+    if (this->context().use_gpu()) {
+      return this->extract_input_chroma_gpu();
+    }
+    return this->extract_input_chroma_cpu();
+  }
+
+  Result extract_input_chroma_gpu()
   {
     GPUShader *shader = context().get_shader("compositor_keying_extract_chroma");
     GPU_shader_bind(shader);
@@ -178,7 +177,41 @@ class KeyingOperation : public NodeOperation {
     return output;
   }
 
+  Result extract_input_chroma_cpu()
+  {
+    Result &input = get_input("Image");
+
+    Result output = context().create_result(ResultType::Color);
+    output.allocate_texture(input.domain());
+
+    parallel_for(input.domain().size, [&](const int2 texel) {
+      const float4 color = input.load_pixel(texel);
+      float4 color_ycca;
+      rgb_to_ycc(color.x,
+                 color.y,
+                 color.z,
+                 &color_ycca.x,
+                 &color_ycca.y,
+                 &color_ycca.z,
+                 BLI_YCC_ITU_BT709);
+      color_ycca /= 255.0f;
+      color_ycca.w = color.w;
+
+      output.store_pixel(texel, color_ycca);
+    });
+
+    return output;
+  }
+
   Result replace_input_chroma(Result &new_chroma)
+  {
+    if (this->context().use_gpu()) {
+      return this->replace_input_chroma_gpu(new_chroma);
+    }
+    return this->replace_input_chroma_cpu(new_chroma);
+  }
+
+  Result replace_input_chroma_gpu(Result &new_chroma)
   {
     GPUShader *shader = context().get_shader("compositor_keying_replace_chroma");
     GPU_shader_bind(shader);
@@ -202,7 +235,53 @@ class KeyingOperation : public NodeOperation {
     return output;
   }
 
+  Result replace_input_chroma_cpu(Result &new_chroma)
+  {
+    Result &input = get_input("Image");
+
+    Result output = context().create_result(ResultType::Color);
+    output.allocate_texture(input.domain());
+
+    parallel_for(input.domain().size, [&](const int2 texel) {
+      const float4 color = input.load_pixel(texel);
+      float4 color_ycca;
+      rgb_to_ycc(color.x,
+                 color.y,
+                 color.z,
+                 &color_ycca.x,
+                 &color_ycca.y,
+                 &color_ycca.z,
+                 BLI_YCC_ITU_BT709);
+
+      const float2 new_chroma_cb_cr = new_chroma.load_pixel(texel).yz();
+      color_ycca.y = new_chroma_cb_cr.x * 255.0f;
+      color_ycca.z = new_chroma_cb_cr.y * 255.0f;
+
+      float4 color_rgba;
+      ycc_to_rgb(color_ycca.x,
+                 color_ycca.y,
+                 color_ycca.z,
+                 &color_rgba.x,
+                 &color_rgba.y,
+                 &color_rgba.z,
+                 BLI_YCC_ITU_BT709);
+      color_rgba.w = color.w;
+
+      output.store_pixel(texel, color_rgba);
+    });
+
+    return output;
+  }
+
   Result compute_matte(Result &input)
+  {
+    if (this->context().use_gpu()) {
+      return this->compute_matte_gpu(input);
+    }
+    return this->compute_matte_cpu(input);
+  }
+
+  Result compute_matte_gpu(Result &input)
   {
     GPUShader *shader = context().get_shader("compositor_keying_compute_matte");
     GPU_shader_bind(shader);
@@ -228,6 +307,64 @@ class KeyingOperation : public NodeOperation {
     return output;
   }
 
+  Result compute_matte_cpu(Result &input)
+  {
+    const float key_balance = node_storage(bnode()).screen_balance;
+
+    Result &key = get_input("Key Color");
+
+    Result output = context().create_result(ResultType::Float);
+    output.allocate_texture(input.domain());
+
+    auto compute_saturation_indices = [](const float3 &v) {
+      int index_of_max = ((v.x > v.y) ? ((v.x > v.z) ? 0 : 2) : ((v.y > v.z) ? 1 : 2));
+      int2 other_indices = (int2(index_of_max) + int2(1, 2)) % 3;
+      int min_index = math::min(other_indices.x, other_indices.y);
+      int max_index = math::max(other_indices.x, other_indices.y);
+      return int3(index_of_max, max_index, min_index);
+    };
+
+    auto compute_saturation = [&](const float4 &color, const int3 &indices) {
+      float weighted_average = math::interpolate(color[indices.y], color[indices.z], key_balance);
+      return (color[indices.x] - weighted_average) * math::abs(1.0f - weighted_average);
+    };
+
+    parallel_for(input.domain().size, [&](const int2 texel) {
+      float4 input_color = input.load_pixel(texel);
+
+      /* We assume that the keying screen will not be overexposed in the image, so if the input
+       * brightness is high, we assume the pixel is opaque. */
+      if (math::reduce_min(input_color) > 1.0f) {
+        output.store_pixel(texel, float4(1.0f));
+        return;
+      }
+
+      float4 key_color = key.load_pixel(texel);
+      int3 key_saturation_indices = compute_saturation_indices(key_color.xyz());
+      float input_saturation = compute_saturation(input_color, key_saturation_indices);
+      float key_saturation = compute_saturation(key_color, key_saturation_indices);
+
+      float matte;
+      if (input_saturation < 0.0f) {
+        /* Means main channel of pixel is different from screen, assume this is completely a
+         * foreground. */
+        matte = 1.0f;
+      }
+      else if (input_saturation >= key_saturation) {
+        /* Matched main channels and higher saturation on pixel is treated as completely
+         * background. */
+        matte = 0.0f;
+      }
+      else {
+        matte = 1.0f - math::clamp(input_saturation / key_saturation, 0.0f, 1.0f);
+      }
+
+      output.store_pixel(texel, float4(matte));
+    });
+
+    return output;
+  }
+
   Result compute_tweaked_matte(Result &input_matte)
   {
     Result &output_edges = get_result("Edges");
@@ -250,16 +387,29 @@ class KeyingOperation : public NodeOperation {
       return output_matte;
     }
 
+    if (this->context().use_gpu()) {
+      return this->compute_tweaked_matte_gpu(input_matte);
+    }
+    return this->compute_tweaked_matte_cpu(input_matte);
+  }
+
+  Result compute_tweaked_matte_gpu(Result &input_matte)
+  {
     GPUShader *shader = context().get_shader("compositor_keying_tweak_matte");
     GPU_shader_bind(shader);
 
+    Result &output_edges = get_result("Edges");
+    const bool core_matte_exists = node().input_by_identifier("Core Matte")->is_logically_linked();
+    const bool garbage_matte_exists =
+        node().input_by_identifier("Garbage Matte")->is_logically_linked();
+
     GPU_shader_uniform_1b(shader, "compute_edges", output_edges.should_compute());
     GPU_shader_uniform_1b(shader, "apply_core_matte", core_matte_exists);
     GPU_shader_uniform_1b(shader, "apply_garbage_matte", garbage_matte_exists);
     GPU_shader_uniform_1i(shader, "edge_search_radius", node_storage(bnode()).edge_kernel_radius);
     GPU_shader_uniform_1f(shader, "edge_tolerance", node_storage(bnode()).edge_kernel_tolerance);
-    GPU_shader_uniform_1f(shader, "black_level", black_level);
-    GPU_shader_uniform_1f(shader, "white_level", white_level);
+    GPU_shader_uniform_1f(shader, "black_level", node_storage(bnode()).clip_black);
+    GPU_shader_uniform_1f(shader, "white_level", node_storage(bnode()).clip_white);
 
     input_matte.bind_as_texture(shader, "input_matte_tx");
 
@@ -288,6 +438,82 @@ class KeyingOperation : public NodeOperation {
     return output_matte;
   }
 
+  Result compute_tweaked_matte_cpu(Result &input_matte)
+  {
+    const bool apply_core_matte =
+        this->node().input_by_identifier("Core Matte")->is_logically_linked();
+    const bool apply_garbage_matte =
+        this->node().input_by_identifier("Garbage Matte")->is_logically_linked();
+
+    Result &output_edges = this->get_result("Edges");
+    const bool compute_edges = output_edges.should_compute();
+    const int edge_search_radius = node_storage(bnode()).edge_kernel_radius;
+    const float edge_tolerance = node_storage(bnode()).edge_kernel_tolerance;
+    const float black_level = node_storage(bnode()).clip_black;
+    const float white_level = node_storage(bnode()).clip_white;
+
+    Result &garbage_matte_image = get_input("Garbage Matte");
+    Result &core_matte_image = get_input("Core Matte");
+
+    Result output_matte = context().create_result(ResultType::Float);
+    output_matte.allocate_texture(input_matte.domain());
+
+    output_edges.allocate_texture(input_matte.domain());
+
+    parallel_for(input_matte.domain().size, [&](const int2 texel) {
+      float matte = input_matte.load_pixel(texel).x;
+
+      /* Search the neighborhood around the current matte value and identify if it lies along the
+       * edges of the matte. This is needs to be computed only when we need to compute the edges
+       * output or tweak the levels of the matte. */
+      bool is_edge = false;
+      if (compute_edges || black_level != 0.0f || white_level != 1.0f) {
+        /* Count the number of neighbors whose matte is sufficiently similar to the current matte,
+         * as controlled by the edge_tolerance factor. */
+        int count = 0;
+        for (int j = -edge_search_radius; j <= edge_search_radius; j++) {
+          for (int i = -edge_search_radius; i <= edge_search_radius; i++) {
+            float neighbor_matte = input_matte.load_pixel_extended(texel + int2(i, j)).x;
+            count += int(math::distance(matte, neighbor_matte) < edge_tolerance);
+          }
+        }
+
+        /* If the number of neighbors that are sufficiently similar to the center matte is less
+         * that 90% of the total number of neighbors, then that means the variance is high in that
+         * areas and it is considered an edge. */
+        is_edge = count < ((edge_search_radius * 2 + 1) * (edge_search_radius * 2 + 1)) * 0.9f;
+      }
+
+      float tweaked_matte = matte;
+
+      /* Remap the matte using the black and white levels, but only for areas that are not on the
+       * edge of the matte to preserve details. Also check for equality between levels to avoid
+       * zero division. */
+      if (!is_edge && white_level != black_level) {
+        tweaked_matte = math::clamp(
+            (matte - black_level) / (white_level - black_level), 0.0f, 1.0f);
+      }
+
+      /* Exclude unwanted areas using the provided garbage matte, 1 means unwanted, so invert the
+       * garbage matte and take the minimum. */
+      if (apply_garbage_matte) {
+        float garbage_matte = garbage_matte_image.load_pixel(texel).x;
+        tweaked_matte = math::min(tweaked_matte, 1.0f - garbage_matte);
+      }
+
+      /* Include wanted areas that were incorrectly keyed using the provided core matte. */
+      if (apply_core_matte) {
+        float core_matte = core_matte_image.load_pixel(texel).x;
+        tweaked_matte = math::max(tweaked_matte, core_matte);
+      }
+
+      output_matte.store_pixel(texel, float4(tweaked_matte));
+      output_edges.store_pixel(texel, float4(is_edge ? 1.0f : 0.0f));
+    });
+
+    return output_matte;
+  }
+
   Result compute_blurred_matte(Result &input_matte)
   {
     const float blur_size = node_storage(bnode()).blur_post;
@@ -345,6 +571,16 @@ class KeyingOperation : public NodeOperation {
   }
 
   void compute_image(Result &matte)
+  {
+    if (this->context().use_gpu()) {
+      this->compute_image_gpu(matte);
+    }
+    else {
+      this->compute_image_cpu(matte);
+    }
+  }
+
+  void compute_image_gpu(Result &matte)
   {
     GPUShader *shader = context().get_shader("compositor_keying_compute_image");
     GPU_shader_bind(shader);
@@ -372,6 +608,43 @@ class KeyingOperation : public NodeOperation {
     matte.unbind_as_texture();
     output.unbind_as_image();
   }
+
+  void compute_image_cpu(Result &matte_image)
+  {
+    const float despill_factor = node_storage(bnode()).despill_factor;
+    const float despill_balance = node_storage(bnode()).despill_balance;
+
+    Result &input = get_input("Image");
+    Result &key = get_input("Key Color");
+
+    Result &output = get_result("Image");
+    output.allocate_texture(matte_image.domain());
+
+    auto compute_saturation_indices = [](const float3 &v) {
+      int index_of_max = ((v.x > v.y) ? ((v.x > v.z) ? 0 : 2) : ((v.y > v.z) ? 1 : 2));
+      int2 other_indices = (int2(index_of_max) + int2(1, 2)) % 3;
+      int min_index = math::min(other_indices.x, other_indices.y);
+      int max_index = math::max(other_indices.x, other_indices.y);
+      return int3(index_of_max, max_index, min_index);
+    };
+
+    parallel_for(input.domain().size, [&](const int2 texel) {
+      float4 key_color = key.load_pixel(texel);
+      float4 color = input.load_pixel(texel);
+      float matte = matte_image.load_pixel(texel).x;
+
+      /* Alpha multiply the matte to the image. */
+      color *= matte;
+
+      /* Color despill. */
+      int3 indices = compute_saturation_indices(key_color.xyz());
+      float weighted_average = math::interpolate(
+          color[indices.y], color[indices.z], despill_balance);
+      color[indices.x] -= math::max(0.0f, (color[indices.x] - weighted_average) * despill_factor);
+
+      output.store_pixel(texel, color);
+    });
+  }
 };
 
 static NodeOperation *get_compositor_operation(Context &context, DNode node)