Upload New File

src/methods/EL-MURR-Theresa/malvar.py

+import numpy as np
+from scipy.signal import correlate2d
+from src.forward_model import CFA
+
+def malvar_he_cutler(y: np.ndarray, op: CFA ) -> np.ndarray:
+    """Performs demosaicing using the malvar-he-cutler algorithm
+
+    Args:
+        op (CFA): CFA operator.
+        y (np.ndarray): Mosaicked image.
+
+    Returns:
+        np.ndarray: Demosaicked image.
+    """
+
+    red_mask, green_mask, blue_mask =  [op.mask[:, :, 0], op.mask[:, :, 1], op.mask[:, :, 2]]
+    mosaicked_image = np.float32(y)
+    demosaicked_image = np.empty(op.input_shape)
+
+    if op.cfa == 'quad_bayer':
+        filters = get_quad_bayer_filters()
+    else:
+        filters = get_default_filters()
+
+    demosaicked_image = apply_demosaicking_filters(
+        mosaicked_image,demosaicked_image, red_mask, green_mask, blue_mask, filters
+    )
+
+    return demosaicked_image
+
+def get_quad_bayer_filters():
+    coefficient_scale = 0.03125
+    return {
+        "G_at_R_and_B": np.array([
+            [0, 0, 0, 0, -1, -1, 0, 0, 0, 0],
+            [0, 0, 0, 0, -1, -1, 0, 0, 0, 0],
+            [0, 0, 0, 0, 2, 2, 0, 0, 0, 0],
+            [0, 0, 0, 0, 2, 2, 0, 0, 0, 0],
+            [-1, -1, 2, 2, 4, 4, 2, 2, -1, -1],
+            [-1, -1, 2, 2, 4, 4, 2, 2, -1, -1],
+            [0, 0, 0, 0, 2, 2, 0, 0, 0, 0],
+            [0, 0, 0, 0, 2, 2, 0, 0, 0, 0],
+            [0, 0, 0, 0, -1, -1, 0, 0, 0, 0],
+            [0, 0, 0, 0, -1, -1, 0, 0, 0, 0]
+        ]) * coefficient_scale,
+        "R_at_GR_and_B_at_GB": np.array([
+            [0, 0, 0, 0, 0.5, 0.5, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0.5, 0.5, 0, 0, 0, 0],
+            [0, 0, -1, -1, 0, 0, -1, -1, 0, 0],
+            [0, 0, -1, -1, 0, 0, -1, -1, 0, 0],
+            [-1, -1, 4, 4, 5, 5, 4, 4, -1, -1],
+            [-1, -1, 4, 4, 5, 5, 4, 4, -1, -1],
+            [0, 0, -1, -1, 0, 0, -1, -1, 0, 0],
+            [0, 0, -1, -1, 0, 0, -1, -1, 0, 0],
+            [0, 0, 0, 0, 0.5, 0.5, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0.5, 0.5, 0, 0, 0, 0]
+        ]) * coefficient_scale,
+        "R_at_GB_and_B_at_GR": np.array([
+            [0, 0, 0, 0, -1, -1, 0, 0, 0, 0],
+            [0, 0, 0, 0, -1, -1, 0, 0, 0, 0],
+            [0, 0, -1, -1, 4, 4, -1, -1, 0, 0],
+            [0, 0, -1, -1, 4, 4, -1, -1, 0, 0],
+            [0.5, 0.5, 0, 0, 5, 5, 0, 0, 0.5, 0.5],
+            [0.5, 0.5, 0, 0, 5, 5, 0, 0, 0.5, 0.5],
+            [0, 0, -1, -1, 4, 4, -1, -1, 0, 0],
+            [0, 0, -1, -1, 4, 4, -1, -1, 0, 0],
+            [0, 0, 0, 0, -1, -1, 0, 0, 0, 0],
+            [0, 0, 0, 0, -1, -1, 0, 0, 0, 0]
+        ]) * coefficient_scale,
+        "R_at_B_and_B_at_R": np.array([
+            [0, 0, 0, 0, -1.5, -1.5, 0, 0, 0, 0],
+            [0, 0, 0, 0, -1.5, -1.5, 0, 0, 0, 0],
+            [0, 0, 2, 2, 0, 0, 2, 2, 0, 0],
+            [0, 0, 2, 2, 0, 0, 2, 2, 0, 0],
+            [-1.5, -1.5, 0, 0, 6, 6, 0, 0, -1.5, -1.5],
+            [-1.5, -1.5, 0, 0, 6, 6, 0, 0, -1.5, -1.5],
+            [0, 0, 2, 2, 0, 0, 2, 2, 0, 0],
+            [0, 0, 2, 2, 0, 0, 2, 2, 0, 0],
+            [0, 0, 0, 0, -1.5, -1.5, 0, 0, 0, 0],
+            [0, 0, 0, 0, -1.5, -1.5, 0, 0, 0, 0]
+        ]) * coefficient_scale,
+    }
+
+def get_default_filters():
+    coefficient_scale = 0.125
+    return {
+        "G_at_R_and_B": np.array([
+            [0, 0, -1, 0, 0],
+            [0, 0, 2, 0, 0],
+            [-1, 2, 4, 2, -1],
+            [0, 0, 2, 0, 0],
+            [0, 0, -1, 0, 0]
+        ]) * coefficient_scale,
+        "R_at_GR_and_B_at_GB": np.array([
+            [0, 0, 0.5, 0, 0],
+            [0, -1, 0, -1, 0],
+            [-1, 4, 5, 4, -1],
+            [0, -1, 0, -1, 0],
+            [0, 0, 0.5, 0, 0]
+        ]) * coefficient_scale,
+        "R_at_GB_and_B_at_GR": np.array([
+            [0, 0, -1, 0, 0],
+            [0, -1, 4, -1, 0],
+            [0.5, 0, 5, 0, 0.5],
+            [0, -1, 4, -1, 0],
+            [0, 0, -1, 0, 0]
+        ]) * coefficient_scale,
+        "R_at_B_and_B_at_R": np.array([
+            [0, 0, -1.5, 0, 0],
+            [0, 2, 0, 2, 0],
+            [-1.5, 0, 6, 0, -1.5],
+            [0, 2, 0, 2, 0],
+            [0, 0, -1.5, 0, 0]
+        ]) * coefficient_scale,
+    }
+
+def apply_demosaicking_filters(image, res, red_mask, green_mask, blue_mask, filters):
+    red_channel = image * red_mask
+    green_channel = image * green_mask
+    blue_channel = image * blue_mask
+
+    # Create the green channel after applying a filter
+    green_channel = np.where(
+        np.logical_or(red_mask == 1, blue_mask == 1),
+        correlate2d(image, filters['G_at_R_and_B'], mode="same", boundary="symm"),
+        green_channel
+    )
+
+
+    # Define masks for extracting pixel values
+    red_row_mask = np.any(red_mask == 1, axis=1)[:, np.newaxis].astype(np.float32)
+    red_col_mask = np.any(red_mask == 1, axis=0)[np.newaxis].astype(np.float32)
+
+    blue_row_mask = np.any(blue_mask == 1, axis=1)[:, np.newaxis].astype(np.float32)
+    blue_col_mask = np.any(blue_mask == 1, axis=0)[np.newaxis].astype(np.float32)
+
+    def update_channel(channel, row_mask, col_mask, filter_key):
+        return np.where(
+            np.logical_and(row_mask == 1, col_mask == 1),
+            correlate2d(image, filters[filter_key], mode="same", boundary="symm"),
+            channel
+        )
+
+# Update the red channel and blue channel
+    red_channel = update_channel(red_channel, red_row_mask, blue_col_mask, 'R_at_GR_and_B_at_GB')
+    red_channel = update_channel(red_channel, blue_row_mask, red_col_mask, 'R_at_GB_and_B_at_GR')
+
+    blue_channel = update_channel(blue_channel, blue_row_mask, red_col_mask, 'R_at_GR_and_B_at_GB')
+    blue_channel = update_channel(blue_channel, red_row_mask, blue_col_mask, 'R_at_GB_and_B_at_GR')
+
+    # Update R channel and B channel again
+    red_channel = update_channel(red_channel, blue_row_mask, blue_col_mask, 'R_at_B_and_B_at_R')
+    blue_channel = update_channel(blue_channel, red_row_mask, red_col_mask, 'R_at_B_and_B_at_R')
+    res[:, :, 0] = red_channel
+    res[:, :, 1] = green_channel
+    res[:, :, 2] = blue_channel
+    return res
\ No newline at end of file