Merge branch 'master' into 'master'

TheoLafond

See merge request !43

src/methods/TheoLafond/reconstruct.py

+"""The main file for the reconstruction.
+This file should NOT be modified except the body of the 'run_reconstruction' function.
+Students can call their functions (declared in others files of src/methods/your_name).
+"""
+
+
+import numpy as np
+from scipy.ndimage import generic_filter, convolve
+import warnings
+from tqdm import tqdm
+from src.forward_model import CFA
+
+
+
+def regression_filter_func(chans):
+    """computes the linear regression coefficients between channel :
+    red and green
+    blue and green
+    green and red
+    green and blue
+
+
+    Returns:
+        a is high order coef and b is 0 order.
+        _type_: a_red, b_red, a_blue, b_blue, a_green_red, b_green_red, a_green_blue, b_green_blue
+    """
+    red = chans[:,:,0].flatten()
+    green = chans[:,:,1].flatten()
+    blue = chans[:,:,2].flatten()
+    nonzeros_red = np.nonzero(red)
+    nonzeros_blue = np.nonzero(blue)
+    return np.concatenate((np.polyfit(red[nonzeros_red], green[nonzeros_red], deg=1), np.polyfit(blue[nonzeros_blue], green[nonzeros_blue], deg=1), np.polyfit(green[nonzeros_red], red[nonzeros_red], deg=1), np.polyfit(green[nonzeros_blue], blue[nonzeros_blue], deg=1)), axis=0)
+
+def personal_generic_filter(a, func, footprint, output, msg=""):
+    """Apply func on each footprint of a.
+    """
+
+    # Suppress RankWarning
+    with warnings.catch_warnings():
+        warnings.simplefilter('ignore', np.RankWarning)
+        for i in tqdm(range(a.shape[0]), desc=msg):
+            for j in range(a.shape[1]):
+                fen = np.ones((footprint.shape[0], footprint.shape[1], a.shape[2]))
+                if i+footprint.shape[0] > a.shape[0]:
+                    i_end = a.shape[0]
+                    i_start = i_end - footprint.shape[0]
+                else:
+                    i_start = i
+                    i_end = i + footprint.shape[0]
+                if j+footprint.shape[1] > a.shape[1]:
+                    j_end = a.shape[1]
+                    j_start = j_end - footprint.shape[1]
+                else:
+                    j_start = j
+                    j_end = j + footprint.shape[1]
+                fen[:i_end-i_start, :j_end-j_start] = a[i_start:i_end, j_start:j_end]
+                output[i, j, :] = func(a[i_start:i_end, j_start:j_end])
+
+def combine_directions(bh,bv):
+    combo = np.zeros(bh.shape)
+    combo[:,:,0] = bh[:,:,0] + bv[:,:,0]
+    combo[:,:,0][(bh[:,:,0]!=0) * (bv[:,:,0]!=0)] = (bh[:,:,0][(bh[:,:,0]!=0) * (bv[:,:,0]!=0)] + bv[:,:,0][(bh[:,:,0]!=0) * (bv[:,:,0]!=0)])/2
+    combo[:,:,0][(bh[:,:,0]==0) * (bv[:,:,0]==0)] = 0
+
+    combo[:,:,1] = bh[:,:,1]/2 + bv[:,:,1]/2
+    combo[:,:,2] = bh[:,:,2] + bv[:,:,2]
+    combo[:,:,2][(bh[:,:,2]!=0) * (bv[:,:,2]!=0)] = (bh[:,:,2][(bh[:,:,2]!=0) * (bv[:,:,2]!=0)] + bv[:,:,2][(bh[:,:,2]!=0) * (bv[:,:,2]!=0)])/2
+    combo[:,:,2][(bh[:,:,2]==0) * (bv[:,:,2]==0)] = 0
+    
+
+    for i in range(combo.shape[0]):
+        for j in range(combo.shape[1]):
+            moy = []
+            if i != 0 :
+                moy.append(combo[i-1,j,:])
+            if i != combo.shape[0]-1 :
+                moy.append(combo[i+1,j,:])
+            if j != 0 :
+                moy.append(combo[i,j-1,:])
+            if j != combo.shape[1]-1 :
+                moy.append(combo[i,j+1,:])
+            moy = np.stack(moy).mean(axis=0)
+            if combo[i,j,0] == 0:
+                combo[i,j,0] = moy[0]
+            if combo[i,j,2] == 0:
+                combo[i,j,2] = moy[2]
+    return combo
+
+
+def run_reconstruction(y: np.ndarray, cfa: str) -> np.ndarray:
+    """Performs demosaicking on y.
+
+    Args:
+        y (np.ndarray): Mosaicked image to be reconstructed.
+        cfa (str): Name of the CFA. Can be bayer or quad_bayer.
+
+    Returns:
+        np.ndarray: Demosaicked image.
+    """
+    # Performing the reconstruction.
+
+    input_shape = (y.shape[0], y.shape[1], 3)
+    op = CFA(cfa, input_shape)
+    adjoint = op.adjoint(y)
+
+    bh = np.zeros(adjoint.shape)
+    # Horizontal interpolation
+    for i in range(3):
+        for j in range(adjoint.shape[0]):
+            non_zero = np.nonzero(adjoint[j, :, i])[0]
+            if len(non_zero) == 0:
+                continue
+            bh[j, :, i] = np.interp(np.arange(adjoint.shape[1]), non_zero, adjoint[j, :, i][non_zero])
+
+    bv = np.zeros(adjoint.shape)
+    # Vertical interpolation
+    for i in range(3):
+        for j in range(adjoint.shape[1]):
+            non_zero = np.nonzero(adjoint[:, j, i])[0]
+            if len(non_zero) == 0:
+                continue
+            bv[:, j, i] = np.interp(np.arange(adjoint.shape[0]), non_zero, adjoint[:, j, i][non_zero])
+
+    # Residuals interpolation
+    M = 3
+    N = 5
+    kernel_hor = np.ones((M, N))/(M*N)
+    kernel_ver = np.ones((N, M))/(M*N)
+
+    # Horizontal Regression filtering
+    regh = np.zeros((bh.shape[0], bh.shape[1], 8))
+    personal_generic_filter(bh, regression_filter_func, kernel_hor, regh,msg="Regression filtering horizontal")
+    ch = np.zeros(bh.shape)
+    for i in range(regh.shape[-1]):
+        regh[:, :, i] = convolve(regh[:, :, i], kernel_hor)
+    ch[:,:,0] = regh[:,:,0]*bh[:,:,0] + regh[:,:,1]
+    ch[:,:,1] = (regh[:,:,4]*bh[:,:,1] + regh[:,:,5] + regh[:,:,6]*bh[:,:,1] + regh[:,:,7])/2
+    ch[:,:,2] = regh[:,:,2]*bh[:,:,2] + regh[:,:,3]
+    # return ch[:,:,1]
+    dh = ch - adjoint
+    dh[adjoint==0] = 0
+    # interpolation
+    for i in range(3):
+        for j in range(adjoint.shape[0]):
+            non_zero = np.nonzero(dh[j, :, i])[0]
+            if len(non_zero) == 0:
+                continue
+            ch[j, :, i] -= np.interp(np.arange(adjoint.shape[1]), non_zero, dh[j, :, i][non_zero])
+    ch[bh == 0] = 0
+    bh = ch.copy()
+    # return bh[:,:,0]
+
+    # Vertical Regression filtering
+    regv = np.zeros((bv.shape[0], bv.shape[1], 8))
+    personal_generic_filter(bv, regression_filter_func, kernel_ver, regv,msg="Regression filtering vertical")
+    cv = np.zeros(bv.shape)
+    for i in range(regv.shape[-1]):
+        regv[:, :, i] = convolve(regv[:, :, i], kernel_ver)
+    cv[:,:,0] = regv[:,:,0]*bv[:,:,0] + regv[:,:,1]
+    cv[:,:,1] = (regv[:,:,4]*bv[:,:,1] + regv[:,:,5] + regv[:,:,6]*bv[:,:,1] + regv[:,:,7])/2
+    cv[:,:,2] = regv[:,:,2]*bv[:,:,2] + regv[:,:,3]
+    dv = cv - adjoint
+    dv[adjoint==0] = 0
+    # interpolation
+    for i in range(3):
+        for j in range(adjoint.shape[1]):
+            non_zero = np.nonzero(dv[:, j, i])[0]
+            if len(non_zero) == 0:
+                continue
+            cv[:, j, i] -= np.interp(np.arange(adjoint.shape[0]), non_zero, dv[:, j, i][non_zero])
+    cv[bv == 0] = 0
+    bv = cv.copy()
+
+    return combine_directions(bh,bv)
+
+    
+
+
+if __name__ == "__main__":
+    from src.utils import load_image, save_image, psnr, ssim
+    from src.forward_model import CFA
+    image_path = 'images/img_4.png'
+
+    img = load_image(image_path)
+
+    cfa_name = 'bayer' # bayer or quad_bayer
+    op = CFA(cfa_name, img.shape)
+    y = op.direct(img)
+    res = run_reconstruction(y, cfa_name)
+    print('PSNR:', psnr(img, res))
+    print('SSIM:', ssim(img, res))
+    save_image('output/bh.png', res)
+
+
+
+####
+####
+####
+
+####      ####                ####        #############
+####      ######              ####      ##################
+####      ########            ####      ####################
+####      ##########          ####      ####        ########
+####      ############        ####      ####            ####
+####      ####  ########      ####      ####            ####
+####      ####    ########    ####      ####            ####
+####      ####      ########  ####      ####            ####
+####      ####  ##    ######  ####      ####          ######
+####      ####  ####      ##  ####      ####    ############
+####      ####  ######        ####      ####    ##########
+####      ####  ##########    ####      ####    ########
+####      ####      ########  ####      ####
+####      ####        ############      ####
+####      ####          ##########      ####
+####      ####            ########      ####
+####      ####              ######      ####
+
+# 2023
+# Authors: Mauro Dalla Mura and Matthieu Muller