Skip to content

dncnn.py

ConvBlock

Bases: Sequential

Convolution block: conv => BN => act.

Source code in src/autoden/models/dncnn.py 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class ConvBlock(nn.Sequential):
    """Convolution block: conv => BN => act."""

    def __init__(self, in_ch: int, out_ch: int, kernel_size: int, pad_mode: str = "replicate", last_block: bool = False):
        pad_size = (kernel_size - 1) // 2
        if last_block:
            post_conv = []
        else:
            post_conv = [nn.BatchNorm2d(out_ch), nn.LeakyReLU(0.2, inplace=True)]
        super().__init__(
            nn.Conv2d(in_ch, out_ch, kernel_size=kernel_size, padding=pad_size, padding_mode=pad_mode, bias=False),
            *post_conv,
        )

DnCNN

Bases: Sequential

Implementation of the DnCNN architecture from [1].

[1] Zhang, et al., "Beyond a Gaussian denoiser: Residual learning of deep CNN for image denoising," IEEE Trans. on Image Processing, 2017.

Source code in src/autoden/models/dncnn.py 
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
class DnCNN(nn.Sequential):
    """Implementation of the DnCNN architecture from [1].

    [1] Zhang, et al., "Beyond a Gaussian denoiser: Residual learning of deep CNN
        for image denoising," IEEE Trans. on Image Processing, 2017.
    """

    def __init__(
        self,
        n_channels_in: int,
        n_channels_out: int,
        n_layers: int = 20,
        n_features: int = 32,
        kernel_size: int = 3,
        pad_mode: str = "replicate",
        device: str = "cuda" if pt.cuda.is_available() else "cpu",
    ):
        # From zhang-2017-beyon-gauss-denois:
        #
        #  Thus, for Gaussian denoising with a certain noise level, we
        #  set the receptive field size of DnCNN to 35 × 35 with the
        #  corresponding depth of 17. For other general image denoising
        #  tasks, we adopt a larger receptive field and set the depth
        #  to be 20.
        #
        # Hence, we set the standard depth to 20.
        layers = [
            ConvBlock(
                n_channels_in if i_l == 0 else n_features,
                n_channels_out if i_l == (n_layers - 1) else n_features,
                kernel_size=kernel_size,
                pad_mode=pad_mode,
                last_block=(i_l == (n_layers - 1)),
            )
            for i_l in range(n_layers)
        ]

        super().__init__(*layers)
        self.n_ch_in = n_channels_in
        self.n_ch_out = n_channels_out
        self.n_layers = n_layers
        self.device = device

        self.to(self.device)