From a0fc38d023b66caaca0b896f061a9004d674a4f9 Mon Sep 17 00:00:00 2001
From: Benjamin Loison <benjamin_loison@users.noreply.gitea.lemnoslife.com>
Date: Sat, 27 Apr 2024 20:47:33 +0200
Subject: [PATCH] Add and use `getRawColorChannel` to `utils.py`

---
 datasets/raise/extract_noise.py | 23 +----------------------
 datasets/raise/utils.py         | 26 ++++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/datasets/raise/extract_noise.py b/datasets/raise/extract_noise.py
index b2a15c7..80b58f8 100755
--- a/datasets/raise/extract_noise.py
+++ b/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))
diff --git a/datasets/raise/utils.py b/datasets/raise/utils.py
index 1ce8916..2b8623a 100644
--- a/datasets/raise/utils.py
+++ b/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
\ No newline at end of file