update EEH module: replace bilateral filter with gaussian filter

Author: QiuJueqin

modules/eeh.py

@@ -7,52 +7,44 @@
 import numpy as np
 
 from .basic_module import BasicModule, register_dependent_modules
-from .helpers import bilateral_filter, gen_gaussian_kernel
+from .helpers import gaussian_filter, gen_gaussian_kernel
 
 
 @register_dependent_modules('csc')
 class EEH(BasicModule):
     def __init__(self, cfg):
         super().__init__(cfg)
 
-        self.intensity_weights_lut = self.get_intensity_weights_lut(intensity_sigma=1.0)  # x1024
-        spatial_weights = gen_gaussian_kernel(kernel_size=5, sigma=1.0)
-        self.spatial_weights = (1024 * spatial_weights / spatial_weights.max()).astype(np.int32)  # x1024
+        kernel = gen_gaussian_kernel(kernel_size=7, sigma=5.0)
+        self.kernel = (1024 * kernel / kernel.max()).astype(np.int32)  # x1024
 
         flat_slope = self.params.middle_threshold / (self.params.middle_threshold - self.params.flat_threshold + 1E-6)
         edge_slope = self.params.edge_gain / 256
 
         self.flat_slope = np.array(256 * flat_slope, dtype=np.int32)  # x256
         self.edge_slope = np.array(256 * edge_slope, dtype=np.int32)  # x256
-        self.flat_intercept = np.array(-flat_slope * self.params.flat_threshold, dtype=np.int32)
-        self.edge_intercept = np.array((1 - edge_slope) * self.params.edge_threshold, dtype=np.int32)
+        self.flat_intercept = -np.array(256 * flat_slope * self.params.flat_threshold, dtype=np.int32)  # x256
+        self.edge_intercept = np.array(256 * (1 - edge_slope) * self.params.edge_threshold, dtype=np.int32)  # x256
 
     def execute(self, data):
         y_image = data['y_image'].astype(np.int32)
 
-        bf_y_image = bilateral_filter(y_image, self.spatial_weights, self.intensity_weights_lut, right_shift=10)
+        gf_y_image = gaussian_filter(y_image, self.kernel)
 
-        delta = y_image - bf_y_image
+        delta = y_image - gf_y_image
         sign_map = np.sign(delta)
         abs_delta = np.abs(delta)
 
-        flat_delta = np.right_shift(self.flat_slope * abs_delta, 8) + self.flat_intercept
-        edge_delta = np.right_shift(self.edge_slope * abs_delta, 8) + self.edge_intercept
+        flat_delta = np.right_shift(self.flat_slope * abs_delta + self.flat_intercept, 8)
+        edge_delta = np.right_shift(self.edge_slope * abs_delta + self.edge_intercept, 8)
         enhanced_delta = sign_map * (
                 (abs_delta > self.params.flat_threshold) * (abs_delta <= self.params.middle_threshold) * flat_delta +
                 (abs_delta > self.params.middle_threshold) * (abs_delta <= self.params.edge_threshold) * abs_delta +
                 (abs_delta > self.params.edge_threshold) * edge_delta
         )
         enhanced_delta = np.clip(enhanced_delta, -self.params.delta_threshold, self.params.delta_threshold)
 
-        eeh_y_image = np.clip(bf_y_image + enhanced_delta, 0, self.cfg.saturation_values.sdr)
+        eeh_y_image = np.clip(gf_y_image + enhanced_delta, 0, self.cfg.saturation_values.sdr)
 
         data['y_image'] = eeh_y_image.astype(np.uint8)
         data['edge_map'] = delta
-
-    @staticmethod
-    def get_intensity_weights_lut(intensity_sigma):
-        intensity_diff = np.arange(255 ** 2)
-        exp_lut = 1024 * np.exp(-intensity_diff / (2.0 * (255 * intensity_sigma) ** 2))
-        return exp_lut.astype(np.int32)  # x1024
-

modules/helpers.py

@@ -139,6 +139,32 @@ def mean_filter(array, filter_size=3):
     return (sum(shifted_arrays) / filter_size ** 2).astype(array.dtype)
 
 
+def gaussian_filter(array, kernel):
+    """
+    A faster reimplementation of the bilateral-filter
+    :param array: array to be filter: np.ndarray(H, W, ...), must be np.int dtype
+    :param kernel: np.ndarray(h, w)
+    :return: filtered array: np.ndarray(H, W, ...)
+    """
+
+    kh, kw = kernel.shape[:2]
+    kernel = kernel.flatten()
+
+    padded_array = pad(array, pads=(kh // 2, kw // 2))
+    shifted_arrays = shift_array(padded_array, window_size=(kh, kw))
+
+    gf_array = np.zeros_like(array)
+    weights = np.zeros_like(array)
+
+    for i, shifted_array in enumerate(shifted_arrays):
+        gf_array += kernel[i] * shifted_array
+        weights += kernel[i]
+
+    gf_array = (gf_array / weights).astype(array.dtype)
+
+    return gf_array
+
+
 def bilateral_filter(array, spatial_weights, intensity_weights_lut, right_shift=0):
     """
     A faster reimplementation of the bilateral-filter
@@ -170,4 +196,4 @@ def bilateral_filter(array, spatial_weights, intensity_weights_lut, right_shift=
 
     bf_array = (bf_array / weights).astype(array.dtype)
 
-    return bf_array
+    return bf_array