Add and use getColorChannel

datasets/raise/extract_noise.py

@@ -6,13 +6,12 @@ from PIL import Image
 import os
 from tqdm import tqdm
 import csv
-import rawpy
-from utils import Color, denoise, iterativeMean, isARawImage
+from utils import Color, denoise, iterativeMean, isARawImage, escapeFilePath, getColorChannel
 import matplotlib.pyplot as plt
 from scipy.ndimage import gaussian_filter
 
 imagesFolderPath = 'rafael/arw'
-imagesFolderPathFileName = imagesFolderPath.replace('/', '_')
+imagesFolderPathFileName = escapeFilePath(imagesFolderPath)
 # Among:
 # `denoise` possible denoisers and `mean`.
 denoiser = 'mean'
@@ -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()

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
+from utils import denoise, iterativeMean, getColorChannel, escapeFilePath, Color
 from skimage import img_as_float
 import sys
 import os
@@ -18,7 +18,7 @@ NUMBER_OF_SUBGROUPS = 2
 DENOISER = 'wavelet'
 IMAGES_FOLDER = 'flat-field/TIF'
 
-setting = IMAGES_FOLDER.replace('/', '_') + f'_{DENOISER}'
+setting = escapeFilePath(IMAGES_FOLDER) + f'_{DENOISER}'
 
 imagesFileNames = os.listdir(IMAGES_FOLDER)
 # To not have a bias (chronological for instance) when split to make subgroups.
@@ -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)

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()
@@ -69,4 +70,16 @@ def getRawColorChannel(raw, color):
     return imageNpArray
 
 def isARawImage(imageFilePath):
-    return any([imageFileName.lower().endswith(f'.{rawImageFileExtension}') for rawImageFileExtension in RAW_IMAGE_FILE_EXTENSIONS])
+    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('/', '_')