Add and use getRawColorChannel to utils.py

datasets/raise/extract_noise.py

@@ -59,28 +59,7 @@ def getImageNpArray(imageFileName, computeExtremes, color):
         imageFilePath = f'{imagesFolderPath}/{imageFileName}'
     if imageFileName.endswith('.NEF') or imageFileName.endswith('.ARW'):
         with rawpy.imread(imageFilePath) as raw:
-            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))
-            # RG
-            # GB
-            rawImageVisible = raw.raw_image_visible.copy()
-
-            redRawImageVisible = rawImageVisible[::2, ::2]
-            greenRightRawImageVisible = rawImageVisible[::2, 1::2]
-
-            greenBottomRawImageVisible = rawImageVisible[1::2, ::2]
-            blueRawImageVisible = rawImageVisible[1::2, 1::2]
-
-            match color:
-                case Color.RED:
-                    imageNpArray = redRawImageVisible
-                case Color.GREEN_RIGHT:
-                    imageNpArray = greenRightRawImageVisible
-                case Color.GREEN_BOTTOM:
-                    imageNpArray = greenBottomRawImageVisible
-                case Color.BLUE:
-                    imageNpArray = blueRawImageVisible
+            imageNpArray = getRawColorChannel(raw)
     else:
         imagePil = Image.open(imageFilePath)
         imageNpArray = img_as_float(np.array(imagePil))

datasets/raise/utils.py

@@ -1,5 +1,6 @@
 from enum import Enum, auto
 import skimage.restoration
+import numpy as np
 
 class Color(Enum):
     RED = auto()
@@ -36,3 +37,28 @@ class iterativeMean:
         else:
             self.mean = ((self.mean * self.numberOfElementsInMean) + element) / (self.numberOfElementsInMean + 1)
         self.numberOfElementsInMean += 1
+
+def getRawColorChannel(raw):
+    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))
+    # RG
+    # GB
+    rawImageVisible = raw.raw_image_visible.copy()
+
+    redRawImageVisible = rawImageVisible[::2, ::2]
+    greenRightRawImageVisible = rawImageVisible[::2, 1::2]
+
+    greenBottomRawImageVisible = rawImageVisible[1::2, ::2]
+    blueRawImageVisible = rawImageVisible[1::2, 1::2]
+
+    match color:
+        case Color.RED:
+            imageNpArray = redRawImageVisible
+        case Color.GREEN_RIGHT:
+            imageNpArray = greenRightRawImageVisible
+        case Color.GREEN_BOTTOM:
+            imageNpArray = greenBottomRawImageVisible
+        case Color.BLUE:
+            imageNpArray = blueRawImageVisible
+    return imageNpArray