Skip to content
Snippets Groups Projects
Select Git revision
  • master default protected
1 result
Blame Permalink
forward_model.py 3.21 KiB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123 






"""A file containing the forward operator.
This file should NOT be modified.
"""


import numpy as np

from src.checks import check_cfa, check_rgb


class CFA():
    def __init__(self, cfa: str, input_shape: tuple) -> None:
        """Constructor of the forward operator's class.

        Args:
            cfa (str): Name of the pattern. Either bayer or quad_bayer.
            input_shape (tuple): Shape of the input images of the operator.
        """
        check_cfa(cfa)
    
        self.cfa = cfa
        self.input_shape = input_shape
        self.output_shape = input_shape[:-1]

        if self.cfa == 'bayer':
            self.mask = get_bayer_mask(input_shape)

        elif self.cfa == 'quad_bayer':
            self.mask = get_quad_bayer_mask(input_shape)


    def direct(self, x: np.ndarray) -> np.ndarray:
        """Applies the CFA operation to the image x.

        Args:
            x (np.ndarray): Input image.

        Returns:
            np.ndarray: Output image.
        """
        check_rgb(x)
    
        return np.sum(x * self.mask, axis=2)


    def adjoint(self, y: np.ndarray) -> np.ndarray:
        """Applies the adjoint of the CFA operation.

        Args:
            y (np.ndarray): Input image.

        Returns:
            np.ndarray: Output image.
        """
        return self.mask * y[..., np.newaxis]


def get_bayer_mask(input_shape: tuple) -> np.ndarray:
    """Return the mask of the Bayer CFA.

    Args:
        input_shape (tuple): Shape of the mask.

    Returns:
        np.ndarray: Mask.
    """
    res = np.kron(np.ones((input_shape[0], input_shape[1], 1)), [0, 1, 0])

    res[::2, 1::2] = [1, 0, 0]
    res[1::2, ::2] = [0, 0, 1]

    return res


def get_quad_bayer_mask(input_shape: tuple) -> np.ndarray:
    """Return the mask of the quad_bayer CFA.

    Args:
        input_shape (tuple): Shape of the mask.

    Returns:
        np.ndarray: Mask.
    """
    res = np.kron(np.ones((input_shape[0], input_shape[1], 1)), [0, 1, 0])

    res[::4, 2::4] = [1, 0, 0]
    res[::4, 3::4] = [1, 0, 0]
    res[1::4, 2::4] = [1, 0, 0]
    res[1::4, 3::4] = [1, 0, 0]

    res[2::4, ::4] = [0, 0, 1]
    res[2::4, 1::4] = [0, 0, 1]
    res[3::4, ::4] = [0, 0, 1]
    res[3::4, 1::4] = [0, 0, 1]

    return res


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

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

# 2023
# Authors: Mauro Dalla Mura and Matthieu Muller