1220 lines
51 KiB
Python
1220 lines
51 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from torch.nn import init
|
||
import functools
|
||
from torch.optim import lr_scheduler
|
||
|
||
import torch.nn.functional as F
|
||
|
||
import numpy as np
|
||
|
||
|
||
###############################################################################
|
||
# Helper Functions
|
||
###############################################################################
|
||
|
||
|
||
# class Identity(nn.Module):
|
||
# def forward(self, x):
|
||
# return x
|
||
|
||
|
||
def get_norm_layer(norm_type='instance'):
|
||
"""Return a normalization layer
|
||
|
||
Parameters:
|
||
norm_type (str) -- the name of the normalization layer: batch | instance | none
|
||
|
||
For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
|
||
For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
|
||
"""
|
||
if norm_type == 'batch':
|
||
norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
|
||
elif norm_type == 'instance':
|
||
norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
|
||
elif norm_type == 'none':
|
||
norm_layer = lambda x: Identity()
|
||
else:
|
||
raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
|
||
return norm_layer
|
||
|
||
|
||
def get_scheduler(optimizer, opt):
|
||
"""Return a learning rate scheduler
|
||
|
||
Parameters:
|
||
optimizer -- the optimizer of the network
|
||
opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．
|
||
opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine
|
||
|
||
For 'linear', we keep the same learning rate for the first <opt.niter> epochs
|
||
and linearly decay the rate to zero over the next <opt.niter_decay> epochs.
|
||
For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.
|
||
See https://pytorch.org/docs/stable/optim.html for more details.
|
||
"""
|
||
if opt.lr_policy == 'linear':
|
||
def lambda_rule(epoch):
|
||
lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.niter) / float(opt.niter_decay + 1)
|
||
return lr_l
|
||
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
|
||
elif opt.lr_policy == 'step':
|
||
scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
|
||
elif opt.lr_policy == 'plateau':
|
||
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
|
||
elif opt.lr_policy == 'cosine':
|
||
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.niter, eta_min=0)
|
||
else:
|
||
return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
|
||
return scheduler
|
||
|
||
|
||
def init_weights(net, init_type='normal', init_gain=0.02):
|
||
"""Initialize network weights.
|
||
|
||
Parameters:
|
||
net (network) -- network to be initialized
|
||
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
|
||
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
|
||
|
||
We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
|
||
work better for some applications. Feel free to try yourself.
|
||
"""
|
||
def init_func(m): # define the initialization function
|
||
classname = m.__class__.__name__
|
||
if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
|
||
if init_type == 'normal':
|
||
init.normal_(m.weight.data, 0.0, init_gain)
|
||
elif init_type == 'xavier':
|
||
init.xavier_normal_(m.weight.data, gain=init_gain)
|
||
elif init_type == 'kaiming':
|
||
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
|
||
elif init_type == 'orthogonal':
|
||
init.orthogonal_(m.weight.data, gain=init_gain)
|
||
else:
|
||
raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
|
||
if hasattr(m, 'bias') and m.bias is not None:
|
||
init.constant_(m.bias.data, 0.0)
|
||
elif classname.find('BatchNorm2d') != -1: # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
|
||
init.normal_(m.weight.data, 1.0, init_gain)
|
||
init.constant_(m.bias.data, 0.0)
|
||
|
||
print('initialize network with %s' % init_type)
|
||
net.apply(init_func) # apply the initialization function <init_func>
|
||
|
||
|
||
def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]):
|
||
"""Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
|
||
Parameters:
|
||
net (network) -- the network to be initialized
|
||
init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
|
||
gain (float) -- scaling factor for normal, xavier and orthogonal.
|
||
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
|
||
|
||
Return an initialized network.
|
||
"""
|
||
if len(gpu_ids) > 0:
|
||
assert(torch.cuda.is_available())
|
||
net.to(gpu_ids[0])
|
||
net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs
|
||
init_weights(net, init_type, init_gain=init_gain)
|
||
return net
|
||
|
||
|
||
def define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]):
|
||
"""Create a generator
|
||
|
||
Parameters:
|
||
input_nc (int) -- the number of channels in input images
|
||
output_nc (int) -- the number of channels in output images
|
||
ngf (int) -- the number of filters in the last conv layer
|
||
netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128
|
||
norm (str) -- the name of normalization layers used in the network: batch | instance | none
|
||
use_dropout (bool) -- if use dropout layers.
|
||
init_type (str) -- the name of our initialization method.
|
||
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
|
||
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
|
||
|
||
Returns a generator
|
||
|
||
Our current implementation provides two types of generators:
|
||
U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)
|
||
The original U-Net paper: https://arxiv.org/abs/1505.04597
|
||
|
||
Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)
|
||
Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.
|
||
We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).
|
||
|
||
|
||
The generator has been initialized by <init_net>. It uses RELU for non-linearity.
|
||
"""
|
||
net = None
|
||
norm_layer = get_norm_layer(norm_type=norm)
|
||
|
||
if netG == 'resnet_9blocks':
|
||
net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
|
||
elif netG == 'resnet_6blocks':
|
||
net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)
|
||
elif netG == 'resnet_9blocks_inter':
|
||
net = ResnetGWithIntermediate(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
|
||
elif netG == 'resnet_9blocks_dehaze':
|
||
net = ResnetDehazeGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
|
||
elif netG == 'unet_128':
|
||
net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
|
||
elif netG == 'unet_256':
|
||
net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
|
||
elif netG == 'unet_trans_256':
|
||
net = UnetTransGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout, r=10, eps=1e-3)
|
||
elif netG == 'haze_refine_2':
|
||
net = HazeRefiner(input_nc, output_nc, 2)
|
||
elif netG == 'haze_refine_10':
|
||
net = HazeRefiner(input_nc, output_nc, 10)
|
||
else:
|
||
raise NotImplementedError('Generator model name [%s] is not recognized' % netG)
|
||
return init_net(net, init_type, init_gain, gpu_ids)
|
||
|
||
|
||
def define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, gpu_ids=[]):
|
||
"""Create a discriminator
|
||
|
||
Parameters:
|
||
input_nc (int) -- the number of channels in input images
|
||
ndf (int) -- the number of filters in the first conv layer
|
||
netD (str) -- the architecture's name: basic | n_layers | pixel
|
||
n_layers_D (int) -- the number of conv layers in the discriminator; effective when netD=='n_layers'
|
||
norm (str) -- the type of normalization layers used in the network.
|
||
init_type (str) -- the name of the initialization method.
|
||
init_gain (float) -- scaling factor for normal, xavier and orthogonal.
|
||
gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
|
||
|
||
Returns a discriminator
|
||
|
||
Our current implementation provides three types of discriminators:
|
||
[basic]: 'PatchGAN' classifier described in the original pix2pix paper.
|
||
It can classify whether 70×70 overlapping patches are real or fake.
|
||
Such a patch-level discriminator architecture has fewer parameters
|
||
than a full-image discriminator and can work on arbitrarily-sized images
|
||
in a fully convolutional fashion.
|
||
|
||
[n_layers]: With this mode, you cna specify the number of conv layers in the discriminator
|
||
with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)
|
||
|
||
[pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.
|
||
It encourages greater color diversity but has no effect on spatial statistics.
|
||
|
||
The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.
|
||
"""
|
||
net = None
|
||
norm_layer = get_norm_layer(norm_type=norm)
|
||
|
||
if netD == 'basic': # default PatchGAN classifier
|
||
net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)
|
||
elif netD == 'n_layers': # more options
|
||
net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)
|
||
elif netD == 'pixel': # classify if each pixel is real or fake
|
||
net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)
|
||
else:
|
||
raise NotImplementedError('Discriminator model name [%s] is not recognized' % netD)
|
||
return init_net(net, init_type, init_gain, gpu_ids)
|
||
|
||
|
||
##############################################################################
|
||
# Classes
|
||
##############################################################################
|
||
class TVLoss(nn.Module):
|
||
'''
|
||
Define Total Variance Loss for images
|
||
which is used for smoothness regularization
|
||
'''
|
||
|
||
def __init__(self):
|
||
super(TVLoss, self).__init__()
|
||
|
||
def __call__(self, input):
|
||
# Tensor with shape (n_Batch, C, H, W)
|
||
origin = input[:, :, :-1, :-1]
|
||
right = input[:, :, :-1, 1:]
|
||
down = input[:, :, 1:, :-1]
|
||
|
||
tv = torch.mean(torch.abs(origin-right)) + torch.mean(torch.abs(origin-down))
|
||
return tv * 0.5
|
||
|
||
|
||
class ARefiner(nn.Module):
|
||
"""docstring for ARefiner"""
|
||
def __init__(self, ngf, n_downsampling):
|
||
super(ARefiner, self).__init__()
|
||
self.n_downsampling = n_downsampling
|
||
self.ngf = ngf
|
||
|
||
self.global_pooling = torch.nn.AdaptiveAvgPool2d(1)
|
||
mult = 2 ** self.n_downsampling
|
||
self.fully_connected = torch.nn.Linear(self.ngf*mult, 3)
|
||
|
||
def forward(self, x):
|
||
# real_I = x['real_I']
|
||
# dcp_J = x['dcp_J']
|
||
tmp = self.global_pooling(x)
|
||
tmp = tmp.view(tmp.shape[0],-1)
|
||
refined_A = F.relu(self.fully_connected(tmp))
|
||
|
||
return refined_A
|
||
|
||
|
||
class HazeRefiner(nn.Module):
|
||
"""docstring for HazeRefiner"""
|
||
def __init__(self, input_nc, output_nc, block_num):
|
||
super(HazeRefiner, self).__init__()
|
||
self.block_num = max(2, block_num)
|
||
|
||
self.block_0 = RefinerBlock(input_nc, output_nc)
|
||
for id_b in range(1, self.block_num):
|
||
setattr(self, 'block_%d'%id_b, RefinerBlock(input_nc+output_nc*id_b, output_nc))
|
||
|
||
def forward(self, x):
|
||
# real_I = x['real_I']
|
||
# dcp_J = x['dcp_J']]
|
||
last_J = x
|
||
cur_refine_J = self.block_0(last_J)
|
||
for id_b in range(1, self.block_num):
|
||
cur_block = getattr(self, 'block_%d'%id_b)
|
||
last_J = torch.cat((last_J, cur_refine_J), 1)
|
||
cur_refine_J = cur_block(last_J)
|
||
|
||
return cur_refine_J
|
||
|
||
|
||
class RefinerBlock(nn.Module):
|
||
"""docstring for RefinerBlock"""
|
||
def __init__(self, input_nc, output_nc):
|
||
super(RefinerBlock, self).__init__()
|
||
|
||
self.relu=nn.LeakyReLU(0.2, inplace=True)
|
||
|
||
self.tanh=nn.Tanh()
|
||
|
||
self.refine1= nn.Conv2d(input_nc, 20, kernel_size=3,stride=1,padding=1)
|
||
self.refine2= nn.Conv2d(20, 20, kernel_size=3,stride=1,padding=1)
|
||
|
||
self.conv1010 = nn.Conv2d(20, 1, kernel_size=1,stride=1,padding=0) # 1mm
|
||
self.conv1020 = nn.Conv2d(20, 1, kernel_size=1,stride=1,padding=0) # 1mm
|
||
self.conv1030 = nn.Conv2d(20, 1, kernel_size=1,stride=1,padding=0) # 1mm
|
||
self.conv1040 = nn.Conv2d(20, 1, kernel_size=1,stride=1,padding=0) # 1mm
|
||
|
||
self.refine3= nn.Conv2d(20+4, output_nc, kernel_size=3,stride=1,padding=1)
|
||
|
||
# self.upsample = F.upsample_nearest
|
||
# self.batch1 = nn.InstanceNorm2d(100, affine=True)
|
||
|
||
def forward(self, x):
|
||
output = self.relu((self.refine1(x)))
|
||
output = self.relu((self.refine2(output)))
|
||
shape_out = output.data.size()
|
||
# print(shape_out)
|
||
shape_out = shape_out[2:4]
|
||
|
||
x101 = F.avg_pool2d(output, 32)
|
||
|
||
x102 = F.avg_pool2d(output, 16)
|
||
|
||
x103 = F.avg_pool2d(output, 8)
|
||
|
||
x104 = F.avg_pool2d(output, 4)
|
||
x1010 = F.interpolate(self.relu(self.conv1010(x101)),size=shape_out, mode='nearest')
|
||
x1020 = F.interpolate(self.relu(self.conv1020(x102)),size=shape_out, mode='nearest')
|
||
x1030 = F.interpolate(self.relu(self.conv1030(x103)),size=shape_out, mode='nearest')
|
||
x1040 = F.interpolate(self.relu(self.conv1040(x104)),size=shape_out, mode='nearest')
|
||
|
||
output = torch.cat((x1010, x1020, x1030, x1040, output), 1)
|
||
output= self.tanh(self.refine3(output))
|
||
|
||
return output
|
||
|
||
|
||
class GANLoss(nn.Module):
|
||
"""Define different GAN objectives.
|
||
|
||
The GANLoss class abstracts away the need to create the target label tensor
|
||
that has the same size as the input.
|
||
"""
|
||
|
||
def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):
|
||
""" Initialize the GANLoss class.
|
||
|
||
Parameters:
|
||
gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.
|
||
target_real_label (bool) - - label for a real image
|
||
target_fake_label (bool) - - label of a fake image
|
||
|
||
Note: Do not use sigmoid as the last layer of Discriminator.
|
||
LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.
|
||
"""
|
||
super(GANLoss, self).__init__()
|
||
self.register_buffer('real_label', torch.tensor(target_real_label))
|
||
self.register_buffer('fake_label', torch.tensor(target_fake_label))
|
||
self.gan_mode = gan_mode
|
||
if gan_mode == 'lsgan':
|
||
self.loss = nn.MSELoss()
|
||
elif gan_mode == 'vanilla':
|
||
self.loss = nn.BCEWithLogitsLoss()
|
||
elif gan_mode in ['wgangp']:
|
||
self.loss = None
|
||
else:
|
||
raise NotImplementedError('gan mode %s not implemented' % gan_mode)
|
||
|
||
def get_target_tensor(self, prediction, target_is_real):
|
||
"""Create label tensors with the same size as the input.
|
||
|
||
Parameters:
|
||
prediction (tensor) - - tpyically the prediction from a discriminator
|
||
target_is_real (bool) - - if the ground truth label is for real images or fake images
|
||
|
||
Returns:
|
||
A label tensor filled with ground truth label, and with the size of the input
|
||
"""
|
||
|
||
if target_is_real:
|
||
target_tensor = self.real_label
|
||
else:
|
||
target_tensor = self.fake_label
|
||
return target_tensor.expand_as(prediction)
|
||
|
||
def __call__(self, prediction, target_is_real):
|
||
"""Calculate loss given Discriminator's output and grount truth labels.
|
||
|
||
Parameters:
|
||
prediction (tensor) - - tpyically the prediction output from a discriminator
|
||
target_is_real (bool) - - if the ground truth label is for real images or fake images
|
||
|
||
Returns:
|
||
the calculated loss.
|
||
"""
|
||
if self.gan_mode in ['lsgan', 'vanilla']:
|
||
target_tensor = self.get_target_tensor(prediction, target_is_real)
|
||
loss = self.loss(prediction, target_tensor)
|
||
elif self.gan_mode == 'wgangp':
|
||
if target_is_real:
|
||
loss = -prediction.mean()
|
||
else:
|
||
loss = prediction.mean()
|
||
return loss
|
||
|
||
|
||
def cal_gradient_penalty(netD, real_data, fake_data, device, type='mixed', constant=1.0, lambda_gp=10.0):
|
||
"""Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028
|
||
|
||
Arguments:
|
||
netD (network) -- discriminator network
|
||
real_data (tensor array) -- real images
|
||
fake_data (tensor array) -- generated images from the generator
|
||
device (str) -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
|
||
type (str) -- if we mix real and fake data or not [real | fake | mixed].
|
||
constant (float) -- the constant used in formula ( | |gradient||_2 - constant)^2
|
||
lambda_gp (float) -- weight for this loss
|
||
|
||
Returns the gradient penalty loss
|
||
"""
|
||
if lambda_gp > 0.0:
|
||
if type == 'real': # either use real images, fake images, or a linear interpolation of two.
|
||
interpolatesv = real_data
|
||
elif type == 'fake':
|
||
interpolatesv = fake_data
|
||
elif type == 'mixed':
|
||
alpha = torch.rand(real_data.shape[0], 1, device=device)
|
||
alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)
|
||
interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
|
||
else:
|
||
raise NotImplementedError('{} not implemented'.format(type))
|
||
interpolatesv.requires_grad_(True)
|
||
disc_interpolates = netD(interpolatesv)
|
||
gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv,
|
||
grad_outputs=torch.ones(disc_interpolates.size()).to(device),
|
||
create_graph=True, retain_graph=True, only_inputs=True)
|
||
gradients = gradients[0].view(real_data.size(0), -1) # flat the data
|
||
gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp # added eps
|
||
return gradient_penalty, gradients
|
||
else:
|
||
return 0.0, None
|
||
|
||
|
||
class ResnetGWithIntermediate(nn.Module):
|
||
"""Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
|
||
|
||
We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
|
||
"""
|
||
|
||
def __init__(self, input_nc, output_nc, ngf=6, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect', filtering='guided', r=10, eps=1e-3):
|
||
"""Construct a Resnet-based generator
|
||
|
||
Parameters:
|
||
input_nc (int) -- the number of channels in input images
|
||
output_nc (int) -- the number of channels in output images
|
||
ngf (int) -- the number of filters in the last conv layer
|
||
norm_layer -- normalization layer
|
||
use_dropout (bool) -- if use dropout layers
|
||
n_blocks (int) -- the number of ResNet blocks
|
||
padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero
|
||
"""
|
||
assert(n_blocks >= 0)
|
||
super(ResnetGWithIntermediate, self).__init__()
|
||
if type(norm_layer) == functools.partial:
|
||
use_bias = norm_layer.func == nn.InstanceNorm2d
|
||
else:
|
||
use_bias = norm_layer == nn.InstanceNorm2d
|
||
|
||
self.filtering=filtering
|
||
|
||
model = [nn.ReflectionPad2d(3),
|
||
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
|
||
norm_layer(ngf),
|
||
nn.ReLU(True)]
|
||
|
||
n_downsampling = 2
|
||
for i in range(n_downsampling): # add downsampling layers
|
||
mult = 2 ** i
|
||
model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
|
||
norm_layer(ngf * mult * 2),
|
||
nn.ReLU(True)]
|
||
|
||
mult = 2 ** n_downsampling
|
||
for i in range(n_blocks): # add ResNet blocks
|
||
|
||
model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
|
||
|
||
model_up_part = []
|
||
for i in range(n_downsampling): # add upsampling layers
|
||
mult = 2 ** (n_downsampling - i)
|
||
model_up_part += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
|
||
kernel_size=3, stride=2,
|
||
padding=1, output_padding=1,
|
||
bias=use_bias),
|
||
norm_layer(int(ngf * mult / 2)),
|
||
nn.ReLU(True)]
|
||
model_up_part += [nn.ReflectionPad2d(3)]
|
||
model_up_part += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
|
||
model_up_part += [nn.Tanh()]
|
||
|
||
self.downsampling = nn.Sequential(*model)
|
||
self.upsampling = nn.Sequential(*model_up_part)
|
||
|
||
if self.filtering is not None:
|
||
if self.filtering == 'max':
|
||
self.last_layer = nn.MaxPool2d(kernel_size=7, stride=1, padding=3)
|
||
elif self.filtering == 'guided':
|
||
self.last_layer = GuidedFilter(r=r, eps=eps)
|
||
|
||
def forward(self, x):
|
||
"""Standard forward"""
|
||
down_out = self.downsampling(x)
|
||
up_out = self.upsampling(down_out)
|
||
# rescale to [0,1]
|
||
up_out = (up_out + 1)/2
|
||
|
||
# rgb2gray
|
||
guidance = 0.2989 * x[:,0,:,:] + 0.5870 * x[:,1,:,:] + 0.1140 * x[:,2,:,:]
|
||
# rescale to [0,1]
|
||
guidance = (guidance + 1) / 2
|
||
guidance = torch.unsqueeze(guidance, dim=1)
|
||
|
||
if up_out.shape[2:4] != guidance.shape[2:4]:
|
||
up_out = F.interpolate(up_out,size=guidance.shape[2:4], mode='nearest')
|
||
|
||
# up_out = self.last_layer(guidance, up_out)
|
||
return self.last_layer(guidance, up_out), up_out
|
||
|
||
|
||
# Guided image filtering for grayscale images
|
||
class GuidedFilter(nn.Module):
|
||
def __init__(self, r=40, eps=1e-3, gpu_ids=None): # only work for gpu case at this moment
|
||
super(GuidedFilter, self).__init__()
|
||
self.r = r
|
||
self.eps = eps
|
||
# self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') # get device name: CPU or GPU
|
||
|
||
self.boxfilter = nn.AvgPool2d(kernel_size=2*self.r+1, stride=1,padding=self.r)
|
||
|
||
def forward(self, I, p):
|
||
"""
|
||
I -- guidance image, should be [0, 1]
|
||
p -- filtering input image, should be [0, 1]
|
||
"""
|
||
|
||
# N = self.boxfilter(self.tensor(p.size()).fill_(1))
|
||
N = self.boxfilter( torch.ones(p.size()) )
|
||
|
||
if I.is_cuda:
|
||
N = N.cuda()
|
||
|
||
# print(N.shape)
|
||
# print(I.shape)
|
||
# print('-----------')
|
||
|
||
mean_I = self.boxfilter(I) / N
|
||
mean_p = self.boxfilter(p) / N
|
||
mean_Ip = self.boxfilter(I*p) / N
|
||
cov_Ip = mean_Ip - mean_I * mean_p
|
||
|
||
mean_II = self.boxfilter(I*I) / N
|
||
var_I = mean_II - mean_I * mean_I
|
||
|
||
a = cov_Ip / (var_I + self.eps)
|
||
b = mean_p - a * mean_I
|
||
mean_a = self.boxfilter(a) / N
|
||
mean_b = self.boxfilter(b) / N
|
||
|
||
return mean_a * I + mean_b
|
||
|
||
|
||
class ResnetDehazeGenerator(nn.Module):
|
||
"""docstring for ResnetGenerator"""
|
||
def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
|
||
super(ResnetDehazeGenerator, self).__init__()
|
||
self.resnetG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer, use_dropout, n_blocks, padding_type)
|
||
self.refiner = HazeRefiner(input_nc, output_nc, block_num=2)
|
||
|
||
def forward(self, x):
|
||
res_out = self.resnetG(x)
|
||
ref_out = self.refiner(res_out)
|
||
|
||
return ref_out
|
||
|
||
|
||
|
||
class ResnetGenerator(nn.Module):
|
||
"""Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.
|
||
|
||
We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)
|
||
"""
|
||
|
||
def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):
|
||
"""Construct a Resnet-based generator
|
||
|
||
Parameters:
|
||
input_nc (int) -- the number of channels in input images
|
||
output_nc (int) -- the number of channels in output images
|
||
ngf (int) -- the number of filters in the last conv layer
|
||
norm_layer -- normalization layer
|
||
use_dropout (bool) -- if use dropout layers
|
||
n_blocks (int) -- the number of ResNet blocks
|
||
padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero
|
||
"""
|
||
assert(n_blocks >= 0)
|
||
super(ResnetGenerator, self).__init__()
|
||
if type(norm_layer) == functools.partial:
|
||
use_bias = norm_layer.func == nn.InstanceNorm2d
|
||
else:
|
||
use_bias = norm_layer == nn.InstanceNorm2d
|
||
|
||
model = [nn.ReflectionPad2d(3),
|
||
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
|
||
norm_layer(ngf),
|
||
nn.ReLU(True)]
|
||
|
||
n_downsampling = 2
|
||
for i in range(n_downsampling): # add downsampling layers
|
||
mult = 2 ** i
|
||
model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
|
||
norm_layer(ngf * mult * 2),
|
||
nn.ReLU(True)]
|
||
|
||
mult = 2 ** n_downsampling
|
||
for i in range(n_blocks): # add ResNet blocks
|
||
|
||
model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]
|
||
|
||
for i in range(n_downsampling): # add upsampling layers
|
||
mult = 2 ** (n_downsampling - i)
|
||
model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
|
||
kernel_size=3, stride=2,
|
||
padding=1, output_padding=1,
|
||
bias=use_bias),
|
||
norm_layer(int(ngf * mult / 2)),
|
||
nn.ReLU(True)]
|
||
model += [nn.ReflectionPad2d(3)]
|
||
model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
|
||
model += [nn.Tanh()]
|
||
|
||
self.model = nn.Sequential(*model)
|
||
|
||
def forward(self, input):
|
||
"""Standard forward"""
|
||
return self.model(input)
|
||
|
||
|
||
class ResnetBlock(nn.Module):
|
||
"""Define a Resnet block"""
|
||
|
||
def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
|
||
"""Initialize the Resnet block
|
||
|
||
A resnet block is a conv block with skip connections
|
||
We construct a conv block with build_conv_block function,
|
||
and implement skip connections in <forward> function.
|
||
Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf
|
||
"""
|
||
super(ResnetBlock, self).__init__()
|
||
self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)
|
||
|
||
def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
|
||
"""Construct a convolutional block.
|
||
|
||
Parameters:
|
||
dim (int) -- the number of channels in the conv layer.
|
||
padding_type (str) -- the name of padding layer: reflect | replicate | zero
|
||
norm_layer -- normalization layer
|
||
use_dropout (bool) -- if use dropout layers.
|
||
use_bias (bool) -- if the conv layer uses bias or not
|
||
|
||
Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))
|
||
"""
|
||
conv_block = []
|
||
p = 0
|
||
if padding_type == 'reflect':
|
||
conv_block += [nn.ReflectionPad2d(1)]
|
||
elif padding_type == 'replicate':
|
||
conv_block += [nn.ReplicationPad2d(1)]
|
||
elif padding_type == 'zero':
|
||
p = 1
|
||
else:
|
||
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
|
||
|
||
conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]
|
||
if use_dropout:
|
||
conv_block += [nn.Dropout(0.5)]
|
||
|
||
p = 0
|
||
if padding_type == 'reflect':
|
||
conv_block += [nn.ReflectionPad2d(1)]
|
||
elif padding_type == 'replicate':
|
||
conv_block += [nn.ReplicationPad2d(1)]
|
||
elif padding_type == 'zero':
|
||
p = 1
|
||
else:
|
||
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
|
||
conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]
|
||
|
||
return nn.Sequential(*conv_block)
|
||
|
||
def forward(self, x):
|
||
"""Forward function (with skip connections)"""
|
||
out = x + self.conv_block(x) # add skip connections
|
||
return out
|
||
|
||
|
||
class UnetGenerator(nn.Module):
|
||
"""Create a Unet-based generator"""
|
||
|
||
def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):
|
||
"""Construct a Unet generator
|
||
Parameters:
|
||
input_nc (int) -- the number of channels in input images
|
||
output_nc (int) -- the number of channels in output images
|
||
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
|
||
image of size 128x128 will become of size 1x1 # at the bottleneck
|
||
ngf (int) -- the number of filters in the last conv layer
|
||
norm_layer -- normalization layer
|
||
|
||
We construct the U-Net from the innermost layer to the outermost layer.
|
||
It is a recursive process.
|
||
"""
|
||
super(UnetGenerator, self).__init__()
|
||
# construct unet structure
|
||
unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True) # add the innermost layer
|
||
for i in range(num_downs - 5): # add intermediate layers with ngf * 8 filters
|
||
unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
|
||
# gradually reduce the number of filters from ngf * 8 to ngf
|
||
unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
|
||
unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
|
||
unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
|
||
self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer) # add the outermost layer
|
||
|
||
def forward(self, input):
|
||
"""Standard forward"""
|
||
return self.model(input)
|
||
|
||
|
||
class DCPDehazeGenerator(nn.Module):
|
||
"""Create a DCP Dehaze generator"""
|
||
def __init__(self, win_size=5, r=15, eps=1e-3):
|
||
super(DCPDehazeGenerator, self).__init__()
|
||
|
||
self.guided_filter = GuidedFilter(r=r, eps=eps)
|
||
self.neighborhood_size = win_size
|
||
self.omega = 0.95
|
||
|
||
def get_dark_channel(self, img, neighborhood_size):
|
||
shape = img.shape
|
||
if len(shape) == 4:
|
||
img_min,_ = torch.min(img, dim=1)
|
||
|
||
padSize = int(np.floor(neighborhood_size/2))
|
||
if neighborhood_size % 2 == 0:
|
||
pads = [padSize, padSize-1 ,padSize ,padSize-1]
|
||
else:
|
||
pads = [padSize, padSize ,padSize ,padSize]
|
||
|
||
img_min = F.pad(img_min, pads, mode='constant', value=1)
|
||
dark_img = -F.max_pool2d(-img_min, kernel_size=neighborhood_size, stride=1)
|
||
else:
|
||
raise NotImplementedError('get_tensor_dark_channel is only for 4-d tensor [N*C*H*W]')
|
||
|
||
dark_img = torch.unsqueeze(dark_img, dim=1)
|
||
return dark_img
|
||
|
||
def atmospheric_light(self, img, dark_img):
|
||
num,chl,height,width = img.shape
|
||
topNum = int(0.01*height*width)
|
||
|
||
A = torch.Tensor(num,chl,1,1)
|
||
if img.is_cuda:
|
||
A = A.cuda()
|
||
|
||
for num_id in range(num):
|
||
curImg = img[num_id,...]
|
||
curDarkImg = dark_img[num_id,0,...]
|
||
|
||
_, indices = curDarkImg.reshape([height*width]).sort(descending=True)
|
||
#curMask = indices < topNum
|
||
|
||
for chl_id in range(chl):
|
||
imgSlice = curImg[chl_id,...].reshape([height*width])
|
||
A[num_id,chl_id,0,0] = torch.mean(imgSlice[indices[0:topNum]])
|
||
|
||
return A
|
||
|
||
|
||
def forward(self, x):
|
||
if x.shape[1] > 1:
|
||
# rgb2gray
|
||
guidance = 0.2989 * x[:,0,:,:] + 0.5870 * x[:,1,:,:] + 0.1140 * x[:,2,:,:]
|
||
else:
|
||
guidance = x
|
||
# rescale to [0,1]
|
||
guidance = (guidance + 1)/2
|
||
guidance = torch.unsqueeze(guidance, dim=1)
|
||
imgPatch = (x + 1)/2
|
||
|
||
num,chl,height,width = imgPatch.shape
|
||
|
||
# dark_img and A with the range of [0,1]
|
||
dark_img = self.get_dark_channel(imgPatch, self.neighborhood_size)
|
||
A = self.atmospheric_light(imgPatch, dark_img)
|
||
|
||
map_A = A.repeat(1,1,height,width)
|
||
# make sure channel of trans_raw == 1
|
||
trans_raw = 1 - self.omega*self.get_dark_channel(imgPatch/map_A, self.neighborhood_size)
|
||
|
||
# get initial results
|
||
T_DCP = self.guided_filter(guidance, trans_raw)
|
||
J_DCP = (imgPatch - map_A)/T_DCP.repeat(1,3,1,1) + map_A
|
||
|
||
# import cv2
|
||
# temp = cv2.cvtColor(J_DCP[0].numpy().transpose([1,2,0]), cv2.COLOR_BGR2RGB)
|
||
# cv2.imshow('J_DCP',temp)
|
||
# cv2.imshow('T_DCP',T_DCP[0].numpy().transpose([1,2,0]))
|
||
# cv2.waitKey(0)
|
||
# exit()
|
||
|
||
return J_DCP, T_DCP, torch.squeeze(A)
|
||
|
||
|
||
class UnetTransGenerator(nn.Module):
|
||
"""Create a Unet-based generator"""
|
||
|
||
def __init__(self, input_nc, output_nc, num_downs, ngf=6, norm_layer=nn.BatchNorm2d, use_dropout=False, r=10, eps=1e-3):
|
||
"""Construct a Unet generator
|
||
Parameters:
|
||
input_nc (int) -- the number of channels in input images
|
||
output_nc (int) -- the number of channels in output images
|
||
num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
|
||
image of size 128x128 will become of size 1x1 # at the bottleneck
|
||
ngf (int) -- the number of filters in the last conv layer
|
||
norm_layer -- normalization layer
|
||
|
||
We construct the U-Net from the innermost layer to the outermost layer.
|
||
It is a recursive process.
|
||
"""
|
||
super(UnetTransGenerator, self).__init__()
|
||
# construct unet structure
|
||
unet_block = UnetAlignedSkipBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True) # add the innermost layer
|
||
for i in range(num_downs - 5): # add intermediate layers with ngf * 8 filters
|
||
unet_block = UnetAlignedSkipBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
|
||
# gradually reduce the number of filters from ngf * 8 to ngf
|
||
unet_block = UnetAlignedSkipBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
|
||
unet_block = UnetAlignedSkipBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
|
||
unet_block = UnetAlignedSkipBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
|
||
self.model = UnetAlignedSkipBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer) # add the outermost layer
|
||
|
||
self.guided_filter = GuidedFilter(r=r, eps=eps)
|
||
|
||
def forward(self, x):
|
||
if x.shape[1] > 1:
|
||
# rgb2gray
|
||
guidance = 0.2989 * x[:,0,:,:] + 0.5870 * x[:,1,:,:] + 0.1140 * x[:,2,:,:]
|
||
else:
|
||
guidance = x
|
||
# rescale to [0,1]
|
||
guidance = (guidance + 1) / 2
|
||
guidance = torch.unsqueeze(guidance, dim=1)
|
||
|
||
trans_raw = (self.model(x) + 1) / 2 # transmission ranges [0,1]
|
||
if trans_raw.shape[2:4] != guidance.shape[2:4]:
|
||
trans_raw = F.interpolate(trans_raw,size=guidance.shape[2:4], mode='nearest')
|
||
|
||
return self.guided_filter(guidance, trans_raw), trans_raw
|
||
|
||
# for trans refination
|
||
class UnetAlignedSkipBlock(nn.Module):
|
||
"""Defines the Unet submodule with skip connection.
|
||
X -------------------identity----------------------
|
||
|-- downsampling -- |submodule| -- upsampling --|
|
||
"""
|
||
|
||
def __init__(self, outer_nc, inner_nc, input_nc=None,
|
||
submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
|
||
"""Construct a Unet submodule with skip connections.
|
||
|
||
Parameters:
|
||
outer_nc (int) -- the number of filters in the outer conv layer
|
||
inner_nc (int) -- the number of filters in the inner conv layer
|
||
input_nc (int) -- the number of channels in input images/features
|
||
submodule (UnetAlignedSkipBlock) -- previously defined submodules
|
||
outermost (bool) -- if this module is the outermost module
|
||
innermost (bool) -- if this module is the innermost module
|
||
norm_layer -- normalization layer
|
||
user_dropout (bool) -- if use dropout layers.
|
||
"""
|
||
super(UnetAlignedSkipBlock, self).__init__()
|
||
self.outermost = outermost
|
||
if type(norm_layer) == functools.partial:
|
||
use_bias = norm_layer.func == nn.InstanceNorm2d
|
||
else:
|
||
use_bias = norm_layer == nn.InstanceNorm2d
|
||
if input_nc is None:
|
||
input_nc = outer_nc
|
||
downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
|
||
stride=2, padding=1, bias=use_bias)
|
||
downrelu = nn.LeakyReLU(0.2, True)
|
||
downnorm = norm_layer(inner_nc)
|
||
uprelu = nn.ReLU(True)
|
||
upnorm = norm_layer(outer_nc)
|
||
|
||
if outermost:
|
||
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
|
||
kernel_size=4, stride=2,
|
||
padding=1)
|
||
down = [downconv]
|
||
up = [uprelu, upconv, nn.Tanh()]
|
||
model = down + [submodule] + up
|
||
elif innermost:
|
||
upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
|
||
kernel_size=4, stride=2,
|
||
padding=1, bias=use_bias)
|
||
down = [downrelu, downconv]
|
||
up = [uprelu, upconv, upnorm]
|
||
model = down + up
|
||
else:
|
||
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
|
||
kernel_size=4, stride=2,
|
||
padding=1, bias=use_bias)
|
||
down = [downrelu, downconv, downnorm]
|
||
up = [uprelu, upconv, upnorm]
|
||
|
||
if use_dropout:
|
||
model = down + [submodule] + up + [nn.Dropout(0.5)]
|
||
else:
|
||
model = down + [submodule] + up
|
||
|
||
self.model = nn.Sequential(*model)
|
||
|
||
def forward(self, x):
|
||
if self.outermost:
|
||
return self.model(x)
|
||
else: # add skip connections
|
||
y = self.model(x)
|
||
# print(x.shape, y.shape)
|
||
if x.shape != y.shape:
|
||
y = F.interpolate(y, size=x.shape[2:4], mode='nearest')
|
||
return torch.cat([x, y], 1)
|
||
|
||
class UnetSkipConnectionBlock(nn.Module):
|
||
"""Defines the Unet submodule with skip connection.
|
||
X -------------------identity----------------------
|
||
|-- downsampling -- |submodule| -- upsampling --|
|
||
"""
|
||
|
||
def __init__(self, outer_nc, inner_nc, input_nc=None,
|
||
submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
|
||
"""Construct a Unet submodule with skip connections.
|
||
|
||
Parameters:
|
||
outer_nc (int) -- the number of filters in the outer conv layer
|
||
inner_nc (int) -- the number of filters in the inner conv layer
|
||
input_nc (int) -- the number of channels in input images/features
|
||
submodule (UnetSkipConnectionBlock) -- previously defined submodules
|
||
outermost (bool) -- if this module is the outermost module
|
||
innermost (bool) -- if this module is the innermost module
|
||
norm_layer -- normalization layer
|
||
user_dropout (bool) -- if use dropout layers.
|
||
"""
|
||
super(UnetSkipConnectionBlock, self).__init__()
|
||
self.outermost = outermost
|
||
if type(norm_layer) == functools.partial:
|
||
use_bias = norm_layer.func == nn.InstanceNorm2d
|
||
else:
|
||
use_bias = norm_layer == nn.InstanceNorm2d
|
||
if input_nc is None:
|
||
input_nc = outer_nc
|
||
downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
|
||
stride=2, padding=1, bias=use_bias)
|
||
downrelu = nn.LeakyReLU(0.2, True)
|
||
downnorm = norm_layer(inner_nc)
|
||
uprelu = nn.ReLU(True)
|
||
upnorm = norm_layer(outer_nc)
|
||
|
||
if outermost:
|
||
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
|
||
kernel_size=4, stride=2,
|
||
padding=1)
|
||
down = [downconv]
|
||
up = [uprelu, upconv, nn.Tanh()]
|
||
model = down + [submodule] + up
|
||
elif innermost:
|
||
upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
|
||
kernel_size=4, stride=2,
|
||
padding=1, bias=use_bias)
|
||
down = [downrelu, downconv]
|
||
up = [uprelu, upconv, upnorm]
|
||
model = down + up
|
||
else:
|
||
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
|
||
kernel_size=4, stride=2,
|
||
padding=1, bias=use_bias)
|
||
down = [downrelu, downconv, downnorm]
|
||
up = [uprelu, upconv, upnorm]
|
||
|
||
if use_dropout:
|
||
model = down + [submodule] + up + [nn.Dropout(0.5)]
|
||
else:
|
||
model = down + [submodule] + up
|
||
|
||
self.model = nn.Sequential(*model)
|
||
|
||
def forward(self, x):
|
||
if self.outermost:
|
||
return self.model(x)
|
||
else: # add skip connections
|
||
return torch.cat([x, self.model(x)], 1)
|
||
|
||
|
||
class NLayerDiscriminator(nn.Module):
|
||
"""Defines a PatchGAN discriminator"""
|
||
|
||
def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
|
||
"""Construct a PatchGAN discriminator
|
||
|
||
Parameters:
|
||
input_nc (int) -- the number of channels in input images
|
||
ndf (int) -- the number of filters in the last conv layer
|
||
n_layers (int) -- the number of conv layers in the discriminator
|
||
norm_layer -- normalization layer
|
||
"""
|
||
super(NLayerDiscriminator, self).__init__()
|
||
if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
|
||
use_bias = norm_layer.func == nn.InstanceNorm2d
|
||
else:
|
||
use_bias = norm_layer == nn.InstanceNorm2d
|
||
|
||
kw = 4
|
||
padw = 1
|
||
sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
|
||
nf_mult = 1
|
||
nf_mult_prev = 1
|
||
for n in range(1, n_layers): # gradually increase the number of filters
|
||
nf_mult_prev = nf_mult
|
||
nf_mult = min(2 ** n, 8)
|
||
sequence += [
|
||
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
|
||
norm_layer(ndf * nf_mult),
|
||
nn.LeakyReLU(0.2, True)
|
||
]
|
||
|
||
nf_mult_prev = nf_mult
|
||
nf_mult = min(2 ** n_layers, 8)
|
||
sequence += [
|
||
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
|
||
norm_layer(ndf * nf_mult),
|
||
nn.LeakyReLU(0.2, True)
|
||
]
|
||
|
||
sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # output 1 channel prediction map
|
||
self.model = nn.Sequential(*sequence)
|
||
|
||
def forward(self, input):
|
||
"""Standard forward."""
|
||
return self.model(input)
|
||
|
||
|
||
class PixelDiscriminator(nn.Module):
|
||
"""Defines a 1x1 PatchGAN discriminator (pixelGAN)"""
|
||
|
||
def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d):
|
||
"""Construct a 1x1 PatchGAN discriminator
|
||
|
||
Parameters:
|
||
input_nc (int) -- the number of channels in input images
|
||
ndf (int) -- the number of filters in the last conv layer
|
||
norm_layer -- normalization layer
|
||
"""
|
||
super(PixelDiscriminator, self).__init__()
|
||
if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
|
||
use_bias = norm_layer.func == nn.InstanceNorm2d
|
||
else:
|
||
use_bias = norm_layer == nn.InstanceNorm2d
|
||
|
||
self.net = [
|
||
nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),
|
||
nn.LeakyReLU(0.2, True),
|
||
nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),
|
||
norm_layer(ndf * 2),
|
||
nn.LeakyReLU(0.2, True),
|
||
nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias)]
|
||
|
||
self.net = nn.Sequential(*self.net)
|
||
|
||
def forward(self, input):
|
||
"""Standard forward."""
|
||
return self.net(input)
|
||
|
||
|
||
# Defines the Multiscale-PatchGAN discriminator with the specified arguments.
|
||
class MultiDiscriminator(nn.Module):
|
||
def __init__(self, input_nc, ndf=64, n_layers=5, norm_layer=nn.BatchNorm2d, use_sigmoid=False, gpu_ids=[]):
|
||
super(MultiDiscriminator, self).__init__()
|
||
self.gpu_ids = gpu_ids
|
||
if type(norm_layer) == functools.partial:
|
||
use_bias = norm_layer.func == nn.InstanceNorm2d
|
||
else:
|
||
use_bias = norm_layer == nn.InstanceNorm2d
|
||
|
||
# cannot deal with use_sigmoid=True case at thie moment
|
||
assert(use_sigmoid == False)
|
||
|
||
kw = 4
|
||
padw = int(np.ceil((kw-1)/2))
|
||
scale1 = [
|
||
nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
|
||
nn.LeakyReLU(0.2, True)
|
||
]
|
||
|
||
nf_mult = 1
|
||
nf_mult_prev = 1
|
||
for n in range(1, 3):
|
||
nf_mult_prev = nf_mult
|
||
nf_mult = min(2**n, 8)
|
||
scale1 += [
|
||
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
|
||
kernel_size=kw, stride=2, padding=padw, bias=use_bias),
|
||
norm_layer(ndf * nf_mult),
|
||
nn.LeakyReLU(0.2, True)
|
||
]
|
||
|
||
self.scale1 = nn.Sequential(*scale1)
|
||
scale1_output = []
|
||
scale1_output += [
|
||
nn.Conv2d(ndf * nf_mult, ndf * nf_mult,
|
||
kernel_size=kw, stride=1, padding=padw, bias=use_bias),
|
||
norm_layer(ndf * nf_mult),
|
||
nn.LeakyReLU(0.2, True)
|
||
]
|
||
scale1_output += [nn.Conv2d(ndf*nf_mult, 1, kernel_size=kw, stride=1, padding=padw)] # compress to 1 channel
|
||
self.scale1_output = nn.Sequential(*scale1_output)
|
||
|
||
scale2 = []
|
||
nf_mult = nf_mult
|
||
for n in range(3, n_layers):
|
||
nf_mult_prev = nf_mult
|
||
nf_mult = min(2**n, 8)
|
||
scale2 += [
|
||
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
|
||
kernel_size=kw, stride=2, padding=padw, bias=use_bias),
|
||
norm_layer(ndf * nf_mult),
|
||
nn.LeakyReLU(0.2, True)
|
||
]
|
||
|
||
nf_mult_prev = nf_mult
|
||
nf_mult = min(2**n_layers, 8)
|
||
scale2 += [
|
||
nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,
|
||
kernel_size=kw, stride=1, padding=padw, bias=use_bias),
|
||
norm_layer(ndf * nf_mult),
|
||
nn.LeakyReLU(0.2, True)
|
||
]
|
||
|
||
scale2 += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]
|
||
|
||
if use_sigmoid:
|
||
scale2 += [nn.Sigmoid()]
|
||
|
||
self.scale2 = nn.Sequential(*scale2)
|
||
|
||
def forward(self, input):
|
||
if len(self.gpu_ids) and isinstance(input.data, torch.cuda.FloatTensor):
|
||
scale1 = nn.parallel.data_parallel(self.scale1, input, self.gpu_ids)
|
||
output1 = nn.parallel.data_parallel(self.scale1_output, scale1, self.gpu_ids)
|
||
output2 = nn.parallel.data_parallel(self.scale2, scale1, self.gpu_ids)
|
||
else:
|
||
scale1 = self.scale1(input)
|
||
output1 = self.scale1_output(scale1)
|
||
output2 = self.scale2(scale1)
|
||
|
||
return output1, output2
|
||
|
||
class VGGLoss(nn.Module):
|
||
def __init__(self):
|
||
super(VGGLoss, self).__init__()
|
||
self.vgg = Vgg19().cuda()
|
||
self.criterion = nn.L1Loss()
|
||
self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
|
||
# self.weights = [0, 0, 1.0/8, 1.0/4, 1.0]
|
||
|
||
def forward(self, x, y):
|
||
x_vgg, y_vgg = self.vgg(x), self.vgg(y)
|
||
loss = 0
|
||
for i in range(len(x_vgg)):
|
||
loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
|
||
return loss
|
||
|
||
|
||
class Vgg19(torch.nn.Module):
|
||
def __init__(self, requires_grad=False):
|
||
super(Vgg19, self).__init__()
|
||
from torchvision import models
|
||
vgg_pretrained_features = models.vgg19(pretrained=True).features
|
||
self.slice1 = torch.nn.Sequential()
|
||
self.slice2 = torch.nn.Sequential()
|
||
self.slice3 = torch.nn.Sequential()
|
||
self.slice4 = torch.nn.Sequential()
|
||
self.slice5 = torch.nn.Sequential()
|
||
for x in range(2):
|
||
self.slice1.add_module(str(x), vgg_pretrained_features[x])
|
||
for x in range(2, 7):
|
||
self.slice2.add_module(str(x), vgg_pretrained_features[x])
|
||
for x in range(7, 12):
|
||
self.slice3.add_module(str(x), vgg_pretrained_features[x])
|
||
for x in range(12, 21):
|
||
self.slice4.add_module(str(x), vgg_pretrained_features[x])
|
||
for x in range(21, 30):
|
||
self.slice5.add_module(str(x), vgg_pretrained_features[x])
|
||
if not requires_grad:
|
||
for param in self.parameters():
|
||
param.requires_grad = False
|
||
|
||
def forward(self, X):
|
||
h_relu1 = self.slice1(X)
|
||
h_relu2 = self.slice2(h_relu1)
|
||
h_relu3 = self.slice3(h_relu2)
|
||
h_relu4 = self.slice4(h_relu3)
|
||
h_relu5 = self.slice5(h_relu4)
|
||
out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
|
||
return out
|