Skip to content
Snippets Groups Projects
function.py 1.68 KiB
Newer Older
Patrick Dai's avatar
Patrick Dai committed
import numpy as np
from scipy.signal import convolve2d

from src.forward_model import CFA


def naive_interpolation(op: CFA, y: np.ndarray) -> np.ndarray:
    """Performs a simple interpolation of the lost pixels.

    Args:
        op (CFA): CFA operator.
        y (np.ndarray): Mosaicked image.

    Returns:
        np.ndarray: Demosaicked image.
    """
    

    res = np.empty(op.input_shape)
    z = op.adjoint(y)

    res[:, :, 0] = convolve2d(z[:, :, 0], ker_bayer_red_blue, mode='same')
    res[:, :, 1] = convolve2d(z[:, :, 1], ker_bayer_green, mode='same')
    res[:, :, 2] = convolve2d(z[:, :, 2], ker_bayer_red_blue, mode='same')
    
    return res




def Spectral_difference(op,y):

    z = op.adjoint(y)
    y_hat  = naive_interpolation(op,y)
    res = np.empty(op.input_shape)
    for l in range(3):
        res[:,:,0]+=z[:,:,l]+convolve2d(z[:,:,0]-y_hat[:,:,l]*op.mask[:,:,0], ker_bayer_red_blue, mode='same')*op.mask[:,:,l]
        res[:,:,1]+=z[:,:,l]+convolve2d(z[:,:,1]-y_hat[:,:,l]*op.mask[:,:,1], ker_bayer_green, mode='same')*op.mask[:,:,l]
        res[:,:,2]+=z[:,:,l]+convolve2d(z[:,:,2]-y_hat[:,:,l]*op.mask[:,:,2], ker_bayer_red_blue, mode='same')*op.mask[:,:,l]

    return res

def quad_bayer_to_bayer(y_quad):
    y_bayer=np.copy(y_quad)

    for i in range(1,y_quad.shape[0],4):
        temp = np.copy(y_bayer[i,:])
        y_bayer[i,:]=y_bayer[i+1,:]
        y_bayer[i+1,:]=temp
    
    for j in range(1,y_quad.shape[1],4):
        temp = np.copy(y_bayer[:,j])
        y_bayer[:,j]=y_bayer[:,j+1]
        y_bayer[:,j+1]=temp



    return y_bayer




ker_bayer_red_blue = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]]) / 4

ker_bayer_green = np.array([[0, 1, 0], [1, 4, 1], [0, 1, 0]]) / 4