From 8c9dfc2e4111efcb71e9f6a2e66f3e313ed0c32e Mon Sep 17 00:00:00 2001
From: Benjamin Loison <12752145+Benjamin-Loison@users.noreply.github.com>
Date: Wed, 17 Apr 2024 20:25:14 +0200
Subject: [PATCH] Output an estimated PRNU per color channel

---
 datasets/raise/extract_noise.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/datasets/raise/extract_noise.py b/datasets/raise/extract_noise.py
index 6121b10..359971f 100755
--- a/datasets/raise/extract_noise.py
+++ b/datasets/raise/extract_noise.py
@@ -13,15 +13,11 @@ 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:
@@ -48,6 +44,9 @@ class Color(Enum):
     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
@@ -110,16 +109,23 @@ 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'):
-    treatImage(imageFileName, computeExtremes = True, color = Color.GREEN)
+    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)