Add and use isARawImage

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.

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
+    return imageNpArray
+
+def isARawImage(imageFilePath):
+    return any([imageFileName.lower().endswith(f'.{rawImageFileExtension}') for rawImageFileExtension in RAW_IMAGE_FILE_EXTENSIONS])