Remove model code, only keep model checkpoint

README.md

@@ -2,6 +2,8 @@
 
 [**Paper**](https://arxiv.org/abs/2101.04061) **|** [**Project Page**](https://xinntao.github.io/projects/gfpgan)    [English](README.md) **|** [简体中文](README_CN.md)
 
+GitHub: https://github.com/TencentARC/GFPGAN
+
 GFPGAN is a blind face restoration algorithm towards real-world face images.
 
 <a href="https://colab.research.google.com/drive/1sVsoBd9AjckIXThgtZhGrHRfFI6UUYOo"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a>

archs/__init__.py

@@ -1,12 +0,0 @@
-import importlib
-from os import path as osp
-
-from basicsr.utils import scandir
-
-# automatically scan and import arch modules for registry
-# scan all the files under the 'archs' folder and collect files ending with
-# '_arch.py'
-arch_folder = osp.dirname(osp.abspath(__file__))
-arch_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(arch_folder) if v.endswith('_arch.py')]
-# import all the arch modules
-_arch_modules = [importlib.import_module(f'archs.{file_name}') for file_name in arch_filenames]

archs/arcface_arch.py

@@ -1,198 +0,0 @@
-import torch.nn as nn
-
-from basicsr.utils.registry import ARCH_REGISTRY
-
-
-def conv3x3(in_planes, out_planes, stride=1):
-    """3x3 convolution with padding"""
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
-
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
-        super(BasicBlock, self).__init__()
-        self.conv1 = conv3x3(inplanes, planes, stride)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.relu = nn.ReLU(inplace=True)
-        self.conv2 = conv3x3(planes, planes)
-        self.bn2 = nn.BatchNorm2d(planes)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class IRBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True):
-        super(IRBlock, self).__init__()
-        self.bn0 = nn.BatchNorm2d(inplanes)
-        self.conv1 = conv3x3(inplanes, inplanes)
-        self.bn1 = nn.BatchNorm2d(inplanes)
-        self.prelu = nn.PReLU()
-        self.conv2 = conv3x3(inplanes, planes, stride)
-        self.bn2 = nn.BatchNorm2d(planes)
-        self.downsample = downsample
-        self.stride = stride
-        self.use_se = use_se
-        if self.use_se:
-            self.se = SEBlock(planes)
-
-    def forward(self, x):
-        residual = x
-        out = self.bn0(x)
-        out = self.conv1(out)
-        out = self.bn1(out)
-        out = self.prelu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        if self.use_se:
-            out = self.se(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.prelu(out)
-
-        return out
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
-        super(Bottleneck, self).__init__()
-        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(planes)
-        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(planes)
-        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
-        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu(out)
-
-        out = self.conv3(out)
-        out = self.bn3(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class SEBlock(nn.Module):
-
-    def __init__(self, channel, reduction=16):
-        super(SEBlock, self).__init__()
-        self.avg_pool = nn.AdaptiveAvgPool2d(1)
-        self.fc = nn.Sequential(
-            nn.Linear(channel, channel // reduction), nn.PReLU(), nn.Linear(channel // reduction, channel),
-            nn.Sigmoid())
-
-    def forward(self, x):
-        b, c, _, _ = x.size()
-        y = self.avg_pool(x).view(b, c)
-        y = self.fc(y).view(b, c, 1, 1)
-        return x * y
-
-
-@ARCH_REGISTRY.register()
-class ResNetArcFace(nn.Module):
-
-    def __init__(self, block, layers, use_se=True):
-        if block == 'IRBlock':
-            block = IRBlock
-        self.inplanes = 64
-        self.use_se = use_se
-        super(ResNetArcFace, self).__init__()
-        self.conv1 = nn.Conv2d(1, 64, kernel_size=3, padding=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(64)
-        self.prelu = nn.PReLU()
-        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
-        self.layer1 = self._make_layer(block, 64, layers[0])
-        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
-        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
-        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
-        self.bn4 = nn.BatchNorm2d(512)
-        self.dropout = nn.Dropout()
-        self.fc5 = nn.Linear(512 * 8 * 8, 512)
-        self.bn5 = nn.BatchNorm1d(512)
-
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                nn.init.xavier_normal_(m.weight)
-            elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
-                nn.init.constant_(m.weight, 1)
-                nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.Linear):
-                nn.init.xavier_normal_(m.weight)
-                nn.init.constant_(m.bias, 0)
-
-    def _make_layer(self, block, planes, blocks, stride=1):
-        downsample = None
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),
-                nn.BatchNorm2d(planes * block.expansion),
-            )
-        layers = []
-        layers.append(block(self.inplanes, planes, stride, downsample, use_se=self.use_se))
-        self.inplanes = planes
-        for _ in range(1, blocks):
-            layers.append(block(self.inplanes, planes, use_se=self.use_se))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.prelu(x)
-        x = self.maxpool(x)
-
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        x = self.bn4(x)
-        x = self.dropout(x)
-        x = x.view(x.size(0), -1)
-        x = self.fc5(x)
-        x = self.bn5(x)
-
-        return x

archs/gfpganv1_arch.py

@@ -1,418 +0,0 @@
-import math
-import random
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-from basicsr.archs.stylegan2_arch import (ConvLayer, EqualConv2d, EqualLinear, ResBlock, ScaledLeakyReLU,
-                                          StyleGAN2Generator)
-from basicsr.ops.fused_act import FusedLeakyReLU
-from basicsr.utils.registry import ARCH_REGISTRY
-
-
-class StyleGAN2GeneratorSFT(StyleGAN2Generator):
-    """StyleGAN2 Generator.
-
-    Args:
-        out_size (int): The spatial size of outputs.
-        num_style_feat (int): Channel number of style features. Default: 512.
-        num_mlp (int): Layer number of MLP style layers. Default: 8.
-        channel_multiplier (int): Channel multiplier for large networks of
-            StyleGAN2. Default: 2.
-        resample_kernel (list[int]): A list indicating the 1D resample kernel
-            magnitude. A cross production will be applied to extent 1D resample
-            kenrel to 2D resample kernel. Default: [1, 3, 3, 1].
-        lr_mlp (float): Learning rate multiplier for mlp layers. Default: 0.01.
-    """
-
-    def __init__(self,
-                 out_size,
-                 num_style_feat=512,
-                 num_mlp=8,
-                 channel_multiplier=2,
-                 resample_kernel=(1, 3, 3, 1),
-                 lr_mlp=0.01,
-                 narrow=1,
-                 sft_half=False):
-        super(StyleGAN2GeneratorSFT, self).__init__(
-            out_size,
-            num_style_feat=num_style_feat,
-            num_mlp=num_mlp,
-            channel_multiplier=channel_multiplier,
-            resample_kernel=resample_kernel,
-            lr_mlp=lr_mlp,
-            narrow=narrow)
-        self.sft_half = sft_half
-
-    def forward(self,
-                styles,
-                conditions,
-                input_is_latent=False,
-                noise=None,
-                randomize_noise=True,
-                truncation=1,
-                truncation_latent=None,
-                inject_index=None,
-                return_latents=False):
-        """Forward function for StyleGAN2Generator.
-
-        Args:
-            styles (list[Tensor]): Sample codes of styles.
-            input_is_latent (bool): Whether input is latent style.
-                Default: False.
-            noise (Tensor | None): Input noise or None. Default: None.
-            randomize_noise (bool): Randomize noise, used when 'noise' is
-                False. Default: True.
-            truncation (float): TODO. Default: 1.
-            truncation_latent (Tensor | None): TODO. Default: None.
-            inject_index (int | None): The injection index for mixing noise.
-                Default: None.
-            return_latents (bool): Whether to return style latents.
-                Default: False.
-        """
-        # style codes -> latents with Style MLP layer
-        if not input_is_latent:
-            styles = [self.style_mlp(s) for s in styles]
-        # noises
-        if noise is None:
-            if randomize_noise:
-                noise = [None] * self.num_layers  # for each style conv layer
-            else:  # use the stored noise
-                noise = [getattr(self.noises, f'noise{i}') for i in range(self.num_layers)]
-        # style truncation
-        if truncation < 1:
-            style_truncation = []
-            for style in styles:
-                style_truncation.append(truncation_latent + truncation * (style - truncation_latent))
-            styles = style_truncation
-        # get style latent with injection
-        if len(styles) == 1:
-            inject_index = self.num_latent
-
-            if styles[0].ndim < 3:
-                # repeat latent code for all the layers
-                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
-            else:  # used for encoder with different latent code for each layer
-                latent = styles[0]
-        elif len(styles) == 2:  # mixing noises
-            if inject_index is None:
-                inject_index = random.randint(1, self.num_latent - 1)
-            latent1 = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
-            latent2 = styles[1].unsqueeze(1).repeat(1, self.num_latent - inject_index, 1)
-            latent = torch.cat([latent1, latent2], 1)
-
-        # main generation
-        out = self.constant_input(latent.shape[0])
-        out = self.style_conv1(out, latent[:, 0], noise=noise[0])
-        skip = self.to_rgb1(out, latent[:, 1])
-
-        i = 1
-        for conv1, conv2, noise1, noise2, to_rgb in zip(self.style_convs[::2], self.style_convs[1::2], noise[1::2],
-                                                        noise[2::2], self.to_rgbs):
-            out = conv1(out, latent[:, i], noise=noise1)
-
-            # the conditions may have fewer levels
-            if i < len(conditions):
-                # SFT part to combine the conditions
-                if self.sft_half:
-                    out_same, out_sft = torch.split(out, int(out.size(1) // 2), dim=1)
-                    out_sft = out_sft * conditions[i - 1] + conditions[i]
-                    out = torch.cat([out_same, out_sft], dim=1)
-                else:
-                    out = out * conditions[i - 1] + conditions[i]
-
-            out = conv2(out, latent[:, i + 1], noise=noise2)
-            skip = to_rgb(out, latent[:, i + 2], skip)
-            i += 2
-
-        image = skip
-
-        if return_latents:
-            return image, latent
-        else:
-            return image, None
-
-
-class ConvUpLayer(nn.Module):
-    """Conv Up Layer. Bilinear upsample + Conv.
-
-    Args:
-        in_channels (int): Channel number of the input.
-        out_channels (int): Channel number of the output.
-        kernel_size (int): Size of the convolving kernel.
-        stride (int): Stride of the convolution. Default: 1
-        padding (int): Zero-padding added to both sides of the input.
-            Default: 0.
-        bias (bool): If ``True``, adds a learnable bias to the output.
-            Default: ``True``.
-        bias_init_val (float): Bias initialized value. Default: 0.
-        activate (bool): Whether use activateion. Default: True.
-    """
-
-    def __init__(self,
-                 in_channels,
-                 out_channels,
-                 kernel_size,
-                 stride=1,
-                 padding=0,
-                 bias=True,
-                 bias_init_val=0,
-                 activate=True):
-        super(ConvUpLayer, self).__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.kernel_size = kernel_size
-        self.stride = stride
-        self.padding = padding
-        self.scale = 1 / math.sqrt(in_channels * kernel_size**2)
-
-        self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size, kernel_size))
-
-        if bias and not activate:
-            self.bias = nn.Parameter(torch.zeros(out_channels).fill_(bias_init_val))
-        else:
-            self.register_parameter('bias', None)
-
-        # activation
-        if activate:
-            if bias:
-                self.activation = FusedLeakyReLU(out_channels)
-            else:
-                self.activation = ScaledLeakyReLU(0.2)
-        else:
-            self.activation = None
-
-    def forward(self, x):
-        # bilinear upsample
-        out = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
-        # conv
-        out = F.conv2d(
-            out,
-            self.weight * self.scale,
-            bias=self.bias,
-            stride=self.stride,
-            padding=self.padding,
-        )
-        # activation
-        if self.activation is not None:
-            out = self.activation(out)
-        return out
-
-
-class ResUpBlock(nn.Module):
-    """Residual block with upsampling.
-
-    Args:
-        in_channels (int): Channel number of the input.
-        out_channels (int): Channel number of the output.
-    """
-
-    def __init__(self, in_channels, out_channels):
-        super(ResUpBlock, self).__init__()
-
-        self.conv1 = ConvLayer(in_channels, in_channels, 3, bias=True, activate=True)
-        self.conv2 = ConvUpLayer(in_channels, out_channels, 3, stride=1, padding=1, bias=True, activate=True)
-        self.skip = ConvUpLayer(in_channels, out_channels, 1, bias=False, activate=False)
-
-    def forward(self, x):
-        out = self.conv1(x)
-        out = self.conv2(out)
-        skip = self.skip(x)
-        out = (out + skip) / math.sqrt(2)
-        return out
-
-
-@ARCH_REGISTRY.register()
-class GFPGANv1(nn.Module):
-    """Unet + StyleGAN2 decoder with SFT."""
-
-    def __init__(
-            self,
-            out_size,
-            num_style_feat=512,
-            channel_multiplier=1,
-            resample_kernel=(1, 3, 3, 1),
-            decoder_load_path=None,
-            fix_decoder=True,
-            # for stylegan decoder
-            num_mlp=8,
-            lr_mlp=0.01,
-            input_is_latent=False,
-            different_w=False,
-            narrow=1,
-            sft_half=False):
-
-        super(GFPGANv1, self).__init__()
-        self.input_is_latent = input_is_latent
-        self.different_w = different_w
-        self.num_style_feat = num_style_feat
-
-        unet_narrow = narrow * 0.5
-        channels = {
-            '4': int(512 * unet_narrow),
-            '8': int(512 * unet_narrow),
-            '16': int(512 * unet_narrow),
-            '32': int(512 * unet_narrow),
-            '64': int(256 * channel_multiplier * unet_narrow),
-            '128': int(128 * channel_multiplier * unet_narrow),
-            '256': int(64 * channel_multiplier * unet_narrow),
-            '512': int(32 * channel_multiplier * unet_narrow),
-            '1024': int(16 * channel_multiplier * unet_narrow)
-        }
-
-        self.log_size = int(math.log(out_size, 2))
-        first_out_size = 2**(int(math.log(out_size, 2)))
-
-        self.conv_body_first = ConvLayer(3, channels[f'{first_out_size}'], 1, bias=True, activate=True)
-
-        # downsample
-        in_channels = channels[f'{first_out_size}']
-        self.conv_body_down = nn.ModuleList()
-        for i in range(self.log_size, 2, -1):
-            out_channels = channels[f'{2**(i - 1)}']
-            self.conv_body_down.append(ResBlock(in_channels, out_channels, resample_kernel))
-            in_channels = out_channels
-
-        self.final_conv = ConvLayer(in_channels, channels['4'], 3, bias=True, activate=True)
-
-        # upsample
-        in_channels = channels['4']
-        self.conv_body_up = nn.ModuleList()
-        for i in range(3, self.log_size + 1):
-            out_channels = channels[f'{2**i}']
-            self.conv_body_up.append(ResUpBlock(in_channels, out_channels))
-            in_channels = out_channels
-
-        # to RGB
-        self.toRGB = nn.ModuleList()
-        for i in range(3, self.log_size + 1):
-            self.toRGB.append(EqualConv2d(channels[f'{2**i}'], 3, 1, stride=1, padding=0, bias=True, bias_init_val=0))
-
-        if different_w:
-            linear_out_channel = (int(math.log(out_size, 2)) * 2 - 2) * num_style_feat
-        else:
-            linear_out_channel = num_style_feat
-
-        self.final_linear = EqualLinear(
-            channels['4'] * 4 * 4, linear_out_channel, bias=True, bias_init_val=0, lr_mul=1, activation=None)
-
-        self.stylegan_decoder = StyleGAN2GeneratorSFT(
-            out_size=out_size,
-            num_style_feat=num_style_feat,
-            num_mlp=num_mlp,
-            channel_multiplier=channel_multiplier,
-            resample_kernel=resample_kernel,
-            lr_mlp=lr_mlp,
-            narrow=narrow,
-            sft_half=sft_half)
-
-        if decoder_load_path:
-            self.stylegan_decoder.load_state_dict(
-                torch.load(decoder_load_path, map_location=lambda storage, loc: storage)['params_ema'])
-        if fix_decoder:
-            for _, param in self.stylegan_decoder.named_parameters():
-                param.requires_grad = False
-
-        # for SFT
-        self.condition_scale = nn.ModuleList()
-        self.condition_shift = nn.ModuleList()
-        for i in range(3, self.log_size + 1):
-            out_channels = channels[f'{2**i}']
-            if sft_half:
-                sft_out_channels = out_channels
-            else:
-                sft_out_channels = out_channels * 2
-            self.condition_scale.append(
-                nn.Sequential(
-                    EqualConv2d(out_channels, out_channels, 3, stride=1, padding=1, bias=True, bias_init_val=0),
-                    ScaledLeakyReLU(0.2),
-                    EqualConv2d(out_channels, sft_out_channels, 3, stride=1, padding=1, bias=True, bias_init_val=1)))
-            self.condition_shift.append(
-                nn.Sequential(
-                    EqualConv2d(out_channels, out_channels, 3, stride=1, padding=1, bias=True, bias_init_val=0),
-                    ScaledLeakyReLU(0.2),
-                    EqualConv2d(out_channels, sft_out_channels, 3, stride=1, padding=1, bias=True, bias_init_val=0)))
-
-    def forward(self,
-                x,
-                return_latents=False,
-                save_feat_path=None,
-                load_feat_path=None,
-                return_rgb=True,
-                randomize_noise=True):
-        conditions = []
-        unet_skips = []
-        out_rgbs = []
-
-        # encoder
-        feat = self.conv_body_first(x)
-        for i in range(self.log_size - 2):
-            feat = self.conv_body_down[i](feat)
-            unet_skips.insert(0, feat)
-
-        feat = self.final_conv(feat)
-
-        # style code
-        style_code = self.final_linear(feat.view(feat.size(0), -1))
-        if self.different_w:
-            style_code = style_code.view(style_code.size(0), -1, self.num_style_feat)
-
-        # decode
-        for i in range(self.log_size - 2):
-            # add unet skip
-            feat = feat + unet_skips[i]
-            # ResUpLayer
-            feat = self.conv_body_up[i](feat)
-            # generate scale and shift for SFT layer
-            scale = self.condition_scale[i](feat)
-            conditions.append(scale.clone())
-            shift = self.condition_shift[i](feat)
-            conditions.append(shift.clone())
-            # generate rgb images
-            if return_rgb:
-                out_rgbs.append(self.toRGB[i](feat))
-
-        if save_feat_path is not None:
-            torch.save(conditions, save_feat_path)
-        if load_feat_path is not None:
-            conditions = torch.load(load_feat_path)
-            conditions = [v.cuda() for v in conditions]
-
-        # decoder
-        image, _ = self.stylegan_decoder([style_code],
-                                         conditions,
-                                         return_latents=return_latents,
-                                         input_is_latent=self.input_is_latent,
-                                         randomize_noise=randomize_noise)
-
-        return image, out_rgbs
-
-
-@ARCH_REGISTRY.register()
-class FacialComponentDiscriminator(nn.Module):
-
-    def __init__(self):
-        super(FacialComponentDiscriminator, self).__init__()
-
-        self.conv1 = ConvLayer(3, 64, 3, downsample=False, resample_kernel=(1, 3, 3, 1), bias=True, activate=True)
-        self.conv2 = ConvLayer(64, 128, 3, downsample=True, resample_kernel=(1, 3, 3, 1), bias=True, activate=True)
-        self.conv3 = ConvLayer(128, 128, 3, downsample=False, resample_kernel=(1, 3, 3, 1), bias=True, activate=True)
-        self.conv4 = ConvLayer(128, 256, 3, downsample=True, resample_kernel=(1, 3, 3, 1), bias=True, activate=True)
-        self.conv5 = ConvLayer(256, 256, 3, downsample=False, resample_kernel=(1, 3, 3, 1), bias=True, activate=True)
-        self.final_conv = ConvLayer(256, 1, 3, bias=True, activate=False)
-
-    def forward(self, x, return_feats=False):
-        feat = self.conv1(x)
-        feat = self.conv3(self.conv2(feat))
-        rlt_feats = []
-        if return_feats:
-            rlt_feats.append(feat.clone())
-        feat = self.conv5(self.conv4(feat))
-        if return_feats:
-            rlt_feats.append(feat.clone())
-        out = self.final_conv(feat)
-
-        if return_feats:
-            return out, rlt_feats
-        else:
-            return out, None

data/__init__.py

@@ -1,11 +0,0 @@
-import importlib
-from os import path as osp
-
-from basicsr.utils import scandir
-
-# automatically scan and import dataset modules for registry
-# scan all the files under the data folder with '_dataset' in file names
-data_folder = osp.dirname(osp.abspath(__file__))
-dataset_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(data_folder) if v.endswith('_dataset.py')]
-# import all the dataset modules
-_dataset_modules = [importlib.import_module(f'data.{file_name}') for file_name in dataset_filenames]

data/ffhq_degradation_dataset.py

@@ -1,213 +0,0 @@
-import cv2
-import math
-import numpy as np
-import os.path as osp
-import torch
-import torch.utils.data as data
-from torchvision.transforms.functional import (adjust_brightness, adjust_contrast, adjust_hue, adjust_saturation,
-                                               normalize)
-
-from basicsr.data import degradations as degradations
-from basicsr.data.data_util import paths_from_folder
-from basicsr.data.transforms import augment
-from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
-from basicsr.utils.registry import DATASET_REGISTRY
-
-
-@DATASET_REGISTRY.register()
-class FFHQDegradationDataset(data.Dataset):
-
-    def __init__(self, opt):
-        super(FFHQDegradationDataset, self).__init__()
-        self.opt = opt
-        # file client (io backend)
-        self.file_client = None
-        self.io_backend_opt = opt['io_backend']
-
-        self.gt_folder = opt['dataroot_gt']
-        self.mean = opt['mean']
-        self.std = opt['std']
-        self.out_size = opt['out_size']
-
-        self.crop_components = opt.get('crop_components', False)  # facial components
-        self.eye_enlarge_ratio = opt.get('eye_enlarge_ratio', 1)
-
-        if self.crop_components:
-            self.components_list = torch.load(opt.get('component_path'))
-
-        if self.io_backend_opt['type'] == 'lmdb':
-            self.io_backend_opt['db_paths'] = self.gt_folder
-            if not self.gt_folder.endswith('.lmdb'):
-                raise ValueError(f"'dataroot_gt' should end with '.lmdb', but received {self.gt_folder}")
-            with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
-                self.paths = [line.split('.')[0] for line in fin]
-        else:
-            self.paths = paths_from_folder(self.gt_folder)
-
-        # degradations
-        self.blur_kernel_size = opt['blur_kernel_size']
-        self.kernel_list = opt['kernel_list']
-        self.kernel_prob = opt['kernel_prob']
-        self.blur_sigma = opt['blur_sigma']
-        self.downsample_range = opt['downsample_range']
-        self.noise_range = opt['noise_range']
-        self.jpeg_range = opt['jpeg_range']
-
-        # color jitter
-        self.color_jitter_prob = opt.get('color_jitter_prob')
-        self.color_jitter_pt_prob = opt.get('color_jitter_pt_prob')
-        self.color_jitter_shift = opt.get('color_jitter_shift', 20)
-        # to gray
-        self.gray_prob = opt.get('gray_prob')
-
-        logger = get_root_logger()
-        logger.info(f'Blur: blur_kernel_size {self.blur_kernel_size}, '
-                    f'sigma: [{", ".join(map(str, self.blur_sigma))}]')
-        logger.info(f'Downsample: downsample_range [{", ".join(map(str, self.downsample_range))}]')
-        logger.info(f'Noise: [{", ".join(map(str, self.noise_range))}]')
-        logger.info(f'JPEG compression: [{", ".join(map(str, self.jpeg_range))}]')
-
-        if self.color_jitter_prob is not None:
-            logger.info(f'Use random color jitter. Prob: {self.color_jitter_prob}, '
-                        f'shift: {self.color_jitter_shift}')
-        if self.gray_prob is not None:
-            logger.info(f'Use random gray. Prob: {self.gray_prob}')
-
-        self.color_jitter_shift /= 255.
-
-    @staticmethod
-    def color_jitter(img, shift):
-        jitter_val = np.random.uniform(-shift, shift, 3).astype(np.float32)
-        img = img + jitter_val
-        img = np.clip(img, 0, 1)
-        return img
-
-    @staticmethod
-    def color_jitter_pt(img, brightness, contrast, saturation, hue):
-        fn_idx = torch.randperm(4)
-        for fn_id in fn_idx:
-            if fn_id == 0 and brightness is not None:
-                brightness_factor = torch.tensor(1.0).uniform_(brightness[0], brightness[1]).item()
-                img = adjust_brightness(img, brightness_factor)
-
-            if fn_id == 1 and contrast is not None:
-                contrast_factor = torch.tensor(1.0).uniform_(contrast[0], contrast[1]).item()
-                img = adjust_contrast(img, contrast_factor)
-
-            if fn_id == 2 and saturation is not None:
-                saturation_factor = torch.tensor(1.0).uniform_(saturation[0], saturation[1]).item()
-                img = adjust_saturation(img, saturation_factor)
-
-            if fn_id == 3 and hue is not None:
-                hue_factor = torch.tensor(1.0).uniform_(hue[0], hue[1]).item()
-                img = adjust_hue(img, hue_factor)
-        return img
-
-    def get_component_coordinates(self, index, status):
-        components_bbox = self.components_list[f'{index:08d}']
-        if status[0]:  # hflip
-            # exchange right and left eye
-            tmp = components_bbox['left_eye']
-            components_bbox['left_eye'] = components_bbox['right_eye']
-            components_bbox['right_eye'] = tmp
-            # modify the width coordinate
-            components_bbox['left_eye'][0] = self.out_size - components_bbox['left_eye'][0]
-            components_bbox['right_eye'][0] = self.out_size - components_bbox['right_eye'][0]
-            components_bbox['mouth'][0] = self.out_size - components_bbox['mouth'][0]
-
-        # get coordinates
-        locations = []
-        for part in ['left_eye', 'right_eye', 'mouth']:
-            mean = components_bbox[part][0:2]
-            half_len = components_bbox[part][2]
-            if 'eye' in part:
-                half_len *= self.eye_enlarge_ratio
-            loc = np.hstack((mean - half_len + 1, mean + half_len))
-            loc = torch.from_numpy(loc).float()
-            locations.append(loc)
-        return locations
-
-    def __getitem__(self, index):
-        if self.file_client is None:
-            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
-
-        # load gt image
-        gt_path = self.paths[index]
-        img_bytes = self.file_client.get(gt_path)
-        img_gt = imfrombytes(img_bytes, float32=True)
-
-        # random horizontal flip
-        img_gt, status = augment(img_gt, hflip=self.opt['use_hflip'], rotation=False, return_status=True)
-        h, w, _ = img_gt.shape
-
-        if self.crop_components:
-            locations = self.get_component_coordinates(index, status)
-            loc_left_eye, loc_right_eye, loc_mouth = locations
-
-        # ------------------------ generate lq image ------------------------ #
-        # blur
-        kernel = degradations.random_mixed_kernels(
-            self.kernel_list,
-            self.kernel_prob,
-            self.blur_kernel_size,
-            self.blur_sigma,
-            self.blur_sigma, [-math.pi, math.pi],
-            noise_range=None)
-        img_lq = cv2.filter2D(img_gt, -1, kernel)
-        # downsample
-        scale = np.random.uniform(self.downsample_range[0], self.downsample_range[1])
-        img_lq = cv2.resize(img_lq, (int(w // scale), int(h // scale)), interpolation=cv2.INTER_LINEAR)
-        # noise
-        if self.noise_range is not None:
-            img_lq = degradations.random_add_gaussian_noise(img_lq, self.noise_range)
-        # jpeg compression
-        if self.jpeg_range is not None:
-            img_lq = degradations.random_add_jpg_compression(img_lq, self.jpeg_range)
-
-        # resize to original size
-        img_lq = cv2.resize(img_lq, (w, h), interpolation=cv2.INTER_LINEAR)
-
-        # random color jitter (only for lq)
-        if self.color_jitter_prob is not None and (np.random.uniform() < self.color_jitter_prob):
-            img_lq = self.color_jitter(img_lq, self.color_jitter_shift)
-        # random to gray (only for lq)
-        if self.gray_prob and np.random.uniform() < self.gray_prob:
-            img_lq = cv2.cvtColor(img_lq, cv2.COLOR_BGR2GRAY)
-            img_lq = np.tile(img_lq[:, :, None], [1, 1, 3])
-            if self.opt.get('gt_gray'):
-                img_gt = cv2.cvtColor(img_gt, cv2.COLOR_BGR2GRAY)
-                img_gt = np.tile(img_gt[:, :, None], [1, 1, 3])
-
-        # BGR to RGB, HWC to CHW, numpy to tensor
-        img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
-
-        # random color jitter (pytorch version) (only for lq)
-        if self.color_jitter_pt_prob is not None and (np.random.uniform() < self.color_jitter_pt_prob):
-            brightness = self.opt.get('brightness', (0.5, 1.5))
-            contrast = self.opt.get('contrast', (0.5, 1.5))
-            saturation = self.opt.get('saturation', (0, 1.5))
-            hue = self.opt.get('hue', (-0.1, 0.1))
-            img_lq = self.color_jitter_pt(img_lq, brightness, contrast, saturation, hue)
-
-        # round and clip
-        img_lq = torch.clamp((img_lq * 255.0).round(), 0, 255) / 255.
-
-        # normalize
-        normalize(img_gt, self.mean, self.std, inplace=True)
-        normalize(img_lq, self.mean, self.std, inplace=True)
-
-        if self.crop_components:
-            return_dict = {
-                'lq': img_lq,
-                'gt': img_gt,
-                'gt_path': gt_path,
-                'loc_left_eye': loc_left_eye,
-                'loc_right_eye': loc_right_eye,
-                'loc_mouth': loc_mouth
-            }
-            return return_dict
-        else:
-            return {'lq': img_lq, 'gt': img_gt, 'gt_path': gt_path}
-
-    def __len__(self):
-        return len(self.paths)

experiments/pretrained_models/GFPGANv1.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:6db3a33dd00dd427b8a70a7e3c6244a5bcccb818736f4861ce1d609024a991de
-size 615378983

experiments/pretrained_models/README.md

@@ -1,7 +0,0 @@
-# Pre-trained Models and Other Data
-
-Download pre-trained models and other data. Put them in this folder.
-
-1. [Pretrained StyleGAN2 model: StyleGAN2_512_Cmul1_FFHQ_B12G4_scratch_800k.pth](https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/StyleGAN2_512_Cmul1_FFHQ_B12G4_scratch_800k.pth)
-1. [Component locations of FFHQ: FFHQ_eye_mouth_landmarks_512.pth](https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/FFHQ_eye_mouth_landmarks_512.pth)
-1. [A simple ArcFace model: arcface_resnet18.pth](https://github.com/TencentARC/GFPGAN/releases/download/v0.1.0/arcface_resnet18.pth)

inference_gfpgan_full.py

@@ -1,130 +0,0 @@
-import argparse
-import cv2
-import glob
-import numpy as np
-import os
-import torch
-from facexlib.utils.face_restoration_helper import FaceRestoreHelper
-from torchvision.transforms.functional import normalize
-
-from archs.gfpganv1_arch import GFPGANv1
-from basicsr.utils import img2tensor, imwrite, tensor2img
-
-
-def restoration(gfpgan,
-                face_helper,
-                img_path,
-                save_root,
-                has_aligned=False,
-                only_center_face=True,
-                suffix=None,
-                paste_back=False):
-    # read image
-    img_name = os.path.basename(img_path)
-    print(f'Processing {img_name} ...')
-    basename, _ = os.path.splitext(img_name)
-    input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
-    face_helper.clean_all()
-
-    if has_aligned:
-        input_img = cv2.resize(input_img, (512, 512))
-        face_helper.cropped_faces = [input_img]
-    else:
-        face_helper.read_image(input_img)
-        # get face landmarks for each face
-        face_helper.get_face_landmarks_5(only_center_face=only_center_face, pad_blur=False)
-        # align and warp each face
-        save_crop_path = os.path.join(save_root, 'cropped_faces', img_name)
-        face_helper.align_warp_face(save_crop_path)
-
-    # face restoration
-    for idx, cropped_face in enumerate(face_helper.cropped_faces):
-        # prepare data
-        cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
-        normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
-        cropped_face_t = cropped_face_t.unsqueeze(0).to('cuda')
-
-        try:
-            with torch.no_grad():
-                output = gfpgan(cropped_face_t, return_rgb=False)[0]
-                # convert to image
-                restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(-1, 1))
-        except RuntimeError as error:
-            print(f'\tFailed inference for GFPGAN: {error}.')
-            restored_face = cropped_face
-
-        restored_face = restored_face.astype('uint8')
-        face_helper.add_restored_face(restored_face)
-
-        if suffix is not None:
-            save_face_name = f'{basename}_{idx:02d}_{suffix}.png'
-        else:
-            save_face_name = f'{basename}_{idx:02d}.png'
-        save_restore_path = os.path.join(save_root, 'restored_faces', save_face_name)
-        imwrite(restored_face, save_restore_path)
-
-        # save cmp image
-        cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
-        imwrite(cmp_img, os.path.join(save_root, 'cmp', f'{basename}_{idx:02d}.png'))
-
-    if not has_aligned and paste_back:
-        face_helper.get_inverse_affine(None)
-        save_restore_path = os.path.join(save_root, 'restored_imgs', img_name)
-        # paste each restored face to the input image
-        face_helper.paste_faces_to_input_image(save_restore_path)
-
-
-if __name__ == '__main__':
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument('--upscale_factor', type=int, default=1)
-    parser.add_argument('--model_path', type=str, default='experiments/pretrained_models/GFPGANv1.pth')
-    parser.add_argument('--test_path', type=str, default='inputs/whole_imgs')
-    parser.add_argument('--suffix', type=str, default=None, help='Suffix of the restored faces')
-    parser.add_argument('--only_center_face', action='store_true')
-    parser.add_argument('--aligned', action='store_true')
-    parser.add_argument('--paste_back', action='store_true')
-
-    args = parser.parse_args()
-    if args.test_path.endswith('/'):
-        args.test_path = args.test_path[:-1]
-    save_root = 'results/'
-    os.makedirs(save_root, exist_ok=True)
-
-    # initialize the GFP-GAN
-    gfpgan = GFPGANv1(
-        out_size=512,
-        num_style_feat=512,
-        channel_multiplier=1,
-        decoder_load_path=None,
-        fix_decoder=True,
-        # for stylegan decoder
-        num_mlp=8,
-        input_is_latent=True,
-        different_w=True,
-        narrow=1,
-        sft_half=True)
-
-    gfpgan.to(device)
-    checkpoint = torch.load(args.model_path, map_location=lambda storage, loc: storage)
-    gfpgan.load_state_dict(checkpoint['params_ema'])
-    gfpgan.eval()
-
-    # initialize face helper
-    face_helper = FaceRestoreHelper(
-        args.upscale_factor, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png')
-
-    img_list = sorted(glob.glob(os.path.join(args.test_path, '*')))
-    for img_path in img_list:
-        restoration(
-            gfpgan,
-            face_helper,
-            img_path,
-            save_root,
-            has_aligned=args.aligned,
-            only_center_face=args.only_center_face,
-            suffix=args.suffix,
-            paste_back=args.paste_back)
-
-    print('Results are in the <results> folder.')

models/__init__.py

@@ -1,12 +0,0 @@
-import importlib
-from os import path as osp
-
-from basicsr.utils import scandir
-
-# automatically scan and import model modules for registry
-# scan all the files under the 'models' folder and collect files ending with
-# '_model.py'
-model_folder = osp.dirname(osp.abspath(__file__))
-model_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(model_folder) if v.endswith('_model.py')]
-# import all the model modules
-_model_modules = [importlib.import_module(f'models.{file_name}') for file_name in model_filenames]

models/gfpgan_model.py

@@ -1,562 +0,0 @@
-import math
-import os.path as osp
-import torch
-from collections import OrderedDict
-from torch.nn import functional as F
-from torchvision.ops import roi_align
-from tqdm import tqdm
-
-from basicsr.archs import build_network
-from basicsr.losses import build_loss
-from basicsr.losses.losses import r1_penalty
-from basicsr.metrics import calculate_metric
-from basicsr.models.base_model import BaseModel
-from basicsr.utils import get_root_logger, imwrite, tensor2img
-from basicsr.utils.registry import MODEL_REGISTRY
-
-
-@MODEL_REGISTRY.register()
-class GFPGANModel(BaseModel):
-    """GFPGAN model for <Towards real-world blind face restoratin with generative facial prior>"""
-
-    def __init__(self, opt):
-        super(GFPGANModel, self).__init__(opt)
-        self.idx = 0
-
-        # define network
-        self.net_g = build_network(opt['network_g'])
-        self.net_g = self.model_to_device(self.net_g)
-        self.print_network(self.net_g)
-
-        # load pretrained model
-        load_path = self.opt['path'].get('pretrain_network_g', None)
-        if load_path is not None:
-            param_key = self.opt['path'].get('param_key_g', 'params')
-            self.load_network(self.net_g, load_path, self.opt['path'].get('strict_load_g', True), param_key)
-
-        self.log_size = int(math.log(self.opt['network_g']['out_size'], 2))
-
-        if self.is_train:
-            self.init_training_settings()
-
-    def init_training_settings(self):
-        train_opt = self.opt['train']
-
-        # ----------- define net_d ----------- #
-        self.net_d = build_network(self.opt['network_d'])
-        self.net_d = self.model_to_device(self.net_d)
-        self.print_network(self.net_d)
-        # load pretrained model
-        load_path = self.opt['path'].get('pretrain_network_d', None)
-        if load_path is not None:
-            self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True))
-
-        # ----------- define net_g with Exponential Moving Average (EMA) ----------- #
-        # net_g_ema only used for testing on one GPU and saving
-        # There is no need to wrap with DistributedDataParallel
-        self.net_g_ema = build_network(self.opt['network_g']).to(self.device)
-        # load pretrained model
-        load_path = self.opt['path'].get('pretrain_network_g', None)
-        if load_path is not None:
-            self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')
-        else:
-            self.model_ema(0)  # copy net_g weight
-
-        self.net_g.train()
-        self.net_d.train()
-        self.net_g_ema.eval()
-
-        # ----------- facial components networks ----------- #
-        if ('network_d_left_eye' in self.opt and 'network_d_right_eye' in self.opt and 'network_d_mouth' in self.opt):
-            self.use_facial_disc = True
-        else:
-            self.use_facial_disc = False
-
-        if self.use_facial_disc:
-            # left eye
-            self.net_d_left_eye = build_network(self.opt['network_d_left_eye'])
-            self.net_d_left_eye = self.model_to_device(self.net_d_left_eye)
-            self.print_network(self.net_d_left_eye)
-            load_path = self.opt['path'].get('pretrain_network_d_left_eye')
-            if load_path is not None:
-                self.load_network(self.net_d_left_eye, load_path, True, 'params')
-            # right eye
-            self.net_d_right_eye = build_network(self.opt['network_d_right_eye'])
-            self.net_d_right_eye = self.model_to_device(self.net_d_right_eye)
-            self.print_network(self.net_d_right_eye)
-            load_path = self.opt['path'].get('pretrain_network_d_right_eye')
-            if load_path is not None:
-                self.load_network(self.net_d_right_eye, load_path, True, 'params')
-            # mouth
-            self.net_d_mouth = build_network(self.opt['network_d_mouth'])
-            self.net_d_mouth = self.model_to_device(self.net_d_mouth)
-            self.print_network(self.net_d_mouth)
-            load_path = self.opt['path'].get('pretrain_network_d_mouth')
-            if load_path is not None:
-                self.load_network(self.net_d_mouth, load_path, True, 'params')
-
-            self.net_d_left_eye.train()
-            self.net_d_right_eye.train()
-            self.net_d_mouth.train()
-
-            # ----------- define facial component gan loss ----------- #
-            self.cri_component = build_loss(train_opt['gan_component_opt']).to(self.device)
-
-        # ----------- define losses ----------- #
-        if train_opt.get('pixel_opt'):
-            self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
-        else:
-            self.cri_pix = None
-
-        if train_opt.get('perceptual_opt'):
-            self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)
-        else:
-            self.cri_perceptual = None
-
-        # L1 loss used in pyramid loss, component style loss and identity loss
-        self.cri_l1 = build_loss(train_opt['L1_opt']).to(self.device)
-
-        # gan loss (wgan)
-        self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device)
-
-        # ----------- define identity loss ----------- #
-        if 'network_identity' in self.opt:
-            self.use_identity = True
-        else:
-            self.use_identity = False
-
-        if self.use_identity:
-            # define identity network
-            self.network_identity = build_network(self.opt['network_identity'])
-            self.network_identity = self.model_to_device(self.network_identity)
-            self.print_network(self.network_identity)
-            load_path = self.opt['path'].get('pretrain_network_identity')
-            if load_path is not None:
-                self.load_network(self.network_identity, load_path, True, None)
-            self.network_identity.eval()
-            for param in self.network_identity.parameters():
-                param.requires_grad = False
-
-        # regularization weights
-        self.r1_reg_weight = train_opt['r1_reg_weight']  # for discriminator
-        self.net_d_iters = train_opt.get('net_d_iters', 1)
-        self.net_d_init_iters = train_opt.get('net_d_init_iters', 0)
-        self.net_d_reg_every = train_opt['net_d_reg_every']
-
-        # set up optimizers and schedulers
-        self.setup_optimizers()
-        self.setup_schedulers()
-
-    def setup_optimizers(self):
-        train_opt = self.opt['train']
-
-        # ----------- optimizer g ----------- #
-        net_g_reg_ratio = 1
-        normal_params = []
-        for _, param in self.net_g.named_parameters():
-            normal_params.append(param)
-        optim_params_g = [{  # add normal params first
-            'params': normal_params,
-            'lr': train_opt['optim_g']['lr']
-        }]
-        optim_type = train_opt['optim_g'].pop('type')
-        lr = train_opt['optim_g']['lr'] * net_g_reg_ratio
-        betas = (0**net_g_reg_ratio, 0.99**net_g_reg_ratio)
-        self.optimizer_g = self.get_optimizer(optim_type, optim_params_g, lr, betas=betas)
-        self.optimizers.append(self.optimizer_g)
-
-        # ----------- optimizer d ----------- #
-        net_d_reg_ratio = self.net_d_reg_every / (self.net_d_reg_every + 1)
-        normal_params = []
-        for _, param in self.net_d.named_parameters():
-            normal_params.append(param)
-        optim_params_d = [{  # add normal params first
-            'params': normal_params,
-            'lr': train_opt['optim_d']['lr']
-        }]
-        optim_type = train_opt['optim_d'].pop('type')
-        lr = train_opt['optim_d']['lr'] * net_d_reg_ratio
-        betas = (0**net_d_reg_ratio, 0.99**net_d_reg_ratio)
-        self.optimizer_d = self.get_optimizer(optim_type, optim_params_d, lr, betas=betas)
-        self.optimizers.append(self.optimizer_d)
-
-        if self.use_facial_disc:
-            # setup optimizers for facial component discriminators
-            optim_type = train_opt['optim_component'].pop('type')
-            lr = train_opt['optim_component']['lr']
-            # left eye
-            self.optimizer_d_left_eye = self.get_optimizer(
-                optim_type, self.net_d_left_eye.parameters(), lr, betas=(0.9, 0.99))
-            self.optimizers.append(self.optimizer_d_left_eye)
-            # right eye
-            self.optimizer_d_right_eye = self.get_optimizer(
-                optim_type, self.net_d_right_eye.parameters(), lr, betas=(0.9, 0.99))
-            self.optimizers.append(self.optimizer_d_right_eye)
-            # mouth
-            self.optimizer_d_mouth = self.get_optimizer(
-                optim_type, self.net_d_mouth.parameters(), lr, betas=(0.9, 0.99))
-            self.optimizers.append(self.optimizer_d_mouth)
-
-    def feed_data(self, data):
-        self.lq = data['lq'].to(self.device)
-        if 'gt' in data:
-            self.gt = data['gt'].to(self.device)
-
-        if 'loc_left_eye' in data:
-            # get facial component locations, shape (batch, 4)
-            self.loc_left_eyes = data['loc_left_eye']
-            self.loc_right_eyes = data['loc_right_eye']
-            self.loc_mouths = data['loc_mouth']
-
-            # uncomment to check data
-            # import torchvision
-            # if self.opt['rank'] == 0:
-            #     import os
-            #     os.makedirs('tmp/gt', exist_ok=True)
-            #     os.makedirs('tmp/lq', exist_ok=True)
-            #     print(self.idx)
-            #     torchvision.utils.save_image(
-            #         self.gt, f'tmp/gt/gt_{self.idx}.png', nrow=4, padding=2, normalize=True, range=(-1, 1))
-            #     torchvision.utils.save_image(
-            #         self.lq, f'tmp/lq/lq{self.idx}.png', nrow=4, padding=2, normalize=True, range=(-1, 1))
-            #     self.idx = self.idx + 1
-
-    def construct_img_pyramid(self):
-        pyramid_gt = [self.gt]
-        down_img = self.gt
-        for _ in range(0, self.log_size - 3):
-            down_img = F.interpolate(down_img, scale_factor=0.5, mode='bilinear', align_corners=False)
-            pyramid_gt.insert(0, down_img)
-        return pyramid_gt
-
-    def get_roi_regions(self, eye_out_size=80, mouth_out_size=120):
-        # hard code
-        face_ratio = int(self.opt['network_g']['out_size'] / 512)
-        eye_out_size *= face_ratio
-        mouth_out_size *= face_ratio
-
-        rois_eyes = []
-        rois_mouths = []
-        for b in range(self.loc_left_eyes.size(0)):  # loop for batch size
-            # left eye and right eye
-            img_inds = self.loc_left_eyes.new_full((2, 1), b)
-            bbox = torch.stack([self.loc_left_eyes[b, :], self.loc_right_eyes[b, :]], dim=0)  # shape: (2, 4)
-            rois = torch.cat([img_inds, bbox], dim=-1)  # shape: (2, 5)
-            rois_eyes.append(rois)
-            # mouse
-            img_inds = self.loc_left_eyes.new_full((1, 1), b)
-            rois = torch.cat([img_inds, self.loc_mouths[b:b + 1, :]], dim=-1)  # shape: (1, 5)
-            rois_mouths.append(rois)
-
-        rois_eyes = torch.cat(rois_eyes, 0).to(self.device)
-        rois_mouths = torch.cat(rois_mouths, 0).to(self.device)
-
-        # real images
-        all_eyes = roi_align(self.gt, boxes=rois_eyes, output_size=eye_out_size) * face_ratio
-        self.left_eyes_gt = all_eyes[0::2, :, :, :]
-        self.right_eyes_gt = all_eyes[1::2, :, :, :]
-        self.mouths_gt = roi_align(self.gt, boxes=rois_mouths, output_size=mouth_out_size) * face_ratio
-        # output
-        all_eyes = roi_align(self.output, boxes=rois_eyes, output_size=eye_out_size) * face_ratio
-        self.left_eyes = all_eyes[0::2, :, :, :]
-        self.right_eyes = all_eyes[1::2, :, :, :]
-        self.mouths = roi_align(self.output, boxes=rois_mouths, output_size=mouth_out_size) * face_ratio
-
-    def _gram_mat(self, x):
-        """Calculate Gram matrix.
-
-        Args:
-            x (torch.Tensor): Tensor with shape of (n, c, h, w).
-
-        Returns:
-            torch.Tensor: Gram matrix.
-        """
-        n, c, h, w = x.size()
-        features = x.view(n, c, w * h)
-        features_t = features.transpose(1, 2)
-        gram = features.bmm(features_t) / (c * h * w)
-        return gram
-
-    def gray_resize_for_identity(self, out, size=128):
-        out_gray = (0.2989 * out[:, 0, :, :] + 0.5870 * out[:, 1, :, :] + 0.1140 * out[:, 2, :, :])
-        out_gray = out_gray.unsqueeze(1)
-        out_gray = F.interpolate(out_gray, (size, size), mode='bilinear', align_corners=False)
-        return out_gray
-
-    def optimize_parameters(self, current_iter):
-        # optimize net_g
-        for p in self.net_d.parameters():
-            p.requires_grad = False
-        self.optimizer_g.zero_grad()
-
-        if self.use_facial_disc:
-            for p in self.net_d_left_eye.parameters():
-                p.requires_grad = False
-            for p in self.net_d_right_eye.parameters():
-                p.requires_grad = False
-            for p in self.net_d_mouth.parameters():
-                p.requires_grad = False
-
-        # image pyramid loss weight
-        if current_iter < self.opt['train'].get('remove_pyramid_loss', float('inf')):
-            pyramid_loss_weight = self.opt['train'].get('pyramid_loss_weight', 1)
-        else:
-            pyramid_loss_weight = 1e-12  # very small loss
-        if pyramid_loss_weight > 0:
-            self.output, out_rgbs = self.net_g(self.lq, return_rgb=True)
-            pyramid_gt = self.construct_img_pyramid()
-        else:
-            self.output, out_rgbs = self.net_g(self.lq, return_rgb=False)
-
-        # get roi-align regions
-        if self.use_facial_disc:
-            self.get_roi_regions(eye_out_size=80, mouth_out_size=120)
-
-        l_g_total = 0
-        loss_dict = OrderedDict()
-        if (current_iter % self.net_d_iters == 0 and current_iter > self.net_d_init_iters):
-            # pixel loss
-            if self.cri_pix:
-                l_g_pix = self.cri_pix(self.output, self.gt)
-                l_g_total += l_g_pix
-                loss_dict['l_g_pix'] = l_g_pix
-
-            # image pyramid loss
-            if pyramid_loss_weight > 0:
-                for i in range(0, self.log_size - 2):
-                    l_pyramid = self.cri_l1(out_rgbs[i], pyramid_gt[i]) * pyramid_loss_weight
-                    l_g_total += l_pyramid
-                    loss_dict[f'l_p_{2**(i+3)}'] = l_pyramid
-
-            # perceptual loss
-            if self.cri_perceptual:
-                l_g_percep, l_g_style = self.cri_perceptual(self.output, self.gt)
-                if l_g_percep is not None:
-                    l_g_total += l_g_percep
-                    loss_dict['l_g_percep'] = l_g_percep
-                if l_g_style is not None:
-                    l_g_total += l_g_style
-                    loss_dict['l_g_style'] = l_g_style
-
-            # gan loss
-            fake_g_pred = self.net_d(self.output)
-            l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
-            l_g_total += l_g_gan
-            loss_dict['l_g_gan'] = l_g_gan
-
-            # facial component loss
-            if self.use_facial_disc:
-                # left eye
-                fake_left_eye, fake_left_eye_feats = self.net_d_left_eye(self.left_eyes, return_feats=True)
-                l_g_gan = self.cri_component(fake_left_eye, True, is_disc=False)
-                l_g_total += l_g_gan
-                loss_dict['l_g_gan_left_eye'] = l_g_gan
-                # right eye
-                fake_right_eye, fake_right_eye_feats = self.net_d_right_eye(self.right_eyes, return_feats=True)
-                l_g_gan = self.cri_component(fake_right_eye, True, is_disc=False)
-                l_g_total += l_g_gan
-                loss_dict['l_g_gan_right_eye'] = l_g_gan
-                # mouth
-                fake_mouth, fake_mouth_feats = self.net_d_mouth(self.mouths, return_feats=True)
-                l_g_gan = self.cri_component(fake_mouth, True, is_disc=False)
-                l_g_total += l_g_gan
-                loss_dict['l_g_gan_mouth'] = l_g_gan
-
-                if self.opt['train'].get('comp_style_weight', 0) > 0:
-                    # get gt feat
-                    _, real_left_eye_feats = self.net_d_left_eye(self.left_eyes_gt, return_feats=True)
-                    _, real_right_eye_feats = self.net_d_right_eye(self.right_eyes_gt, return_feats=True)
-                    _, real_mouth_feats = self.net_d_mouth(self.mouths_gt, return_feats=True)
-
-                    def _comp_style(feat, feat_gt, criterion):
-                        return criterion(self._gram_mat(feat[0]), self._gram_mat(
-                            feat_gt[0].detach())) * 0.5 + criterion(
-                                self._gram_mat(feat[1]), self._gram_mat(feat_gt[1].detach()))
-
-                    # facial component style loss
-                    comp_style_loss = 0
-                    comp_style_loss += _comp_style(fake_left_eye_feats, real_left_eye_feats, self.cri_l1)
-                    comp_style_loss += _comp_style(fake_right_eye_feats, real_right_eye_feats, self.cri_l1)
-                    comp_style_loss += _comp_style(fake_mouth_feats, real_mouth_feats, self.cri_l1)
-                    comp_style_loss = comp_style_loss * self.opt['train']['comp_style_weight']
-                    l_g_total += comp_style_loss
-                    loss_dict['l_g_comp_style_loss'] = comp_style_loss
-
-            # identity loss
-            if self.use_identity:
-                identity_weight = self.opt['train']['identity_weight']
-                # get gray images and resize
-                out_gray = self.gray_resize_for_identity(self.output)
-                gt_gray = self.gray_resize_for_identity(self.gt)
-
-                identity_gt = self.network_identity(gt_gray).detach()
-                identity_out = self.network_identity(out_gray)
-                l_identity = self.cri_l1(identity_out, identity_gt) * identity_weight
-                l_g_total += l_identity
-                loss_dict['l_identity'] = l_identity
-
-            l_g_total.backward()
-            self.optimizer_g.step()
-
-        # EMA
-        self.model_ema(decay=0.5**(32 / (10 * 1000)))
-
-        # ----------- optimize net_d ----------- #
-        for p in self.net_d.parameters():
-            p.requires_grad = True
-        self.optimizer_d.zero_grad()
-        if self.use_facial_disc:
-            for p in self.net_d_left_eye.parameters():
-                p.requires_grad = True
-            for p in self.net_d_right_eye.parameters():
-                p.requires_grad = True
-            for p in self.net_d_mouth.parameters():
-                p.requires_grad = True
-            self.optimizer_d_left_eye.zero_grad()
-            self.optimizer_d_right_eye.zero_grad()
-            self.optimizer_d_mouth.zero_grad()
-
-        fake_d_pred = self.net_d(self.output.detach())
-        real_d_pred = self.net_d(self.gt)
-        l_d = self.cri_gan(real_d_pred, True, is_disc=True) + self.cri_gan(fake_d_pred, False, is_disc=True)
-        loss_dict['l_d'] = l_d
-        # In wgan, real_score should be positive and fake_score should benegative
-        loss_dict['real_score'] = real_d_pred.detach().mean()
-        loss_dict['fake_score'] = fake_d_pred.detach().mean()
-        l_d.backward()
-
-        if current_iter % self.net_d_reg_every == 0:
-            self.gt.requires_grad = True
-            real_pred = self.net_d(self.gt)
-            l_d_r1 = r1_penalty(real_pred, self.gt)
-            l_d_r1 = (self.r1_reg_weight / 2 * l_d_r1 * self.net_d_reg_every + 0 * real_pred[0])
-            loss_dict['l_d_r1'] = l_d_r1.detach().mean()
-            l_d_r1.backward()
-
-        self.optimizer_d.step()
-
-        if self.use_facial_disc:
-            # lefe eye
-            fake_d_pred, _ = self.net_d_left_eye(self.left_eyes.detach())
-            real_d_pred, _ = self.net_d_left_eye(self.left_eyes_gt)
-            l_d_left_eye = self.cri_component(
-                real_d_pred, True, is_disc=True) + self.cri_gan(
-                    fake_d_pred, False, is_disc=True)
-            loss_dict['l_d_left_eye'] = l_d_left_eye
-            l_d_left_eye.backward()
-            # right eye
-            fake_d_pred, _ = self.net_d_right_eye(self.right_eyes.detach())
-            real_d_pred, _ = self.net_d_right_eye(self.right_eyes_gt)
-            l_d_right_eye = self.cri_component(
-                real_d_pred, True, is_disc=True) + self.cri_gan(
-                    fake_d_pred, False, is_disc=True)
-            loss_dict['l_d_right_eye'] = l_d_right_eye
-            l_d_right_eye.backward()
-            # mouth
-            fake_d_pred, _ = self.net_d_mouth(self.mouths.detach())
-            real_d_pred, _ = self.net_d_mouth(self.mouths_gt)
-            l_d_mouth = self.cri_component(
-                real_d_pred, True, is_disc=True) + self.cri_gan(
-                    fake_d_pred, False, is_disc=True)
-            loss_dict['l_d_mouth'] = l_d_mouth
-            l_d_mouth.backward()
-
-            self.optimizer_d_left_eye.step()
-            self.optimizer_d_right_eye.step()
-            self.optimizer_d_mouth.step()
-
-        self.log_dict = self.reduce_loss_dict(loss_dict)
-
-    def test(self):
-        with torch.no_grad():
-            if hasattr(self, 'net_g_ema'):
-                self.net_g_ema.eval()
-                self.output, _ = self.net_g_ema(self.lq)
-            else:
-                logger = get_root_logger()
-                logger.warning('Do not have self.net_g_ema, use self.net_g.')
-                self.net_g.eval()
-                self.output, _ = self.net_g(self.lq)
-                self.net_g.train()
-
-    def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
-        if self.opt['rank'] == 0:
-            self.nondist_validation(dataloader, current_iter, tb_logger, save_img)
-
-    def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
-        dataset_name = dataloader.dataset.opt['name']
-        with_metrics = self.opt['val'].get('metrics') is not None
-        if with_metrics:
-            self.metric_results = {metric: 0 for metric in self.opt['val']['metrics'].keys()}
-        pbar = tqdm(total=len(dataloader), unit='image')
-
-        for idx, val_data in enumerate(dataloader):
-            img_name = osp.splitext(osp.basename(val_data['lq_path'][0]))[0]
-            self.feed_data(val_data)
-            self.test()
-
-            visuals = self.get_current_visuals()
-            sr_img = tensor2img([visuals['sr']], min_max=(-1, 1))
-            gt_img = tensor2img([visuals['gt']], min_max=(-1, 1))
-
-            if 'gt' in visuals:
-                gt_img = tensor2img([visuals['gt']], min_max=(-1, 1))
-                del self.gt
-            # tentative for out of GPU memory
-            del self.lq
-            del self.output
-            torch.cuda.empty_cache()
-
-            if save_img:
-                if self.opt['is_train']:
-                    save_img_path = osp.join(self.opt['path']['visualization'], img_name,
-                                             f'{img_name}_{current_iter}.png')
-                else:
-                    if self.opt['val']['suffix']:
-                        save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,
-                                                 f'{img_name}_{self.opt["val"]["suffix"]}.png')
-                    else:
-                        save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,
-                                                 f'{img_name}_{self.opt["name"]}.png')
-                imwrite(sr_img, save_img_path)
-
-            if with_metrics:
-                # calculate metrics
-                for name, opt_ in self.opt['val']['metrics'].items():
-                    metric_data = dict(img1=sr_img, img2=gt_img)
-                    self.metric_results[name] += calculate_metric(metric_data, opt_)
-            pbar.update(1)
-            pbar.set_description(f'Test {img_name}')
-        pbar.close()
-
-        if with_metrics:
-            for metric in self.metric_results.keys():
-                self.metric_results[metric] /= (idx + 1)
-
-            self._log_validation_metric_values(current_iter, dataset_name, tb_logger)
-
-    def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
-        log_str = f'Validation {dataset_name}\n'
-        for metric, value in self.metric_results.items():
-            log_str += f'\t # {metric}: {value:.4f}\n'
-        logger = get_root_logger()
-        logger.info(log_str)
-        if tb_logger:
-            for metric, value in self.metric_results.items():
-                tb_logger.add_scalar(f'metrics/{metric}', value, current_iter)
-
-    def get_current_visuals(self):
-        out_dict = OrderedDict()
-        out_dict['gt'] = self.gt.detach().cpu()
-        out_dict['sr'] = self.output.detach().cpu()
-        return out_dict
-
-    def save(self, epoch, current_iter):
-        self.save_network([self.net_g, self.net_g_ema], 'net_g', current_iter, param_key=['params', 'params_ema'])
-        self.save_network(self.net_d, 'net_d', current_iter)
-        # save component discriminators
-        if self.use_facial_disc:
-            self.save_network(self.net_d_left_eye, 'net_d_left_eye', current_iter)
-            self.save_network(self.net_d_right_eye, 'net_d_right_eye', current_iter)
-            self.save_network(self.net_d_mouth, 'net_d_mouth', current_iter)
-        self.save_training_state(epoch, current_iter)

requirements.txt

@@ -1,10 +0,0 @@
-facexlib
-lmdb
-numpy
-opencv-python
-pyyaml
-tb-nightly
-torch>=1.7
-torchvision
-tqdm
-yapf

setup.cfg

@@ -1,22 +0,0 @@
-[flake8]
-ignore =
-    # line break before binary operator (W503)
-    W503,
-    # line break after binary operator (W504)
-    W504,
-max-line-length=120
-
-[yapf]
-based_on_style = pep8
-column_limit = 120
-blank_line_before_nested_class_or_def = true
-split_before_expression_after_opening_paren = true
-
-[isort]
-line_length = 120
-multi_line_output = 0
-known_standard_library = pkg_resources,setuptools
-known_first_party = basicsr
-known_third_party = cv2,facexlib,numpy,torch,torchvision,tqdm
-no_lines_before = STDLIB,LOCALFOLDER
-default_section = THIRDPARTY

train.py

@@ -1,10 +0,0 @@
-import os.path as osp
-
-import archs  # noqa: F401
-import data  # noqa: F401
-import models  # noqa: F401
-from basicsr.train import train_pipeline
-
-if __name__ == '__main__':
-    root_path = osp.abspath(osp.join(__file__, osp.pardir))
-    train_pipeline(root_path)

train_gfpgan_v1.yml

@@ -1,210 +0,0 @@
-# general settings
-name: train_GFPGANv1_512
-model_type: GFPGANModel
-num_gpu: 4
-manual_seed: 0
-
-# dataset and data loader settings
-datasets:
-  train:
-    name: FFHQ
-    type: FFHQDegradationDataset
-    # dataroot_gt: datasets/ffhq/ffhq_512.lmdb
-    dataroot_gt: datasets/ffhq/ffhq_512
-    io_backend:
-      # type: lmdb
-      type: disk
-
-    use_hflip: true
-    mean: [0.5, 0.5, 0.5]
-    std: [0.5, 0.5, 0.5]
-    out_size: 512
-
-    blur_kernel_size: 41
-    kernel_list: ['iso', 'aniso']
-    kernel_prob: [0.5, 0.5]
-    blur_sigma: [0.1, 10]
-    downsample_range: [0.8, 8]
-    noise_range: [0, 20]
-    jpeg_range: [60, 100]
-
-    # color jitter and gray
-    color_jitter_prob: 0.3
-    color_jitter_shift: 20
-    color_jitter_pt_prob: 0.3
-    gray_prob: 0.01
-
-    crop_components: true
-    component_path: experiments/pretrained_models/FFHQ_eye_mouth_landmarks_512.pth
-    eye_enlarge_ratio: 1.4
-
-    # data loader
-    use_shuffle: true
-    num_worker_per_gpu: 6
-    batch_size_per_gpu: 3
-    dataset_enlarge_ratio: 100
-    prefetch_mode: ~
-
-  val:
-    # Please modify accordingly to use your own validation
-    # Or comment the val block if do not need validation during training
-    name: validation
-    type: PairedImageDataset
-    dataroot_lq: datasets/faces/validation/input
-    dataroot_gt: datasets/faces/validation/reference
-    io_backend:
-      type: disk
-    mean: [0.5, 0.5, 0.5]
-    std: [0.5, 0.5, 0.5]
-    scale: 1
-
-# network structures
-network_g:
-  type: GFPGANv1
-  out_size: 512
-  num_style_feat: 512
-  channel_multiplier: 1
-  resample_kernel: [1, 3, 3, 1]
-  decoder_load_path: experiments/pretrained_models/StyleGAN2_512_Cmul1_FFHQ_B12G4_scratch_800k.pth
-  fix_decoder: true
-  num_mlp: 8
-  lr_mlp: 0.01
-  input_is_latent: true
-  different_w: true
-  narrow: 1
-  sft_half: true
-
-network_d:
-  type: StyleGAN2Discriminator
-  out_size: 512
-  channel_multiplier: 1
-  resample_kernel: [1, 3, 3, 1]
-
-network_d_left_eye:
-  type: FacialComponentDiscriminator
-
-network_d_right_eye:
-  type: FacialComponentDiscriminator
-
-network_d_mouth:
-  type: FacialComponentDiscriminator
-
-network_identity:
-  type: ResNetArcFace
-  block: IRBlock
-  layers: [2, 2, 2, 2]
-  use_se: False
-
-# path
-path:
-  pretrain_network_g: ~
-  param_key_g: params_ema
-  strict_load_g: ~
-  pretrain_network_d: ~
-  pretrain_network_d_left_eye: ~
-  pretrain_network_d_right_eye: ~
-  pretrain_network_d_mouth: ~
-  pretrain_network_identity: experiments/pretrained_models/arcface_resnet18.pth
-  # resume
-  resume_state: ~
-  ignore_resume_networks: ['network_identity']
-
-# training settings
-train:
-  optim_g:
-    type: Adam
-    lr: !!float 2e-3
-  optim_d:
-    type: Adam
-    lr: !!float 2e-3
-  optim_component:
-    type: Adam
-    lr: !!float 2e-3
-
-  scheduler:
-    type: MultiStepLR
-    milestones: [600000, 700000]
-    gamma: 0.5
-
-  total_iter: 800000
-  warmup_iter: -1  # no warm up
-
-  # losses
-  # pixel loss
-  pixel_opt:
-    type: L1Loss
-    loss_weight: !!float 1e-1
-    reduction: mean
-  # L1 loss used in pyramid loss, component style loss and identity loss
-  L1_opt:
-    type: L1Loss
-    loss_weight: 1
-    reduction: mean
-
-  # image pyramid loss
-  pyramid_loss_weight: 1
-  remove_pyramid_loss: 50000
-  # perceptual loss (content and style losses)
-  perceptual_opt:
-    type: PerceptualLoss
-    layer_weights:
-      # before relu
-      'conv1_2': 0.1
-      'conv2_2': 0.1
-      'conv3_4': 1
-      'conv4_4': 1
-      'conv5_4': 1
-    vgg_type: vgg19
-    use_input_norm: true
-    perceptual_weight: !!float 1
-    style_weight: 50
-    range_norm: true
-    criterion: l1
-  # gan loss
-  gan_opt:
-    type: GANLoss
-    gan_type: wgan_softplus
-    loss_weight: !!float 1e-1
-  # r1 regularization for discriminator
-  r1_reg_weight: 10
-  # facial component loss
-  gan_component_opt:
-    type: GANLoss
-    gan_type: vanilla
-    real_label_val: 1.0
-    fake_label_val: 0.0
-    loss_weight: !!float 1
-  comp_style_weight: 200
-  # identity loss
-  identity_weight: 10
-
-  net_d_iters: 1
-  net_d_init_iters: 0
-  net_d_reg_every: 16
-
-# validation settings
-val:
-  val_freq: !!float 5e3
-  save_img: true
-
-  metrics:
-    psnr: # metric name, can be arbitrary
-      type: calculate_psnr
-      crop_border: 0
-      test_y_channel: false
-
-# logging settings
-logger:
-  print_freq: 100
-  save_checkpoint_freq: !!float 5e3
-  use_tb_logger: true
-  wandb:
-    project: ~
-    resume_id: ~
-
-# dist training settings
-dist_params:
-  backend: nccl
-  port: 29500
-
-find_unused_parameters: true