From 2d67c2bca18faaaca7845f7b34d45b486688024b Mon Sep 17 00:00:00 2001
From: Benjamin Loison <12752145+Benjamin-Loison@users.noreply.github.com>
Date: Tue, 30 Apr 2024 02:04:17 +0200
Subject: [PATCH] Add and use `rescaleRawImageForDenoiser` and `updateExtremes`
 (#59)

---
 datasets/raise/extract_noise.py               | 11 ++---
 .../split_and_compare_prnus_of_subgroups.py   | 45 ++++++++++++-------
 datasets/raise/utils.py                       | 13 ++++++
 3 files changed, 44 insertions(+), 25 deletions(-)

diff --git a/datasets/raise/extract_noise.py b/datasets/raise/extract_noise.py
index 1db2b78..7066ba7 100755
--- a/datasets/raise/extract_noise.py
+++ b/datasets/raise/extract_noise.py
@@ -4,7 +4,7 @@ import numpy as np
 import os
 from tqdm import tqdm
 import csv
-from utils import Color, denoise, iterativeMean, isARawImage, escapeFilePath, getColorChannel, saveNpArray
+from utils import Color, denoise, iterativeMean, isARawImage, escapeFilePath, getColorChannel, saveNpArray, rescaleRawImageForDenoiser, updateExtremes
 import matplotlib.pyplot as plt
 from scipy.ndimage import gaussian_filter
 
@@ -57,16 +57,11 @@ def getImageNpArray(imageFileName, computeExtremes, color):
     imageNpArray = getColorChannel(imageFilePath)
 
     if computeExtremes:
-        colorRawImageVisibleMin = imageNpArray.min()
-        colorRawImageVisibleMax = imageNpArray.max()
-        if minColor is None or colorRawImageVisibleMin < minColor:
-            minColor = colorRawImageVisibleMin
-        if maxColor is None or colorRawImageVisibleMax > maxColor:
-            maxColor = colorRawImageVisibleMax
+        minColor, maxColor = updateExtremes(imageNpArray, minColor, maxColor)
         return
 
     if isARawImage(imageFileName) and denoiser != 'mean':
-        imageNpArray = (imageNpArray - minColor) / (maxColor - minColor)
+        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
diff --git a/datasets/raise/split_and_compare_prnus_of_subgroups.py b/datasets/raise/split_and_compare_prnus_of_subgroups.py
index cab3a16..6fc5e44 100755
--- a/datasets/raise/split_and_compare_prnus_of_subgroups.py
+++ b/datasets/raise/split_and_compare_prnus_of_subgroups.py
@@ -3,7 +3,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from tqdm import tqdm
-from utils import denoise, iterativeMean, getColorChannel, escapeFilePath, Color, mergeSingleColorChannelImagesAccordingToBayerFilter
+from utils import denoise, iterativeMean, getColorChannel, escapeFilePath, Color, mergeSingleColorChannelImagesAccordingToBayerFilter, rescaleRawImageForDenoiser, updateExtremes
 import sys
 import os
 from random import shuffle
@@ -26,24 +26,35 @@ numberOfImagesPerSubgroup = len(imagesFileNames) // NUMBER_OF_SUBGROUPS
 
 subgroupsIterativeMean = [iterativeMean() for _ in range(NUMBER_OF_SUBGROUPS)]
 rmss = []
-for subgroupImageIndex in tqdm(range(numberOfImagesPerSubgroup), 'Subgroup image index'):
-    for subgroupIndex in tqdm(range(NUMBER_OF_SUBGROUPS), 'Subgroup'):
-        imageIndex = (subgroupIndex * NUMBER_OF_SUBGROUPS) + subgroupImageIndex
-        imageFileName = imagesFileNames[imageIndex]
-        imageFilePath = f'{IMAGES_FOLDER}/{imageFileName}'
-        singleColorChannelImages = {color: getColorChannel(imageFilePath, color) for color in Color}
-        multipleColorsImage = mergeSingleColorChannelImagesAccordingToBayerFilter(singleColorChannelImages)
 
-        singleColorChannelDenoisedImages = {color: denoise(singleColorChannelImages[color], DENOISER) for color in Color}
-        multipleColorsDenoisedImage = mergeSingleColorChannelImagesAccordingToBayerFilter(singleColorChannelDenoisedImages)
+minColor = None
+maxColor = None
 
-        imagePrnuEstimateNpArray = multipleColorsImage - multipleColorsDenoisedImage
-        subgroupIterativeMean = subgroupsIterativeMean[subgroupIndex]
-        subgroupIterativeMean.add(imagePrnuEstimateNpArray)
-        if subgroupIndex == NUMBER_OF_SUBGROUPS - 1:
-            assert NUMBER_OF_SUBGROUPS == 2
-            rms = rmsDiffNumpy(subgroupIterativeMean.mean, subgroupsIterativeMean[1 - subgroupIndex].mean)
-            rmss += [rms]
+for computeExtremes in tqdm(([True] if minColor is None or maxColor is None else []) + [False], 'Compute extremes'):
+    for subgroupImageIndex in tqdm(range(numberOfImagesPerSubgroup), 'Subgroup image index'):
+        for subgroupIndex in tqdm(range(NUMBER_OF_SUBGROUPS), 'Subgroup'):
+            imageIndex = (subgroupIndex * NUMBER_OF_SUBGROUPS) + subgroupImageIndex
+            imageFileName = imagesFileNames[imageIndex]
+            imageFilePath = f'{IMAGES_FOLDER}/{imageFileName}'
+            returnSingleColorChannelImage = lambda singleColorChannelImage, _minColor, _maxColor: singleColorChannelImage
+            rescaleIfNeeded = returnSingleColorChannelImage if computeExtremes else rescaleRawImageForDenoiser
+            singleColorChannelImages = {color: rescaleIfNeeded(getColorChannel(imageFilePath, color), minColor, maxColor) for color in Color}
+            multipleColorsImage = mergeSingleColorChannelImagesAccordingToBayerFilter(singleColorChannelImages)
+
+            if computeExtremes:
+                minColor, maxColor = updateExtremes(multipleColorsImage, minColor, maxColor)
+                continue
+
+            singleColorChannelDenoisedImages = {color: denoise(singleColorChannelImages[color], DENOISER) for color in Color}
+            multipleColorsDenoisedImage = mergeSingleColorChannelImagesAccordingToBayerFilter(singleColorChannelDenoisedImages)
+
+            imagePrnuEstimateNpArray = multipleColorsImage - multipleColorsDenoisedImage
+            subgroupIterativeMean = subgroupsIterativeMean[subgroupIndex]
+            subgroupIterativeMean.add(imagePrnuEstimateNpArray)
+            if subgroupIndex == NUMBER_OF_SUBGROUPS - 1:
+                assert NUMBER_OF_SUBGROUPS == 2
+                rms = rmsDiffNumpy(subgroupIterativeMean.mean, subgroupsIterativeMean[1 - subgroupIndex].mean)
+                rmss += [rms]
 
 mostImagesSubgroupPrnuEstimatesNpArray = [subgroupIterativeMean.mean for subgroupIterativeMean in subgroupsIterativeMean]
 minimum = np.min(mostImagesSubgroupPrnuEstimatesNpArray)
diff --git a/datasets/raise/utils.py b/datasets/raise/utils.py
index 825c73b..3a8fa3b 100644
--- a/datasets/raise/utils.py
+++ b/datasets/raise/utils.py
@@ -124,3 +124,16 @@ def mergeSingleColorChannelImagesAccordingToBayerFilter(singleColorChannelImages
 def saveNpArray(fileName, npArray):
     with open(f'{fileName}.npy', 'wb') as f:
         np.save(f, npArray)
+
+def rescaleRawImageForDenoiser(imageNpArray, minColor, maxColor):
+    imageNpArray = (imageNpArray - minColor) / (maxColor - minColor)
+    return imageNpArray
+
+def updateExtremes(imageNpArray, minColor, maxColor):
+    colorRawImageVisibleMin = imageNpArray.min()
+    colorRawImageVisibleMax = imageNpArray.max()
+    if minColor is None or colorRawImageVisibleMin < minColor:
+        minColor = colorRawImageVisibleMin
+    if maxColor is None or colorRawImageVisibleMax > maxColor:
+        maxColor = colorRawImageVisibleMax
+    return minColor, maxColor