Skip to content
Snippets Groups Projects
reconstruct.py 2.71 KiB
Newer Older
Simon's avatar
Simon committed
"""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 src.forward_model import CFA
from scipy.signal import medfilt2d


def freeman_median_demosaicking(op: CFA, y: np.ndarray) -> np.ndarray:
    """Performs the Freeman's method with median for demosaicking.

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

    Returns:
        np.ndarray: Demosaicked image.
    """
    z = op.adjoint(y)
    res = np.empty(op.input_shape)

    mask_r = np.zeros_like(z[:, :, 0], dtype = float)
    mask_g = np.zeros_like(z[:, :, 1], dtype = float)
    mask_b = np.zeros_like(z[:, :, 2], dtype = float)

    D_rg = z[:, :, 0] - z[:, :, 1]
    M1 = medfilt2d(D_rg, kernel_size=5)

    D_gb = z[:, :, 1] - z[:, :, 2]
    M2 = medfilt2d(D_gb, kernel_size=5)

    D_rb = z[:, :, 0] - z[:, :, 2]
    M3 = medfilt2d(D_rb, kernel_size=5)

    res[:, :, 0] = z[:, :, 0] + (M1 * mask_g + z[:, :, 1]) + (M3 * mask_b + z[:, :, 2])
    res[:, :, 1] = z[:, :, 1] + (z[:, :, 0] - M1 * mask_r) + (M2 * mask_b + z[:, :, 2])
    res[:, :, 2] = z[:, :, 2] + (z[:, :, 1] - M2 * mask_g) + (z[:, :, 0] - M3 * mask_r)

    return res



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.
    """
    input_shape = (y.shape[0], y.shape[1], 3)
    op = CFA(cfa, input_shape)

    res = freeman_median_demosaicking(op, y)

    return res


####
####
####

####      ####                ####        #############
####      ######              ####      ##################
####      ########            ####      ####################
####      ##########          ####      ####        ########
####      ############        ####      ####            ####
####      ####  ########      ####      ####            ####
####      ####    ########    ####      ####            ####
####      ####      ########  ####      ####            ####
####      ####  ##    ######  ####      ####          ######
####      ####  ####      ##  ####      ####    ############
####      ####  ######        ####      ####    ##########
####      ####  ##########    ####      ####    ########
####      ####      ########  ####      ####
####      ####        ############      ####
####      ####          ##########      ####
####      ####            ########      ####
####      ####              ######      ####

# 2023
# Authors: Mauro Dalla Mura and Matthieu Muller