Add and use getColorChannel

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