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

---
 datasets/raise/extract_noise.py |  4 ++--
 datasets/raise/utils.py         | 12 ++++++++++--
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/datasets/raise/extract_noise.py b/datasets/raise/extract_noise.py
index 80b58f8..c63692e 100755
--- a/datasets/raise/extract_noise.py
+++ b/datasets/raise/extract_noise.py
@@ -7,7 +7,7 @@ import os
 from tqdm import tqdm
 import csv
 import rawpy
-from utils import Color, denoise, iterativeMean
+from utils import Color, denoise, iterativeMean, isARawImage
 import matplotlib.pyplot as plt
 from scipy.ndimage import gaussian_filter
 
@@ -73,7 +73,7 @@ def getImageNpArray(imageFileName, computeExtremes, color):
             maxColor = colorRawImageVisibleMax
         return
 
-    if (imageFileName.endswith('.NEF') or imageFileName.endswith('.ARW')) and denoiser != 'mean':
+    if isARawImage(imageFileName) and denoiser != 'mean':
         imageNpArray = (imageNpArray - minColor) / (maxColor - minColor)
         # Pay attention to range of values expected by the denoiser.
         # Indeed if provide the thousands valued raw image, then the denoiser only returns values between 0 and 1 and making the difference between both look pointless.
diff --git a/datasets/raise/utils.py b/datasets/raise/utils.py
index 2b8623a..bed7f50 100644
--- a/datasets/raise/utils.py
+++ b/datasets/raise/utils.py
@@ -38,7 +38,12 @@ class iterativeMean:
             self.mean = ((self.mean * self.numberOfElementsInMean) + element) / (self.numberOfElementsInMean + 1)
         self.numberOfElementsInMean += 1
 
-def getRawColorChannel(raw):
+RAW_IMAGE_FILE_EXTENSIONS = [
+    'arw',
+    'nef',
+]
+
+def getRawColorChannel(raw, color):
     colorDesc = raw.color_desc.decode('ascii')
     assert colorDesc == 'RGBG'
     assert np.array_equal(raw.raw_pattern, np.array([[0, 1], [3, 2]], dtype = np.uint8))
@@ -61,4 +66,7 @@ def getRawColorChannel(raw):
             imageNpArray = greenBottomRawImageVisible
         case Color.BLUE:
             imageNpArray = blueRawImageVisible
-    return imageNpArray
\ No newline at end of file
+    return imageNpArray
+
+def isARawImage(imageFilePath):
+    return any([imageFileName.lower().endswith(f'.{rawImageFileExtension}') for rawImageFileExtension in RAW_IMAGE_FILE_EXTENSIONS])
\ No newline at end of file