|
""" |
|
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 |
|
|
