final version

src/methods/anouk_enz/reconstruct.py

@@ -5,9 +5,11 @@ Students can call their functions (declared in others files of src/methods/your_
 
 
 import numpy as np
-
+import cv2 as cv
+from skimage.filters import rank
 from src.forward_model import CFA
-
+from skimage.morphology import disk
+import matplotlib.pyplot as plt
 
 def run_reconstruction(y: np.ndarray, cfa: str) -> np.ndarray:
     """Performs demosaicking on y.
@@ -21,10 +23,223 @@ def run_reconstruction(y: np.ndarray, cfa: str) -> np.ndarray:
     """
     # Performing the reconstruction.
     # TODO
+
     input_shape = (y.shape[0], y.shape[1], 3)
     op = CFA(cfa, input_shape)
 
-    return np.zeros(op.input_shape)
+    z = op.adjoint(y)
+
+
+    ################# MY CODE ############################
+    if op.cfa == 'quad_bayer' : 
+        # Transform into a Bayer-patterned image 
+        y = QuadBayer_to_Bayer(y, op, vizualize=False)
+
+        # Demosaïck the image 
+        res = demosaicking_bayer(y)
+       
+    if op.cfa == 'bayer':
+        res = demosaicking_bayer(y)
+    ####################################################
+
+    return res
+
+def get_color(mask, i, j):
+    """
+        Returns the index of the chanel that contains already a color at a pixel localization. 
+
+        Args :
+            mask (np.ndarray) : mask of the used cfa.
+            i (int): index of the line of the pixel.
+            j (int): index of the column of the pixel.
+
+        Returns 
+            int: the index of the color avaible at the indicated pixel.
+    """
+    if mask[i,j,0] == 1:
+        return 0
+    if mask[i,j,1] == 1:
+        return 1
+    if mask[i,j,2] == 1:
+        return 2
+    return None
+
+def demosaicking_bayer(y):
+    """
+        Performs High-Quality Linear Interpolation on a Bayer-patterned image.
+
+        Args : 
+            y (np.ndarray): Mosaicked image to be reconstructed.
+       
+            Returns 
+                np.ndarray: Demosaicked image. 
+    """
+    input_shape = (y.shape[0], y.shape[1], 3)
+    op = CFA("bayer", input_shape)
+    res = np.empty(op.input_shape)
+
+    for i in range (2, input_shape[0]-2):
+        for j in range (2, input_shape[0]-2):
+            patch = y[i-2:i+3, j-2:j+3]
+            channel = get_color(op.mask, i, j)
+            # If the Red value is avaible 
+            if channel == 0:
+                res[i, j, 0] = y[i,j]
+                # Interpolate the Green value
+                res[i, j, 1] = np.sum(patch*G_at_R)/8
+                #Interpolate the Blue value 
+                res[i, j, 2] = np.sum(patch* B_at_R)/8
+
+            # If the Green value is avaible 
+            if channel == 1 : 
+                res[i, j, 1] = y[i,j]
+                # Interpolation of the Red value
+                if get_color(op.mask, i, j+1) == 0:
+                    res[i,j,0] = np.sum(patch*R_at_G_row)/8
+                else : 
+                    res[i,j,0] = np.sum(patch*R_at_G_column)/8
+                
+                # Interpolation of the Blue value
+                if get_color(op.mask, i, j+1) == 2:
+                    res[i,j,2] = np.sum(patch* B_at_G_row)/8
+                else : 
+                    res[i,j,2] = np.sum(patch * B_at_G_column)/8
+
+            # If the Blue value is avaible 
+            if channel == 2:
+                res[i, j, 2] = y[i,j]
+                # Interpolate the Red value
+                res[i, j, 1] = np.sum(patch* G_at_B)/8
+                #Interpolate the Red value 
+                res[i, j, 0] = np.sum(patch* R_at_B)/8
+    return res
+
+def QuadBayer_to_Bayer(y, op, vizualize=False):
+    """
+    Applies the swapping method to transform a QuadBayer-patterned image into a Bayer-patterned image
+
+    Args : 
+        y (np.ndarray): Mosaicked image to be tranformed 
+        op (CFA) : the CFA object 
+        vizualize (Boolean) : show the evolution of the mask if True
+
+    Returns :
+        (np.ndarray): a bayer-patterned image 
+    """
+    input_shape = (y.shape[0], y.shape[1], 3)
+
+    if vizualize : 
+        fig, axs = plt.subplots(1, 4, figsize=(15,15))
+        axs[0].imshow(op.mask[0:10, 0:10, :])
+        axs[0].set_title('Original mask')
+
+    # Step 1 : Swap 2 columns every 2 columns
+    temp = np.zeros((input_shape[0], 1))
+    temp_mask = np.zeros((input_shape[0], 1, 3))
+    for col in range (1, input_shape[0]-1, 4):
+        temp = np.copy(y[:, col])
+        y[:, col] =  np.copy(y[:, col+1])
+        y[:, col+1] = np.copy(temp)
+
+        if vizualize: 
+            temp_mask = np.copy(op.mask[:, col,:])
+            op.mask[:, col,:] =  np.copy(op.mask[:, col+1,:])
+            op.mask[:, col+1,:] = np.copy(temp_mask)
+    if vizualize:         
+        axs[1].imshow(op.mask[0:10, 0:10, :])
+        axs[1].set_title('Mask after first step')
+        
+    #Step 2 : Swap 2 lines every 2 lines 
+    temp = np.zeros((1, input_shape[1]))
+    temp_mask = np.zeros((1, input_shape[1],3))
+    for line in range (1, input_shape[1], 4):
+        temp = np.copy(y[line, :])
+        y[line, :] =  np.copy(y[line+1, :])
+        y[line+1, :] = np.copy(temp)
+
+        if vizualize: 
+            temp_mask = np.copy(op.mask[line, :, :])
+            op.mask[line, :, :] =  np.copy(op.mask[line+1, :, :])
+            op.mask[line+1, :, :] = np.copy(temp_mask)
+    if vizualize: 
+        axs[2].imshow(op.mask[0:10, 0:10, :])
+        axs[2].set_title('Mask after second step')
+
+    # 3: Swap back some diagonal greens
+    temp_mask = np.zeros((1,1,3))
+    for i in range(0, input_shape[0], 4):
+        for j in range(2, input_shape[1], 4):
+            temp = y[i, j]
+            y[i, j] = y[i+1, j-1]
+            y[i+1, j-1] = temp
+            if vizualize: 
+                temp_mask = op.mask[i, j, :]
+                op.mask[i, j, :] = np.copy(op.mask[i+1, j-1, :])
+                op.mask[i+1, j-1, :] = np.copy(temp_mask)
+    if vizualize: 
+        axs[3].imshow(op.mask[0:10, 0:10, :])
+        axs[3].set_title('Mask after third step')
+        plt.tight_layout()
+        plt.show()
+
+    return y
+
+
+# Defining the mask
+
+# Interpolate the Green 
+# Green at Red 
+G_at_R = [[ 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]]
+# Green at Blue
+G_at_B = [[ 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]]
+
+#Interpolate the Red 
+# R at G in R row, B column
+R_at_G_row = [[ 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]] 
+# R at G in B row, R column
+R_at_G_column = [[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]]
+# Red at Blue
+R_at_B = [[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]]
+
+#Intelropation for Blue 
+# Blue at Green in B row, R column
+B_at_G_row = [[ 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]] 
+
+B_at_G_column = [[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]]
+
+B_at_R = [[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]]
 
 
 ####