Compositor: Full frame Keying node

8 files changed, 251 insertions, 8 deletions

diff --git a/source/blender/compositor/operations/COM_KeyingBlurOperation.cc b/source/blender/compositor/operations/COM_KeyingBlurOperation.cc
index 994b00cd3f4..d5ebd5e9df7 100644
--- a/source/blender/compositor/operations/COM_KeyingBlurOperation.cc
+++ b/source/blender/compositor/operations/COM_KeyingBlurOperation.cc
@@ -96,4 +96,67 @@ bool KeyingBlurOperation::determineDependingAreaOfInterest(rcti *input,
   return NodeOperation::determineDependingAreaOfInterest(&newInput, readOperation, output);
 }
 
+void KeyingBlurOperation::get_area_of_interest(const int UNUSED(input_idx),
+                                               const rcti &output_area,
+                                               rcti &r_input_area)
+{
+  switch (m_axis) {
+    case BLUR_AXIS_X:
+      r_input_area.xmin = output_area.xmin - m_size;
+      r_input_area.ymin = output_area.ymin;
+      r_input_area.xmax = output_area.xmax + m_size;
+      r_input_area.ymax = output_area.ymax;
+      break;
+    case BLUR_AXIS_Y:
+      r_input_area.xmin = output_area.xmin;
+      r_input_area.ymin = output_area.ymin - m_size;
+      r_input_area.xmax = output_area.xmax;
+      r_input_area.ymax = output_area.ymax + m_size;
+      break;
+    default:
+      BLI_assert_msg(0, "Unknown axis");
+      break;
+  }
+}
+
+void KeyingBlurOperation::update_memory_buffer_partial(MemoryBuffer *output,
+                                                       const rcti &area,
+                                                       Span<MemoryBuffer *> inputs)
+{
+  const MemoryBuffer *input = inputs[0];
+  BuffersIterator<float> it = output->iterate_with(inputs, area);
+
+  int coord_max;
+  int elem_stride;
+  std::function<int()> get_current_coord;
+  switch (m_axis) {
+    case BLUR_AXIS_X:
+      get_current_coord = [&] { return it.x; };
+      coord_max = this->getWidth();
+      elem_stride = input->elem_stride;
+      break;
+    case BLUR_AXIS_Y:
+      get_current_coord = [&] { return it.y; };
+      coord_max = this->getHeight();
+      elem_stride = input->row_stride;
+      break;
+  }
+
+  for (; !it.is_end(); ++it) {
+    const int coord = get_current_coord();
+    const int start_coord = MAX2(0, coord - m_size + 1);
+    const int end_coord = MIN2(coord_max, coord + m_size);
+    const int count = end_coord - start_coord;
+
+    float sum = 0.0f;
+    const float *start = it.in(0) + (start_coord - coord) * elem_stride;
+    const float *end = start + count * elem_stride;
+    for (const float *elem = start; elem < end; elem += elem_stride) {
+      sum += *elem;
+    }
+
+    *it.out = sum / count;
+  }
+}
+
 }  // namespace blender::compositor
diff --git a/source/blender/compositor/operations/COM_KeyingBlurOperation.h b/source/blender/compositor/operations/COM_KeyingBlurOperation.h
index b055d7713f1..b290b905e63 100644
--- a/source/blender/compositor/operations/COM_KeyingBlurOperation.h
+++ b/source/blender/compositor/operations/COM_KeyingBlurOperation.h
@@ -18,14 +18,14 @@
 
 #pragma once
 
-#include "COM_NodeOperation.h"
+#include "COM_MultiThreadedOperation.h"
 
 namespace blender::compositor {
 
 /**
  * Class with implementation of blurring for keying node
  */
-class KeyingBlurOperation : public NodeOperation {
+class KeyingBlurOperation : public MultiThreadedOperation {
  protected:
   int m_size;
   int m_axis;
@@ -54,6 +54,13 @@ class KeyingBlurOperation : public NodeOperation {
   bool determineDependingAreaOfInterest(rcti *input,
                                         ReadBufferOperation *readOperation,
                                         rcti *output) override;
+
+  void get_area_of_interest(const int input_idx,
+                            const rcti &output_area,
+                            rcti &r_input_area) override;
+  void update_memory_buffer_partial(MemoryBuffer *output,
+                                    const rcti &area,
+                                    Span<MemoryBuffer *> inputs) override;
 };
 
 }  // namespace blender::compositor
diff --git a/source/blender/compositor/operations/COM_KeyingClipOperation.cc b/source/blender/compositor/operations/COM_KeyingClipOperation.cc
index 4029be4e077..817c920ed91 100644
--- a/source/blender/compositor/operations/COM_KeyingClipOperation.cc
+++ b/source/blender/compositor/operations/COM_KeyingClipOperation.cc
@@ -130,4 +130,89 @@ bool KeyingClipOperation::determineDependingAreaOfInterest(rcti *input,
   return NodeOperation::determineDependingAreaOfInterest(&newInput, readOperation, output);
 }
 
+void KeyingClipOperation::get_area_of_interest(const int input_idx,
+                                               const rcti &output_area,
+                                               rcti &r_input_area)
+{
+  BLI_assert(input_idx == 0);
+  UNUSED_VARS_NDEBUG(input_idx);
+  r_input_area.xmin = output_area.xmin - m_kernelRadius;
+  r_input_area.xmax = output_area.xmax + m_kernelRadius;
+  r_input_area.ymin = output_area.ymin - m_kernelRadius;
+  r_input_area.ymax = output_area.ymax + m_kernelRadius;
+}
+
+void KeyingClipOperation::update_memory_buffer_partial(MemoryBuffer *output,
+                                                       const rcti &area,
+                                                       Span<MemoryBuffer *> inputs)
+{
+  const MemoryBuffer *input = inputs[0];
+  BuffersIterator<float> it = output->iterate_with(inputs, area);
+
+  const int delta = m_kernelRadius;
+  const float tolerance = m_kernelTolerance;
+  const int width = this->getWidth();
+  const int height = this->getHeight();
+  const int row_stride = input->row_stride;
+  const int elem_stride = input->elem_stride;
+  for (; !it.is_end(); ++it) {
+    const int x = it.x;
+    const int y = it.y;
+
+    const int start_x = MAX2(0, x - delta + 1);
+    const int start_y = MAX2(0, y - delta + 1);
+    const int end_x = MIN2(x + delta, width);
+    const int end_y = MIN2(y + delta, height);
+    const int x_len = end_x - start_x;
+    const int y_len = end_y - start_y;
+
+    const int total_count = x_len * y_len - 1;
+    const int threshold_count = ceil((float)total_count * 0.9f);
+    bool ok = false;
+    if (delta == 0) {
+      ok = true;
+    }
+
+    const float *main_elem = it.in(0);
+    const float value = *main_elem;
+    const float *row = input->get_elem(start_x, start_y);
+    const float *end_row = row + y_len * row_stride;
+    int count = 0;
+    for (; ok == false && row < end_row; row += row_stride) {
+      const float *end_elem = row + x_len * elem_stride;
+      for (const float *elem = row; ok == false && elem < end_elem; elem += elem_stride) {
+        if (UNLIKELY(elem == main_elem)) {
+          continue;
+        }
+
+        const float current_value = *elem;
+        if (fabsf(current_value - value) < tolerance) {
+          count++;
+          if (count >= threshold_count) {
+            ok = true;
+          }
+        }
+      }
+    }
+
+    if (m_isEdgeMatte) {
+      *it.out = ok ? 0.0f : 1.0f;
+    }
+    else {
+      if (!ok) {
+        *it.out = value;
+      }
+      else if (value < m_clipBlack) {
+        *it.out = 0.0f;
+      }
+      else if (value >= m_clipWhite) {
+        *it.out = 1.0f;
+      }
+      else {
+        *it.out = (value - m_clipBlack) / (m_clipWhite - m_clipBlack);
+      }
+    }
+  }
+}
+
 }  // namespace blender::compositor
diff --git a/source/blender/compositor/operations/COM_KeyingClipOperation.h b/source/blender/compositor/operations/COM_KeyingClipOperation.h
index 0a21fb48c99..1a17d591781 100644
--- a/source/blender/compositor/operations/COM_KeyingClipOperation.h
+++ b/source/blender/compositor/operations/COM_KeyingClipOperation.h
@@ -18,14 +18,14 @@
 
 #pragma once
 
-#include "COM_NodeOperation.h"
+#include "COM_MultiThreadedOperation.h"
 
 namespace blender::compositor {
 
 /**
  * Class with implementation of black/white clipping for keying node
  */
-class KeyingClipOperation : public NodeOperation {
+class KeyingClipOperation : public MultiThreadedOperation {
  protected:
   float m_clipBlack;
   float m_clipWhite;
@@ -68,6 +68,13 @@ class KeyingClipOperation : public NodeOperation {
   bool determineDependingAreaOfInterest(rcti *input,
                                         ReadBufferOperation *readOperation,
                                         rcti *output) override;
+
+  void get_area_of_interest(const int input_idx,
+                            const rcti &output_area,
+                            rcti &r_input_area) override;
+  void update_memory_buffer_partial(MemoryBuffer *output,
+                                    const rcti &area,
+                                    Span<MemoryBuffer *> inputs) override;
 };
 
 }  // namespace blender::compositor
diff --git a/source/blender/compositor/operations/COM_KeyingDespillOperation.cc b/source/blender/compositor/operations/COM_KeyingDespillOperation.cc
index d31a88cb91e..a03859632c9 100644
--- a/source/blender/compositor/operations/COM_KeyingDespillOperation.cc
+++ b/source/blender/compositor/operations/COM_KeyingDespillOperation.cc
@@ -82,4 +82,32 @@ void KeyingDespillOperation::executePixelSampled(float output[4],
   }
 }
 
+void KeyingDespillOperation::update_memory_buffer_partial(MemoryBuffer *output,
+                                                          const rcti &area,
+                                                          Span<MemoryBuffer *> inputs)
+{
+  for (BuffersIterator<float> it = output->iterate_with(inputs, area); !it.is_end(); ++it) {
+    const float *pixel_color = it.in(0);
+    const float *screen_color = it.in(1);
+
+    const int screen_primary_channel = max_axis_v3(screen_color);
+    const int other_1 = (screen_primary_channel + 1) % 3;
+    const int other_2 = (screen_primary_channel + 2) % 3;
+
+    const int min_channel = MIN2(other_1, other_2);
+    const int max_channel = MAX2(other_1, other_2);
+
+    const float average_value = m_colorBalance * pixel_color[min_channel] +
+                                (1.0f - m_colorBalance) * pixel_color[max_channel];
+    const float amount = (pixel_color[screen_primary_channel] - average_value);
+
+    copy_v4_v4(it.out, pixel_color);
+
+    const float amount_despill = m_despillFactor * amount;
+    if (amount_despill > 0.0f) {
+      it.out[screen_primary_channel] = pixel_color[screen_primary_channel] - amount_despill;
+    }
+  }
+}
+
 }  // namespace blender::compositor
diff --git a/source/blender/compositor/operations/COM_KeyingDespillOperation.h b/source/blender/compositor/operations/COM_KeyingDespillOperation.h
index 279ac60e6e9..16bed651d3a 100644
--- a/source/blender/compositor/operations/COM_KeyingDespillOperation.h
+++ b/source/blender/compositor/operations/COM_KeyingDespillOperation.h
@@ -18,14 +18,14 @@
 
 #pragma once
 
-#include "COM_NodeOperation.h"
+#include "COM_MultiThreadedOperation.h"
 
 namespace blender::compositor {
 
 /**
  * Class with implementation of keying despill node
  */
-class KeyingDespillOperation : public NodeOperation {
+class KeyingDespillOperation : public MultiThreadedOperation {
  protected:
   SocketReader *m_pixelReader;
   SocketReader *m_screenReader;
@@ -48,6 +48,10 @@ class KeyingDespillOperation : public NodeOperation {
   }
 
   void executePixelSampled(float output[4], float x, float y, PixelSampler sampler) override;
+
+  void update_memory_buffer_partial(MemoryBuffer *output,
+                                    const rcti &area,
+                                    Span<MemoryBuffer *> inputs) override;
 };
 
 }  // namespace blender::compositor
diff --git a/source/blender/compositor/operations/COM_KeyingOperation.cc b/source/blender/compositor/operations/COM_KeyingOperation.cc
index e786e4b8219..3edb5a5d34e 100644
--- a/source/blender/compositor/operations/COM_KeyingOperation.cc
+++ b/source/blender/compositor/operations/COM_KeyingOperation.cc
@@ -110,4 +110,49 @@ void KeyingOperation::executePixelSampled(float output[4], float x, float y, Pix
   }
 }
 
+void KeyingOperation::update_memory_buffer_partial(MemoryBuffer *output,
+                                                   const rcti &area,
+                                                   Span<MemoryBuffer *> inputs)
+{
+  for (BuffersIterator<float> it = output->iterate_with(inputs, area); !it.is_end(); ++it) {
+    const float *pixel_color = it.in(0);
+    const float *screen_color = it.in(1);
+
+    const int primary_channel = max_axis_v3(screen_color);
+    const float min_pixel_color = min_fff(pixel_color[0], pixel_color[1], pixel_color[2]);
+
+    if (min_pixel_color > 1.0f) {
+      /* Overexposure doesn't happen on screen itself and usually happens
+       * on light sources in the shot, this need to be checked separately
+       * because saturation and falloff calculation is based on the fact
+       * that pixels are not overexposed.
+       */
+      it.out[0] = 1.0f;
+    }
+    else {
+      const float saturation = get_pixel_saturation(pixel_color, m_screenBalance, primary_channel);
+      const float screen_saturation = get_pixel_saturation(
+          screen_color, m_screenBalance, primary_channel);
+
+      if (saturation < 0) {
+        /* Means main channel of pixel is different from screen,
+         * assume this is completely a foreground.
+         */
+        it.out[0] = 1.0f;
+      }
+      else if (saturation >= screen_saturation) {
+        /* Matched main channels and higher saturation on pixel
+         * is treated as completely background.
+         */
+        it.out[0] = 0.0f;
+      }
+      else {
+        /* Nice alpha falloff on edges. */
+        const float distance = 1.0f - saturation / screen_saturation;
+        it.out[0] = distance;
+      }
+    }
+  }
+}
+
 }  // namespace blender::compositor
diff --git a/source/blender/compositor/operations/COM_KeyingOperation.h b/source/blender/compositor/operations/COM_KeyingOperation.h
index 3d41ecaa0f6..e134ad54896 100644
--- a/source/blender/compositor/operations/COM_KeyingOperation.h
+++ b/source/blender/compositor/operations/COM_KeyingOperation.h
@@ -20,7 +20,7 @@
 
 #include <string.h>
 
-#include "COM_NodeOperation.h"
+#include "COM_MultiThreadedOperation.h"
 
 #include "BLI_listbase.h"
 
@@ -29,7 +29,7 @@ namespace blender::compositor {
 /**
  * Class with implementation of keying node
  */
-class KeyingOperation : public NodeOperation {
+class KeyingOperation : public MultiThreadedOperation {
  protected:
   SocketReader *m_pixelReader;
   SocketReader *m_screenReader;
@@ -48,6 +48,10 @@ class KeyingOperation : public NodeOperation {
   }
 
   void executePixelSampled(float output[4], float x, float y, PixelSampler sampler) override;
+
+  void update_memory_buffer_partial(MemoryBuffer *output,
+                                    const rcti &area,
+                                    Span<MemoryBuffer *> inputs) override;
 };
 
 }  // namespace blender::compositor