the method is working and report upload

reconstruct.py → src/methods/TheoLafond/reconstruct.py

@@ -13,10 +13,15 @@ from src.forward_model import CFA
 
 
 def regression_filter_func(chans):
-    """_summary_
+    """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()
@@ -27,7 +32,7 @@ def regression_filter_func(chans):
     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=""):
-    """Computes the regression filter on the input channel between the red with the green as a guide and the blue with the green as a guide.
+    """Apply func on each footprint of a.
     """
 
     # Suppress RankWarning
@@ -51,6 +56,36 @@ def personal_generic_filter(a, func, footprint, output, msg=""):
                 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.
@@ -63,7 +98,7 @@ def run_reconstruction(y: np.ndarray, cfa: str) -> np.ndarray:
         np.ndarray: Demosaicked image.
     """
     # Performing the reconstruction.
-    # TODO
+
     input_shape = (y.shape[0], y.shape[1], 3)
     op = CFA(cfa, input_shape)
     adjoint = op.adjoint(y)
@@ -86,85 +121,57 @@ def run_reconstruction(y: np.ndarray, cfa: str) -> np.ndarray:
                 continue
             bv[:, j, i] = np.interp(np.arange(adjoint.shape[0]), non_zero, adjoint[:, j, i][non_zero])
 
-    
-
     # Residuals interpolation
-    K = 1
-    for k in range(K):
-        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 #"+str(k+1)+"/"+str(K))
-        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[:,:,2]*bh[:,:,1] + regh[:,:,3]
-        ch[:,:,2] = regh[:,:,4]*bh[:,:,2] + regh[:,:,5] + regh[:,:,6]*bh[:,:,2] + regh[:,:,7]
-        dh = adjoint - ch
-        dh[adjoint==0] = 0
-        # 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
-                ch[j, :, i] += np.interp(np.arange(adjoint.shape[1]), non_zero, adjoint[j, :, i][non_zero])
-        bh = ch.copy()
-
-        # 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 #"+str(k+1))
-        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[:,:,2]*bv[:,:,1] + regv[:,:,3]
-        cv[:,:,2] = regv[:,:,4]*bv[:,:,2] + regv[:,:,5] + regv[:,:,6]*bv[:,:,2] + regv[:,:,7]
-        dv = adjoint - cv
-        dv[adjoint==0] = 0
-        # 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
-                cv[:, j, i] += np.interp(np.arange(adjoint.shape[0]), non_zero, adjoint[:, j, i][non_zero])
-        bv = cv.copy()
-
-    combo = np.zeros(adjoint.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
-    
+    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()
 
-    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
+    return combine_directions(bh,bv)
 
     
 
@@ -172,16 +179,16 @@ def run_reconstruction(y: np.ndarray, cfa: str) -> np.ndarray:
 if __name__ == "__main__":
     from src.utils import load_image, save_image, psnr, ssim
     from src.forward_model import CFA
-    image_path = 'images/img_1.png'
+    image_path = 'images/img_4.png'
 
-    img = load_image(image_path)[:256,:256,:]
+    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))
+    print('PSNR:', psnr(img, res))
+    print('SSIM:', ssim(img, res))
     save_image('output/bh.png', res)