From bc06612501a305e50e7b3107a3264ab31bb326dd Mon Sep 17 00:00:00 2001
From: Benjamin Loison <12752145+Benjamin-Loison@users.noreply.github.com>
Date: Mon, 8 Apr 2024 04:52:29 +0200
Subject: [PATCH] Make `extract_noise.py` easier to configure by providing
 `denoiser`

---
 datasets/raise/extract_noise.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/datasets/raise/extract_noise.py b/datasets/raise/extract_noise.py
index 1fa12d7..dd3aba5 100755
--- a/datasets/raise/extract_noise.py
+++ b/datasets/raise/extract_noise.py
@@ -1,6 +1,6 @@
 #!/usr/bin/python3
 
-from skimage.restoration import denoise_tv_chambolle
+import skimage.restoration
 from skimage import img_as_float
 import numpy as np
 from PIL import Image
@@ -8,7 +8,10 @@ import os
 from tqdm import tqdm
 
 imagesFolderPath = 'flat-field'
-npArrayFilePath = 'mean.npy'
+denoiser = 'wavelet'
+npArrayFilePath = f'mean_{imagesFolderPath}_{denoiser}.npy'
+
+denoise = getattr(skimage.restoration, f'denoise_{denoiser}')
 
 mean = None
 numberOfImagesInMean = 0
@@ -20,7 +23,8 @@ def treatImage(imageFileName):
     imageFilePath = f'{imagesFolderPath}/{imageFileName}'
     imagePil = Image.open(imageFilePath)
     imageNpArray = img_as_float(np.array(imagePil))
-    imageDenoisedNpArray = denoise_tv_chambolle(imageNpArray, weight=0.2, channel_axis=-1)
+    # As the arguments differ from a denoiser to the other should use a match instead.
+    imageDenoisedNpArray = denoise(imageNpArray, channel_axis=-1, convert2ycbcr=True, rescale_sigma=True)
     imageNoiseNpArray = imageNpArray - imageDenoisedNpArray
     if mean is None:
         mean = imageNoiseNpArray