From 5ea0a862c0c5c7dd31c38920bed7260e3a0dcab0 Mon Sep 17 00:00:00 2001
From: Benjamin Loison <benjamin_loison@users.noreply.gitea.lemnoslife.com>
Date: Sat, 27 Apr 2024 21:49:32 +0200
Subject: [PATCH] Add and use `getColorChannel`

---
 datasets/raise/extract_noise.py                        | 10 ++--------
 datasets/raise/split_and_compare_prnus_of_subgroups.py |  6 +++---
 datasets/raise/utils.py                                | 10 ++++++++++
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/datasets/raise/extract_noise.py b/datasets/raise/extract_noise.py
index 8ccd865..8f0f194 100755
--- a/datasets/raise/extract_noise.py
+++ b/datasets/raise/extract_noise.py
@@ -6,8 +6,7 @@ from PIL import Image
 import os
 from tqdm import tqdm
 import csv
-import rawpy
-from utils import Color, denoise, iterativeMean, isARawImage, escapeFilePath
+from utils import Color, denoise, iterativeMean, isARawImage, escapeFilePath, getColorChannel
 import matplotlib.pyplot as plt
 from scipy.ndimage import gaussian_filter
 
@@ -57,12 +56,7 @@ def getImageNpArray(imageFileName, computeExtremes, color):
         imageFilePath = f'{imagesFolderPath}/nef/{imageFileName}'
     else:
         imageFilePath = f'{imagesFolderPath}/{imageFileName}'
-    if imageFileName.endswith('.NEF') or imageFileName.endswith('.ARW'):
-        with rawpy.imread(imageFilePath) as raw:
-            imageNpArray = getRawColorChannel(raw)
-    else:
-        imagePil = Image.open(imageFilePath)
-        imageNpArray = img_as_float(np.array(imagePil))
+    imageNpArray = getColorChannel(imageFilePath)
 
     if computeExtremes:
         colorRawImageVisibleMin = imageNpArray.min()
diff --git a/datasets/raise/split_and_compare_prnus_of_subgroups.py b/datasets/raise/split_and_compare_prnus_of_subgroups.py
index 195c15d..dbd279e 100755
--- a/datasets/raise/split_and_compare_prnus_of_subgroups.py
+++ b/datasets/raise/split_and_compare_prnus_of_subgroups.py
@@ -4,7 +4,7 @@ from PIL import Image
 import numpy as np
 import matplotlib.pyplot as plt
 from tqdm import tqdm
-from utils import denoise, iterativeMean, escapeFilePath
+from utils import denoise, iterativeMean, getColorChannel, escapeFilePath, Color
 from skimage import img_as_float
 import sys
 import os
@@ -33,8 +33,8 @@ for subgroupImageIndex in tqdm(range(numberOfImagesPerSubgroup), 'Subgroup image
         imageIndex = (subgroupIndex * NUMBER_OF_SUBGROUPS) + subgroupImageIndex
         imageFileName = imagesFileNames[imageIndex]
         imageFilePath = f'{IMAGES_FOLDER}/{imageFileName}'
-        imagePil = Image.open(imageFilePath)
-        imageNpArray = img_as_float(np.array(imagePil))
+        imageNpArray = getColorChannel(imageFilePath, Color.GREEN_RIGHT)
+
         imagePrnuEstimateNpArray = imageNpArray - denoise(imageNpArray, DENOISER)
         subgroupIterativeMean = subgroupsIterativeMean[subgroupIndex]
         subgroupIterativeMean.add(imagePrnuEstimateNpArray)
diff --git a/datasets/raise/utils.py b/datasets/raise/utils.py
index eff109a..fbf6235 100644
--- a/datasets/raise/utils.py
+++ b/datasets/raise/utils.py
@@ -1,6 +1,7 @@
 from enum import Enum, auto
 import skimage.restoration
 import numpy as np
+import rawpy
 
 class Color(Enum):
     RED = auto()
@@ -71,5 +72,14 @@ def getRawColorChannel(raw, color):
 def isARawImage(imageFilePath):
     return any([imageFileName.lower().endswith(f'.{rawImageFileExtension}') for rawImageFileExtension in RAW_IMAGE_FILE_EXTENSIONS])
 
+def getColorChannel(imageFilePath, color):
+    if isARawImage(imageFilePath):
+        with rawpy.imread(imageFilePath) as raw:
+            imageNpArray = getRawColorChannel(raw, color)
+    else:
+        imagePil = Image.open(imageFilePath)
+        imageNpArray = img_as_float(np.array(imagePil))
+    return imageNpArray
+
 def escapeFilePath(filePath):
     return filePath.replace('/', '_')
\ No newline at end of file