Clean low pass denoiser

See [Robust_image_source_identification_on_modern_smartphones/issues/69#issuecomment-1892](#69 (comment)).

datasets/raise/utils.py

@@ -3,12 +3,15 @@ import skimage.restoration
 import numpy as np
 import rawpy
 from tqdm import tqdm
-from PIL import Image
+from PIL import Image, ImageDraw
 from skimage import img_as_float
 from datetime import datetime
 import builtins as __builtin__
 from scipy.ndimage import gaussian_filter
 import os
+from skimage.restoration import estimate_sigma
+from scipy import fftpack
+import imageio
 
 class Color(Enum):
     RED = auto()
@@ -24,16 +27,23 @@ class Denoiser(Enum):
     BILATERAL = auto()
     TV_CHAMBOLLE = auto()
     MEAN = auto()
+    NL_MEANS = auto()
+    LOW_PASS = auto()
 
     def __str__(self):
         return self.name.lower()
 
-# Among:
-# - `wavelet`
-# - `bilateral`
-# - `tv_chambolle`
+# Source: https://stackoverflow.com/a/43346070
+def getGaussianKernel(size, sigma):
+    axis = np.linspace(-(size - 1) / 2., (size - 1) / 2., size)
+    gauss = np.exp(-0.5 * np.square(axis) / np.square(sigma))
+    kernel = np.outer(gauss, gauss)
+    return kernel / np.sum(kernel)
+
+# `denoiser` can be all values of `Denoiser` except `MEAN`.
 def denoise(imageNpArray, denoiser):
-    skImageRestorationDenoise = getattr(skimage.restoration, f'denoise_{denoiser}')
+    if denoiser != Denoiser.LOW_PASS:
+        skImageRestorationDenoise = getattr(skimage.restoration, f'denoise_{denoiser}')
 
     match denoiser:
         case Denoiser.WAVELET:
@@ -42,6 +52,36 @@ def denoise(imageNpArray, denoiser):
             imageDenoisedNpArray = skImageRestorationDenoise(imageNpArray, sigma_color=0.05, sigma_spatial=15)
         case Denoiser.TV_CHAMBOLLE:
             imageDenoisedNpArray = skImageRestorationDenoise(imageNpArray, weight=0.2)
+        case Denoiser.NL_MEANS:
+            sigmaEst = np.mean(estimate_sigma(imageNpArray, channel_axis = -1))
+            imageDenoisedNpArray = skImageRestorationDenoise(imageNpArray, patch_size = 7, patch_distance = 4, h = 0.8 * sigmaEst, fast_mode = True)
+        case Denoiser.LOW_PASS:
+            # Based on [the Stack Overflow answer 54662336](https://stackoverflow.com/a/54662336).
+            # fft of image
+            fft1 = fftpack.fftshift(fftpack.fft2(imageNpArray))
+
+            # Padding as follows may not be perfectly correct but seems fine enough.
+            GAUSSIAN_SIGMA = 50
+
+            imageNpArrayShape = imageNpArray.shape
+            lowPassNpArraySize = min(imageNpArrayShape)
+            lowPassNpArray = getGaussianKernel(lowPassNpArraySize, GAUSSIAN_SIGMA)
+
+            wantedHeight, wantedWidth = imageNpArrayShape
+            verticalPadSize, horizontalPadSize = [(wantedDimension - lowPassNpArraySize) // 2 for wantedDimension in [wantedHeight, wantedWidth]]
+            rectangleLowPassNpArray = np.pad(lowPassNpArray, [(verticalPadSize, verticalPadSize), (horizontalPadSize, horizontalPadSize)])
+
+            # Unsure that doing such *normalization* is appropriate.
+            rectangleLowPassNpArray /= rectangleLowPassNpArray.max()
+
+            # multiply both the images
+            filtered = np.multiply(fft1, rectangleLowPassNpArray)
+
+            # inverse fft
+            ifft2 = np.real(fftpack.ifft2(fftpack.ifftshift(filtered)))
+            ifft2 = np.maximum(0, np.minimum(ifft2, 255))
+
+            imageDenoisedNpArray = ifft2.astype(np.uint8)
     return imageDenoisedNpArray
 
 class iterativeMean:
@@ -91,7 +131,7 @@ def isARawImage(imageFilePath):
 def getColorChannel(imageFilePath, color):
     if isARawImage(imageFilePath):
         numpyFilePath = f'{imageFilePath}.{color}.npy'
-        if os.path.isfile(numpyFilePath):
+        if False and os.path.isfile(numpyFilePath):
             imageNpArray = np.load(numpyFilePath)
         else:
             with rawpy.imread(imageFilePath) as raw:
@@ -144,7 +184,8 @@ def updateExtremes(imageNpArray, minColor, maxColor):
     return minColor, maxColor
 
 def print(*toPrint):
-    __builtin__.print(datetime.now(), *toPrint)
+    #__builtin__.print(datetime.now(), *toPrint)
+    tqdm.write(f'{datetime.now()} {str(toPrint)}')
 
 def getColorMeans(imagesFileNames, colors, singleColorChannelCropResolution = None):
     colorMeans = {}
@@ -159,19 +200,18 @@ def getColorMeans(imagesFileNames, colors, singleColorChannelCropResolution = No
         colorMeans[color] = colorIterativeMean.mean
     return colorMeans
 
-def getImageNpArray(imageFilePath, computeExtremes, color, denoiser):
-    global minColor, maxColor
+def getImageNpArray(imageFilePath, computeExtremes, color, denoiser, minColor, maxColor):
     imageNpArray = getColorChannel(imageFilePath, color)
 
     if computeExtremes:
         minColor, maxColor = updateExtremes(imageNpArray, minColor, maxColor)
-        return
+        return None, minColor, maxColor
 
     if isARawImage(imageFilePath) and denoiser != Denoiser.MEAN:
         imageNpArray = rescaleRawImageForDenoiser(imageNpArray, minColor, maxColor)
         # Pay attention to range of values expected by the denoiser.
         # Indeed if provide the thousands valued raw image, then the denoiser only returns values between 0 and 1 and making the difference between both look pointless.
-    return imageNpArray
+    return imageNpArray, minColor, maxColor
 
 def getImageCrop(image, cropResolution):
     imageCrop = image[:cropResolution[0], :cropResolution[1]]