Optimized box blur implementation (#1974)

backend/src/packages/chaiNNer_standard/image_filter/blur/box_blur.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import math
 from math import ceil
 
 import cv2
@@ -11,6 +12,39 @@
 from .. import blur_group
 
 
+def get_kernel_1d(radius: float) -> np.ndarray:
+    kernel = np.ones(ceil(radius) * 2 + 1, np.float32)
+
+    d = radius % 1
+    if d != 0:
+        kernel[0] *= d
+        kernel[-1] *= d
+
+    # normalize
+    kernel /= np.sum(kernel)
+
+    return kernel
+
+
+def get_kernel_2d(radius_x: float, radius_y) -> np.ndarray:
+    # Create kernel of dims h * w, rounded up to the closest odd integer
+    kernel = np.ones((ceil(radius_y) * 2 + 1, ceil(radius_x) * 2 + 1), np.float32) / (
+        (2 * radius_y + 1) * (2 * radius_x + 1)
+    )
+
+    # Modify edges of kernel by fractional amount if kernel size (2r+1) is not odd integer
+    x_d = radius_x % 1
+    y_d = radius_y % 1
+    if y_d != 0:
+        kernel[0, :] *= y_d
+        kernel[-1, :] *= y_d
+    if x_d != 0:
+        kernel[:, 0] *= x_d
+        kernel[:, -1] *= x_d
+
+    return kernel
+
+
 @blur_group.register(
     schema_id="chainner:image:blur",
     name="Box Blur",
@@ -46,23 +80,26 @@ def box_blur_node(
     radius_x: float,
     radius_y: float,
 ) -> np.ndarray:
-    if radius_x == 0 and radius_y == 0:
+    if radius_x == 0 or radius_y == 0:
         return img
 
-    # Create kernel of dims h * w, rounded up to the closest odd integer
-    kernel = np.ones((ceil(radius_y) * 2 + 1, ceil(radius_x) * 2 + 1), np.float32) / (
-        (2 * radius_y + 1) * (2 * radius_x + 1)
-    )
-
-    # Modify edges of kernel by fractional amount if kernel size (2r+1) is not odd integer
-    x_d = radius_x % 1
-    y_d = radius_y % 1
-    if y_d != 0:
-        kernel[0, :] *= y_d
-        kernel[-1, :] *= y_d
-    if x_d != 0:
-        kernel[:, 0] *= x_d
-        kernel[:, -1] *= x_d
+    # Separable filter is faster for relatively small kernels, but after a certain size it becomes
+    # slower than filter2D's DFT implementation. The exact cutoff depends on the hardware.
+    avg_radius = math.sqrt(radius_x * radius_y)
+    use_sep = avg_radius < 70
 
-    # Linear filter with reflected padding
-    return cv2.filter2D(img, -1, kernel, borderType=cv2.BORDER_REFLECT_101)
+    if use_sep:
+        return cv2.sepFilter2D(
+            img,
+            -1,
+            get_kernel_1d(radius_x),
+            get_kernel_1d(radius_y),
+            borderType=cv2.BORDER_REFLECT_101,
+        )
+    else:
+        return cv2.filter2D(
+            img,
+            -1,
+            get_kernel_2d(radius_x, radius_y),
+            borderType=cv2.BORDER_REFLECT_101,
+        )