datasets/raise/extract_noise.py

@@ -8,16 +8,19 @@ 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:
@@ -38,17 +41,7 @@ WALL = range(2_807, 2_912)
 minColor = None
 maxColor = None
 
-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):
+def treatImage(imageFileName, computeExtremes = False):
     global mean, numberOfImagesInMean, minColor, maxColor
     if raiseNotFlatFields:
         imageFileName = imageFileName.replace('.png', '.NEF')
@@ -70,16 +63,8 @@ def treatImage(imageFileName, computeExtremes = False, color = None):
             greenBottomRawImageVisible = rawImageVisible[1::2, ::2]
             blueRawImageVisible = rawImageVisible[1::2, 1::2]
 
-            match color:
-                case Color.RED:
-                    colorRawImageVisible = redRawImageVisible
-                case Color.GREEN_RIGHT:
-                    colorRawImageVisible = greenRightRawImageVisible
-                case Color.GREEN_BOTTOM:
-                    colorRawImageVisible = greenBottomRawImageVisible
-                case Color.BLUE:
-                    colorRawImageVisible = blueRawImageVisible
-
+            # Actual color to estimate the PRNU with.
+            colorRawImageVisible = greenBottomRawImageVisible
             if computeExtremes:
                 colorRawImageVisibleMin = colorRawImageVisible.min()
                 colorRawImageVisibleMax = colorRawImageVisible.max()
@@ -109,23 +94,16 @@ def treatImage(imageFileName, computeExtremes = False, color = None):
         mean = ((mean * numberOfImagesInMean) + imageNoiseNpArray) / (numberOfImagesInMean + 1)
     numberOfImagesInMean += 1
 
-# Assuming same intensity scale across color channels.
 for imageFileName in tqdm(imagesFileNames, 'Computing extremes of images'):
-    for color in Color:
-        treatImage(imageFileName, computeExtremes = True, color = color)
+    treatImage(imageFileName, computeExtremes = True)
 
 # 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 color in Color:
-    mean = None
-    numberOfImagesInMean = 0
+for imageFileName in tqdm(imagesFileNames, 'Denoising images'):
+    treatImage(imageFileName)
 
-    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)
+with open(npArrayFilePath, 'wb') as f:
+    np.save(f, mean)