Merge branch 'master' into 'master'

Submission Jeanne Lemoine

See merge request !27

src/methods/lemoine/reconstruct.py

+"""A file containing the Malvar et al. reconstruction.
+"""
+import matplotlib.pyplot as plt
+
+import numpy as np
+from src.forward_model import CFA
+from scipy.ndimage.filters import convolve
+
+# Masks definition
+M0 = np.multiply(1/8,[              #  Recover G at R and B locations                                   
+    [ 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]
+] )
+
+M1 = np.multiply(1/8,[              # Recover R at G in R row, B col; Recover B at G in B row, R col                            
+    [ 0, 0, 1/2, 0,  0],
+    [ 0, -1,  0, -1,  0],
+    [-1, 4,  5, 4, -1],
+    [ 0, -1,  0, -1,  0],
+    [ 0, 0, 1/2, 0,  0]
+] )
+
+M2 = M1.T                           # Recover R at G in B row, R col; Recover B at G in R row, B col
+
+M3 = np.multiply(1/8,[              # Recover R at B in B row, B col; Recover B at R in R row, R col
+    [ 0, 0, -3/2, 0,  0],
+    [ 0, 2,  0, 2,  0],
+    [-3/2, 0, 6, 0, -3/2],
+    [ 0, 2,  0, 2,  0],
+    [ 0, 0, -3/2, 0,  0]
+] )
+
+def run_reconstruction(y: np.ndarray, cfa: str) -> np.ndarray:
+    """Performs a high quality linear interpolation of the lost pixels from Malvar et al's (2004) paper.
+
+    Args:
+        op (CFA): CFA operator.
+        y (np.ndarray): Mosaicked image.
+
+    Returns:
+        np.ndarray: Demosaicked image.
+    """
+
+    input_shape = (y.shape[0], y.shape[1], 3)
+    op = CFA(cfa, input_shape)
+
+    # If CFA is Quadratic it needs be transformed into simple Bayer pattern
+    if op.cfa == 'quad_bayer':
+        op.mask, y = adapt_quad(op, y)
+
+    # Applies the adjoint operation to y
+    z = op.adjoint(y)
+
+    # Construct the CFA masks (R,G,B) of z
+    R = np.where(z[:,:,0] != 0., 1, 0) # Red
+    G = np.where(z[:,:,1] != 0., 1, 0) # Green
+    B = np.where(z[:,:,2] != 0., 1, 0) # Blue
+
+    # R, G, B channels of the image to be reconstructed
+    R_demosaic = z[:,:,0]
+    G_demosaic = z[:,:,1]
+    B_demosaic = z[:,:,2]
+
+    # Padding to recover the edges
+    y_pad = np.pad(y, 2, 'constant', constant_values = 0)
+
+    # Recover red and blue columns/rows
+    # Red rows
+    Rr = np.any(R == 1, axis = 1)[None].T * np.ones(R.shape)
+    # Red columns
+    Rc = np.any(R == 1, axis = 0)[None] * np.ones(R.shape)
+    # Blue rows
+    Br = np.any(B == 1, axis = 1)[None].T * np.ones(B.shape)
+    # Blue columns
+    Bc = np.any(B == 1, axis = 0)[None] * np.ones(B.shape)
+
+    # Recover the lost pixels
+    # Red channel
+    R_demosaic = np.where(np.logical_and(Rr == 1, Bc == 1), convolve(y_pad[2:y_pad.shape[0]-2,2:y_pad.shape[1]-2], M1), R_demosaic )
+    R_demosaic = np.where(np.logical_and(Br == 1, Rc == 1), convolve(y_pad[2:y_pad.shape[0]-2,2:y_pad.shape[1]-2], M2), R_demosaic )
+
+    R_demosaic = np.where(np.logical_and(Br == 1, Bc == 1), convolve(y_pad[2:y_pad.shape[0]-2,2:y_pad.shape[1]-2], M3), R_demosaic )
+
+    # Green channel 
+    G_demosaic = np.where(np.logical_or(R == 1, B == 1), convolve(y_pad[2:y_pad.shape[0]-2,2:y_pad.shape[1]-2], M0), G_demosaic )
+
+    # Blue channel
+    B_demosaic = np.where(np.logical_and(Br == 1, Rc == 1), convolve(y_pad[2:y_pad.shape[0]-2,2:y_pad.shape[1]-2], M1), B_demosaic )
+    B_demosaic = np.where(np.logical_and(Rr == 1, Bc == 1), convolve(y_pad[2:y_pad.shape[0]-2,2:y_pad.shape[1]-2], M2), B_demosaic )
+
+    B_demosaic = np.where(np.logical_and(Rr == 1, Rc == 1), convolve(y_pad[2:y_pad.shape[0]-2,2:y_pad.shape[1]-2], M3), B_demosaic )
+
+    # Stack the 3 channels
+    demosaiced = np.stack((R_demosaic, G_demosaic, B_demosaic), axis=2)
+
+    return np.clip(demosaiced, 0, 1, demosaiced)
+
+
+
+def adapt_quad(op: CFA, y: np.ndarray):
+    """Performs an adaptation of the quadratic bayer pattern to a simple bayer pattern.
+
+    Args:
+        op (CFA): CFA operator with Quadratic Bayer pattern.
+        y (np.ndarray): Mosaicked image.
+
+    Returns:
+        bayer (np.ndarray): CFA operator with simple Bayer pattern.
+        new (np.ndarray): Demosaicked image re-arranged.
+    """
+    L, l, d = op.mask.shape
+
+    bayer = op.mask.copy()
+    new = y.copy()
+    
+    # Swap 2 columns every 2 columns
+    for j in range(1, l, 4): 
+        bayer[:,j], bayer[:,j+1] = bayer[:,j+1].copy(), bayer[:,j].copy()
+        new[:,j], new[:,j+1] = new[:,j+1].copy(), new[:,j].copy()
+
+    # Swap 2 lines every 2 lines
+    for i in range(1, L, 4): 
+        bayer[i, :], bayer[i+1,:] = bayer[i+1,:].copy(), bayer[i,:].copy()
+        new[i, :], new[i+1,:] = new[i+1,:].copy(), new[i,:].copy()
+
+    # Swap back some diagonal greens      
+    for i in range(1, L, 4):
+        for j in range(1, 1, 4):
+            bayer[i,j], bayer[i+1,j+1] = op.mask[i+1,j+1].copy(), op.mask[i,j].copy()
+            new[i,j], new[i+1,j+1] = new[i+1,j+1].copy(), new[i,j].copy()
+
+    return bayer, new
\ No newline at end of file