HexB/CodeFormer

1"""Modified from https://github.com/chaofengc/PSFRGAN2"""3import numpy as np4import torch.nn as nn5from torch.nn import functional as F6 7 8class NormLayer(nn.Module):9    """Normalization Layers.10 11    Args:12        channels: input channels, for batch norm and instance norm.13        input_size: input shape without batch size, for layer norm.14    """15 16    def __init__(self, channels, normalize_shape=None, norm_type='bn'):17        super(NormLayer, self).__init__()18        norm_type = norm_type.lower()19        self.norm_type = norm_type20        if norm_type == 'bn':21            self.norm = nn.BatchNorm2d(channels, affine=True)22        elif norm_type == 'in':23            self.norm = nn.InstanceNorm2d(channels, affine=False)24        elif norm_type == 'gn':25            self.norm = nn.GroupNorm(32, channels, affine=True)26        elif norm_type == 'pixel':27            self.norm = lambda x: F.normalize(x, p=2, dim=1)28        elif norm_type == 'layer':29            self.norm = nn.LayerNorm(normalize_shape)30        elif norm_type == 'none':31            self.norm = lambda x: x * 1.032        else:33            assert 1 == 0, f'Norm type {norm_type} not support.'34 35    def forward(self, x, ref=None):36        if self.norm_type == 'spade':37            return self.norm(x, ref)38        else:39            return self.norm(x)40 41 42class ReluLayer(nn.Module):43    """Relu Layer.44 45    Args:46        relu type: type of relu layer, candidates are47            - ReLU48            - LeakyReLU: default relu slope 0.249            - PRelu50            - SELU51            - none: direct pass52    """53 54    def __init__(self, channels, relu_type='relu'):55        super(ReluLayer, self).__init__()56        relu_type = relu_type.lower()57        if relu_type == 'relu':58            self.func = nn.ReLU(True)59        elif relu_type == 'leakyrelu':60            self.func = nn.LeakyReLU(0.2, inplace=True)61        elif relu_type == 'prelu':62            self.func = nn.PReLU(channels)63        elif relu_type == 'selu':64            self.func = nn.SELU(True)65        elif relu_type == 'none':66            self.func = lambda x: x * 1.067        else:68            assert 1 == 0, f'Relu type {relu_type} not support.'69 70    def forward(self, x):71        return self.func(x)72 73 74class ConvLayer(nn.Module):75 76    def __init__(self,77                 in_channels,78                 out_channels,79                 kernel_size=3,80                 scale='none',81                 norm_type='none',82                 relu_type='none',83                 use_pad=True,84                 bias=True):85        super(ConvLayer, self).__init__()86        self.use_pad = use_pad87        self.norm_type = norm_type88        if norm_type in ['bn']:89            bias = False90 91        stride = 2 if scale == 'down' else 192 93        self.scale_func = lambda x: x94        if scale == 'up':95            self.scale_func = lambda x: nn.functional.interpolate(x, scale_factor=2, mode='nearest')96 97        self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.) / 2)))98        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias)99 100        self.relu = ReluLayer(out_channels, relu_type)101        self.norm = NormLayer(out_channels, norm_type=norm_type)102 103    def forward(self, x):104        out = self.scale_func(x)105        if self.use_pad:106            out = self.reflection_pad(out)107        out = self.conv2d(out)108        out = self.norm(out)109        out = self.relu(out)110        return out111 112 113class ResidualBlock(nn.Module):114    """115    Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html116    """117 118    def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', scale='none'):119        super(ResidualBlock, self).__init__()120 121        if scale == 'none' and c_in == c_out:122            self.shortcut_func = lambda x: x123        else:124            self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)125 126        scale_config_dict = {'down': ['none', 'down'], 'up': ['up', 'none'], 'none': ['none', 'none']}127        scale_conf = scale_config_dict[scale]128 129        self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type)130        self.conv2 = ConvLayer(c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type='none')131 132    def forward(self, x):133        identity = self.shortcut_func(x)134 135        res = self.conv1(x)136        res = self.conv2(res)137        return identity + res138 139 140class ParseNet(nn.Module):141 142    def __init__(self,143                 in_size=128,144                 out_size=128,145                 min_feat_size=32,146                 base_ch=64,147                 parsing_ch=19,148                 res_depth=10,149                 relu_type='LeakyReLU',150                 norm_type='bn',151                 ch_range=[32, 256]):152        super().__init__()153        self.res_depth = res_depth154        act_args = {'norm_type': norm_type, 'relu_type': relu_type}155        min_ch, max_ch = ch_range156 157        ch_clip = lambda x: max(min_ch, min(x, max_ch))  # noqa: E731158        min_feat_size = min(in_size, min_feat_size)159 160        down_steps = int(np.log2(in_size // min_feat_size))161        up_steps = int(np.log2(out_size // min_feat_size))162 163        # =============== define encoder-body-decoder ====================164        self.encoder = []165        self.encoder.append(ConvLayer(3, base_ch, 3, 1))166        head_ch = base_ch167        for i in range(down_steps):168            cin, cout = ch_clip(head_ch), ch_clip(head_ch * 2)169            self.encoder.append(ResidualBlock(cin, cout, scale='down', **act_args))170            head_ch = head_ch * 2171 172        self.body = []173        for i in range(res_depth):174            self.body.append(ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args))175 176        self.decoder = []177        for i in range(up_steps):178            cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2)179            self.decoder.append(ResidualBlock(cin, cout, scale='up', **act_args))180            head_ch = head_ch // 2181 182        self.encoder = nn.Sequential(*self.encoder)183        self.body = nn.Sequential(*self.body)184        self.decoder = nn.Sequential(*self.decoder)185        self.out_img_conv = ConvLayer(ch_clip(head_ch), 3)186        self.out_mask_conv = ConvLayer(ch_clip(head_ch), parsing_ch)187 188    def forward(self, x):189        feat = self.encoder(x)190        x = feat + self.body(feat)191        x = self.decoder(x)192        out_img = self.out_img_conv(x)193        out_mask = self.out_mask_conv(x)194        return out_mask, out_img195