From f15af68bbab662f5c59ea12e8947562b0f40d47b Mon Sep 17 00:00:00 2001
From: Benjamin Loison <12752145+Benjamin-Loison@users.noreply.github.com>
Date: Wed, 17 Apr 2024 20:15:33 +0200
Subject: [PATCH] Add and use `Color` enumeration

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

diff --git a/datasets/raise/extract_noise.py b/datasets/raise/extract_noise.py
index 1686d34..6121b10 100755
--- a/datasets/raise/extract_noise.py
+++ b/datasets/raise/extract_noise.py
@@ -8,6 +8,7 @@ import os
 from tqdm import tqdm
 import csv
 import rawpy
+from enum import Enum, auto
 
 imagesFolderPath = '/mnt/HDD0/raise'
 imagesFolderPathFileName = imagesFolderPath.replace('/', '_')
@@ -41,7 +42,14 @@ 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()
+
+# `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 +71,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()
@@ -95,7 +111,7 @@ def treatImage(imageFileName, computeExtremes = False):
     numberOfImagesInMean += 1
 
 for imageFileName in tqdm(imagesFileNames, 'Computing extremes of images'):
-    treatImage(imageFileName, computeExtremes = True)
+    treatImage(imageFileName, computeExtremes = True, color = Color.GREEN)
 
 # To skip this step next time.
 # Maybe thanks to `rawpy.RawPy` fields, possibly stating device maximal value, can avoid doing so to some extent.