Output an estimated PRNU per color channel

datasets/raise/extract_noise.py

@@ -8,19 +8,16 @@ import os
 from tqdm import tqdm
 import csv
 import rawpy
+from enum import Enum, auto
 
 imagesFolderPath = '/mnt/HDD0/raise'
 imagesFolderPathFileName = imagesFolderPath.replace('/', '_')
 denoiser = 'wavelet'
-npArrayFilePath = f'mean_{imagesFolderPathFileName}_{denoiser}.npy'
 
 raiseNotFlatFields = False
 
 denoise = getattr(skimage.restoration, f'denoise_{denoiser}')
 
-mean = None
-numberOfImagesInMean = 0
-
 imagesFileNames = os.listdir(imagesFolderPath + ('/png' if raiseNotFlatFields else ''))
 
 if raiseNotFlatFields:
@@ -41,7 +38,17 @@ WALL = range(2_807, 2_912)
 minColor = None
 maxColor = None
 
-def treatImage(imageFileName, computeExtremes = False):
+class Color(Enum):
+    RED = auto()
+    GREEN_RIGHT = auto()
+    GREEN_BOTTOM = auto()
+    BLUE = auto()
+
+    def __str__(self):
+        return self.name.lower()
+
+# `color` is the actual color to estimate PRNU with.
+def treatImage(imageFileName, computeExtremes = False, color = None):
     global mean, numberOfImagesInMean, minColor, maxColor
     if raiseNotFlatFields:
         imageFileName = imageFileName.replace('.png', '.NEF')
@@ -63,8 +70,16 @@ def treatImage(imageFileName, computeExtremes = False):
             greenBottomRawImageVisible = rawImageVisible[1::2, ::2]
             blueRawImageVisible = rawImageVisible[1::2, 1::2]
 
-            # Actual color to estimate the PRNU with.
-            colorRawImageVisible = greenBottomRawImageVisible
+            match color:
+                case Color.RED:
+                    colorRawImageVisible = redRawImageVisible
+                case Color.GREEN_RIGHT:
+                    colorRawImageVisible = greenRightRawImageVisible
+                case Color.GREEN_BOTTOM:
+                    colorRawImageVisible = greenBottomRawImageVisible
+                case Color.BLUE:
+                    colorRawImageVisible = blueRawImageVisible
+
             if computeExtremes:
                 colorRawImageVisibleMin = colorRawImageVisible.min()
                 colorRawImageVisibleMax = colorRawImageVisible.max()
@@ -94,16 +109,23 @@ def treatImage(imageFileName, computeExtremes = False):
         mean = ((mean * numberOfImagesInMean) + imageNoiseNpArray) / (numberOfImagesInMean + 1)
     numberOfImagesInMean += 1
 
+# Assuming same intensity scale across color channels.
 for imageFileName in tqdm(imagesFileNames, 'Computing extremes of images'):
-    treatImage(imageFileName, computeExtremes = True)
+    for color in Color:
+        treatImage(imageFileName, computeExtremes = True, color = color)
 
 # To skip this step next time.
 # Maybe thanks to `rawpy.RawPy` fields, possibly stating device maximal value, can avoid doing so to some extent.
 print(f'{minColor=}')
 print(f'{maxColor=}')
 
-for imageFileName in tqdm(imagesFileNames, 'Denoising images'):
-    treatImage(imageFileName)
+for color in Color:
+    mean = None
+    numberOfImagesInMean = 0
 
-with open(npArrayFilePath, 'wb') as f:
-    np.save(f, mean)
+    for imageFileName in tqdm(imagesFileNames, 'Denoising images'):
+        treatImage(imageFileName)
+
+    npArrayFilePath = f'mean_{imagesFolderPathFileName}_{denoiser}_{color}.npy'
+    with open(npArrayFilePath, 'wb') as f:
+        np.save(f, mean)