Add and use getColorChannel

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()

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)

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('/', '_')