Skip to content
Snippets Groups Projects
reconstruct.py 2.44 KiB
Newer Older
Matthieu Muller's avatar
Matthieu Muller 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).
Matthieu Muller's avatar
Matthieu Muller committed
"""


import numpy as np

from src.forward_model import CFA
Lou-Anne Quellet's avatar
Lou-Anne Quellet committed
from src.methods.quelletl.some_function import interpolation
import cv2
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
from scipy.signal import convolve2d
Matthieu Muller's avatar
Matthieu Muller committed


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.
Matthieu Muller's avatar
Matthieu Muller committed

    Returns:
        np.ndarray: Demosaicked image.
    """
Matthieu Muller's avatar
Matthieu Muller committed
    input_shape = (y.shape[0], y.shape[1], 3)
    op = CFA(cfa, input_shape)
Lou-Anne Quellet's avatar
Lou-Anne Quellet committed
    res = interpolation(y, op)
    
    return res
Matthieu Muller's avatar
Matthieu Muller committed

Lou-Anne Quellet's avatar
Lou-Anne Quellet committed
def quad_to_bayer(y):
    for i in range(1, y.shape[0], 4):
        save = np.copy(y[:,i])
        y[:,i] = y[:,i+1]
        y[:,i+1] = save
    for j in range(1, y.shape[0], 4):
        save = np.copy(y[j,:])
        y[j,:] = y[j+1,:]
        y[j+1,:] = save
    for i in range(1, y.shape[0], 4):
        for j in range(1, y.shape[0], 4):
            save = np.copy(y[i,j])
            y[i,j] = y[i+1,j+1]
            y[i+1,j+1] = save
    return y
Matthieu Muller's avatar
Matthieu Muller committed

Lou-Anne Quellet's avatar
Lou-Anne Quellet committed
ker_bayer_green = np.array([[0, 1, 0], [1, 4, 1], [0, 1, 0]]) / 4
Matthieu Muller's avatar
Matthieu Muller committed

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

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

# 2023
# Authors: Mauro Dalla Mura and Matthieu Muller