| """Modified from https://github.com/chaofengc/PSFRGAN
|
| """
|
| import numpy as np
|
| import torch.nn as nn
|
| from torch.nn import functional as F
|
|
|
|
|
| class NormLayer(nn.Module):
|
| """Normalization Layers.
|
|
|
| Args:
|
| channels: input channels, for batch norm and instance norm.
|
| input_size: input shape without batch size, for layer norm.
|
| """
|
|
|
| def __init__(self, channels, normalize_shape=None, norm_type='bn'):
|
| super(NormLayer, self).__init__()
|
| norm_type = norm_type.lower()
|
| self.norm_type = norm_type
|
| if norm_type == 'bn':
|
| self.norm = nn.BatchNorm2d(channels, affine=True)
|
| elif norm_type == 'in':
|
| self.norm = nn.InstanceNorm2d(channels, affine=False)
|
| elif norm_type == 'gn':
|
| self.norm = nn.GroupNorm(32, channels, affine=True)
|
| elif norm_type == 'pixel':
|
| self.norm = lambda x: F.normalize(x, p=2, dim=1)
|
| elif norm_type == 'layer':
|
| self.norm = nn.LayerNorm(normalize_shape)
|
| elif norm_type == 'none':
|
| self.norm = lambda x: x * 1.0
|
| else:
|
| assert 1 == 0, f'Norm type {norm_type} not support.'
|
|
|
| def forward(self, x, ref=None):
|
| if self.norm_type == 'spade':
|
| return self.norm(x, ref)
|
| else:
|
| return self.norm(x)
|
|
|
|
|
| class ReluLayer(nn.Module):
|
| """Relu Layer.
|
|
|
| Args:
|
| relu type: type of relu layer, candidates are
|
| - ReLU
|
| - LeakyReLU: default relu slope 0.2
|
| - PRelu
|
| - SELU
|
| - none: direct pass
|
| """
|
|
|
| def __init__(self, channels, relu_type='relu'):
|
| super(ReluLayer, self).__init__()
|
| relu_type = relu_type.lower()
|
| if relu_type == 'relu':
|
| self.func = nn.ReLU(True)
|
| elif relu_type == 'leakyrelu':
|
| self.func = nn.LeakyReLU(0.2, inplace=True)
|
| elif relu_type == 'prelu':
|
| self.func = nn.PReLU(channels)
|
| elif relu_type == 'selu':
|
| self.func = nn.SELU(True)
|
| elif relu_type == 'none':
|
| self.func = lambda x: x * 1.0
|
| else:
|
| assert 1 == 0, f'Relu type {relu_type} not support.'
|
|
|
| def forward(self, x):
|
| return self.func(x)
|
|
|
|
|
| class ConvLayer(nn.Module):
|
|
|
| def __init__(self,
|
| in_channels,
|
| out_channels,
|
| kernel_size=3,
|
| scale='none',
|
| norm_type='none',
|
| relu_type='none',
|
| use_pad=True,
|
| bias=True):
|
| super(ConvLayer, self).__init__()
|
| self.use_pad = use_pad
|
| self.norm_type = norm_type
|
| if norm_type in ['bn']:
|
| bias = False
|
|
|
| stride = 2 if scale == 'down' else 1
|
|
|
| self.scale_func = lambda x: x
|
| if scale == 'up':
|
| self.scale_func = lambda x: nn.functional.interpolate(x, scale_factor=2, mode='nearest')
|
|
|
| self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.) / 2)))
|
| self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias)
|
|
|
| self.relu = ReluLayer(out_channels, relu_type)
|
| self.norm = NormLayer(out_channels, norm_type=norm_type)
|
|
|
| def forward(self, x):
|
| out = self.scale_func(x)
|
| if self.use_pad:
|
| out = self.reflection_pad(out)
|
| out = self.conv2d(out)
|
| out = self.norm(out)
|
| out = self.relu(out)
|
| return out
|
|
|
|
|
| class ResidualBlock(nn.Module):
|
| """
|
| Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
|
| """
|
|
|
| def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', scale='none'):
|
| super(ResidualBlock, self).__init__()
|
|
|
| if scale == 'none' and c_in == c_out:
|
| self.shortcut_func = lambda x: x
|
| else:
|
| self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)
|
|
|
| scale_config_dict = {'down': ['none', 'down'], 'up': ['up', 'none'], 'none': ['none', 'none']}
|
| scale_conf = scale_config_dict[scale]
|
|
|
| self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type)
|
| self.conv2 = ConvLayer(c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type='none')
|
|
|
| def forward(self, x):
|
| identity = self.shortcut_func(x)
|
|
|
| res = self.conv1(x)
|
| res = self.conv2(res)
|
| return identity + res
|
|
|
|
|
| class ParseNet(nn.Module):
|
|
|
| def __init__(self,
|
| in_size=128,
|
| out_size=128,
|
| min_feat_size=32,
|
| base_ch=64,
|
| parsing_ch=19,
|
| res_depth=10,
|
| relu_type='LeakyReLU',
|
| norm_type='bn',
|
| ch_range=[32, 256]):
|
| super().__init__()
|
| self.res_depth = res_depth
|
| act_args = {'norm_type': norm_type, 'relu_type': relu_type}
|
| min_ch, max_ch = ch_range
|
|
|
| ch_clip = lambda x: max(min_ch, min(x, max_ch))
|
| min_feat_size = min(in_size, min_feat_size)
|
|
|
| down_steps = int(np.log2(in_size // min_feat_size))
|
| up_steps = int(np.log2(out_size // min_feat_size))
|
|
|
|
|
| self.encoder = []
|
| self.encoder.append(ConvLayer(3, base_ch, 3, 1))
|
| head_ch = base_ch
|
| for i in range(down_steps):
|
| cin, cout = ch_clip(head_ch), ch_clip(head_ch * 2)
|
| self.encoder.append(ResidualBlock(cin, cout, scale='down', **act_args))
|
| head_ch = head_ch * 2
|
|
|
| self.body = []
|
| for i in range(res_depth):
|
| self.body.append(ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args))
|
|
|
| self.decoder = []
|
| for i in range(up_steps):
|
| cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2)
|
| self.decoder.append(ResidualBlock(cin, cout, scale='up', **act_args))
|
| head_ch = head_ch // 2
|
|
|
| self.encoder = nn.Sequential(*self.encoder)
|
| self.body = nn.Sequential(*self.body)
|
| self.decoder = nn.Sequential(*self.decoder)
|
| self.out_img_conv = ConvLayer(ch_clip(head_ch), 3)
|
| self.out_mask_conv = ConvLayer(ch_clip(head_ch), parsing_ch)
|
|
|
| def forward(self, x):
|
| feat = self.encoder(x)
|
| x = feat + self.body(feat)
|
| x = self.decoder(x)
|
| out_img = self.out_img_conv(x)
|
| out_mask = self.out_mask_conv(x)
|
| return out_mask, out_img
|
|
|
