| """ |
| Definition of the FFDNet model and its custom layers |
| |
| Copyright (C) 2018, Matias Tassano <matias.tassano@parisdescartes.fr> |
| |
| This program is free software: you can use, modify and/or |
| redistribute it under the terms of the GNU General Public |
| License as published by the Free Software Foundation, either |
| version 3 of the License, or (at your option) any later |
| version. You should have received a copy of this license along |
| this program. If not, see <http://www.gnu.org/licenses/>. |
| """ |
| import torch.nn as nn |
| from torch.autograd import Variable |
| import denoising.functions as functions |
| |
| class UpSampleFeatures(nn.Module): |
| r"""Implements the last layer of FFDNet |
| """ |
| def __init__(self): |
| super(UpSampleFeatures, self).__init__() |
| def forward(self, x): |
| return functions.upsamplefeatures(x) |
|
|
| class IntermediateDnCNN(nn.Module): |
| r"""Implements the middel part of the FFDNet architecture, which |
| is basically a DnCNN net |
| """ |
| def __init__(self, input_features, middle_features, num_conv_layers): |
| super(IntermediateDnCNN, self).__init__() |
| self.kernel_size = 3 |
| self.padding = 1 |
| self.input_features = input_features |
| self.num_conv_layers = num_conv_layers |
| self.middle_features = middle_features |
| if self.input_features == 5: |
| self.output_features = 4 |
| elif self.input_features == 15: |
| self.output_features = 12 |
| else: |
| raise Exception('Invalid number of input features') |
|
|
| layers = [] |
| layers.append(nn.Conv2d(in_channels=self.input_features,\ |
| out_channels=self.middle_features,\ |
| kernel_size=self.kernel_size,\ |
| padding=self.padding,\ |
| bias=False)) |
| layers.append(nn.ReLU(inplace=True)) |
| for _ in range(self.num_conv_layers-2): |
| layers.append(nn.Conv2d(in_channels=self.middle_features,\ |
| out_channels=self.middle_features,\ |
| kernel_size=self.kernel_size,\ |
| padding=self.padding,\ |
| bias=False)) |
| layers.append(nn.BatchNorm2d(self.middle_features)) |
| layers.append(nn.ReLU(inplace=True)) |
| layers.append(nn.Conv2d(in_channels=self.middle_features,\ |
| out_channels=self.output_features,\ |
| kernel_size=self.kernel_size,\ |
| padding=self.padding,\ |
| bias=False)) |
| self.itermediate_dncnn = nn.Sequential(*layers) |
| def forward(self, x): |
| out = self.itermediate_dncnn(x) |
| return out |
|
|
| class FFDNet(nn.Module): |
| r"""Implements the FFDNet architecture |
| """ |
| def __init__(self, num_input_channels): |
| super(FFDNet, self).__init__() |
| self.num_input_channels = num_input_channels |
| if self.num_input_channels == 1: |
| |
| self.num_feature_maps = 64 |
| self.num_conv_layers = 15 |
| self.downsampled_channels = 5 |
| self.output_features = 4 |
| elif self.num_input_channels == 3: |
| |
| self.num_feature_maps = 96 |
| self.num_conv_layers = 12 |
| self.downsampled_channels = 15 |
| self.output_features = 12 |
| else: |
| raise Exception('Invalid number of input features') |
|
|
| self.intermediate_dncnn = IntermediateDnCNN(\ |
| input_features=self.downsampled_channels,\ |
| middle_features=self.num_feature_maps,\ |
| num_conv_layers=self.num_conv_layers) |
| self.upsamplefeatures = UpSampleFeatures() |
|
|
| def forward(self, x, noise_sigma): |
| concat_noise_x = functions.concatenate_input_noise_map(x.data, noise_sigma.data) |
| concat_noise_x = Variable(concat_noise_x) |
| h_dncnn = self.intermediate_dncnn(concat_noise_x) |
| pred_noise = self.upsamplefeatures(h_dncnn) |
| return pred_noise |
|
|
