diff --git a/README.md b/README.md index 25c066ec44741db933b3c06cce0d7d79a05b0ef9..bab238475e4490fd97bc0d2b5abf1956465fce03 100644 --- a/README.md +++ b/README.md @@ -123,6 +123,8 @@ GAN-Generative Adversarial Network, was praised by "the Father of Convolutional * [InvDN](./docs/en_US/tutorials/invdn.md) * [AOT-GAN](./docs/en_US/tutorials/aotgan.md) * [NAFNet](./docs/en_US/tutorials/nafnet.md) +* [GFPGan](./docs/en_US/tutorials/gfpgan.md) +* [GPEN](./docs/en_US/tutorials/gpen.md) ## Composite Application diff --git a/README_cn.md b/README_cn.md index 1a2a5b848f3058fb23aeb00dd0fd73700dfdc2e8..14118607f36e414651d2c94bf1bd2026ecaf386c 100644 --- a/README_cn.md +++ b/README_cn.md @@ -132,6 +132,7 @@ GAN--生成对抗网络，被“卷积网络之父”**Yann LeCun（杨立昆） * 人脸编码：[Pixel2Style2Pixel](./docs/zh_CN/tutorials/pixel2style2pixel.md) * 人脸增强：[FaceEnhancement](./docs/zh_CN/tutorials/face_enhancement.md) * 人脸解析：[FaceParsing](./docs/zh_CN/tutorials/face_parse.md) + * 盲人脸修复：[GFPGan](./docs/zh_CN/tutorials/gfpgan.md)、[GPEN](./docs/zh_CN/tutorials/gpen.md) * 分辨率提升 * 单张图片超分：[Single Image Super Resolution(SISR)](./docs/zh_CN/tutorials/single_image_super_resolution.md) * 包含模型：RealSR、ESRGAN、LESRCNN、PAN、DRN diff --git a/configs/gpen_256_ffhq.yaml b/configs/gpen_256_ffhq.yaml index a5e0936630bc96fe2fb252843f59046753c2b469..bead76a995becfe788380f7b65d2c3bc7ae58f19 100644 --- a/configs/gpen_256_ffhq.yaml +++ b/configs/gpen_256_ffhq.yaml @@ -6,7 +6,7 @@ find_unused_parameters: True model: name: GPENModel generator: - name: GPEN + name: GPENGenerator size: 256 style_dim: 512 n_mlp: 8 diff --git a/docs/en_US/tutorials/GFPGAN.md b/docs/en_US/tutorials/gfpgan.md similarity index 100% rename from docs/en_US/tutorials/GFPGAN.md rename to docs/en_US/tutorials/gfpgan.md diff --git a/docs/zh_CN/tutorials/face_enhancement.md b/docs/zh_CN/tutorials/face_enhancement.md index 0d3c3e166879f66be8cd6ec9972d2f053c28ccb5..fb14b6130d1b3a7734aca5e5d87e7a2bc5d1efb9 100644 --- a/docs/zh_CN/tutorials/face_enhancement.md +++ b/docs/zh_CN/tutorials/face_enhancement.md @@ -20,9 +20,9 @@ img = faceenhancer.enhance_from_image(img) 注意：请将图片转为float类型输入，目前不支持int8类型 -### 训练(TODO) +### 训练 -未来还将添加训练脚本方便用户训练出更多类型的 GPEN 人脸增强。 +[详见](../../zh_CN/tutorials/gpen.md) ## 人脸增强结果展示 diff --git a/docs/zh_CN/tutorials/GFPGAN.md b/docs/zh_CN/tutorials/gfpgan.md similarity index 100% rename from docs/zh_CN/tutorials/GFPGAN.md rename to docs/zh_CN/tutorials/gfpgan.md diff --git a/ppgan/apps/gpen_predictor.py b/ppgan/apps/gpen_predictor.py index 95f9811b12cca93228aeab5cb062d81beb6664b3..b0648a1724d25f423ec4ea9c99c42cdfca7d6c36 100644 --- a/ppgan/apps/gpen_predictor.py +++ b/ppgan/apps/gpen_predictor.py @@ -21,7 +21,7 @@ import sys sys.path.append(".") from .base_predictor import BasePredictor from ppgan.datasets.gpen_dataset import GFPGAN_degradation -from ppgan.models.generators import GPEN +from ppgan.models.generators import GPENGenerator from ppgan.metrics.fid import FID from ppgan.utils.download import get_path_from_url import cv2 @@ -103,7 +103,7 @@ class GPENPredictor(BasePredictor): checkpoint = paddle.load(weight_path) warnings.filterwarnings("always") - self.generator = GPEN(size, style_dim, n_mlp, channel_multiplier, + self.generator = GPENGenerator(size, style_dim, n_mlp, channel_multiplier, narrow) self.generator.set_state_dict(checkpoint) self.generator.eval() diff --git a/ppgan/models/generators/__init__.py b/ppgan/models/generators/__init__.py index 7b007075fbe26fe0afe5e784c7789e50a0f08ded..76af90ac8c5462039d8c00260ee06858c1e5caeb 100644 --- a/ppgan/models/generators/__init__.py +++ b/ppgan/models/generators/__init__.py @@ -41,7 +41,7 @@ from .msvsr import MSVSR from .generator_singan import SinGANGenerator from .rcan import RCAN from .prenet import PReNet -from .gpen import GPEN +from .generator_gpen import GPENGenerator from .swinir import SwinIR from .gfpganv1_clean_arch import GFPGANv1Clean from .gfpganv1_arch import GFPGANv1, StyleGAN2DiscriminatorGFPGAN diff --git a/ppgan/models/generators/generator_gpen.py b/ppgan/models/generators/generator_gpen.py index de5119edcb7a9bebf6f086fb3a1f97926ed1ec26..d73690bfd07a8d98cd7b10eb26d7f01ca326d8af 100644 --- a/ppgan/models/generators/generator_gpen.py +++ b/ppgan/models/generators/generator_gpen.py @@ -18,10 +18,13 @@ import math import random +import itertools import paddle import paddle.nn as nn import paddle.nn.functional as F +from ppgan.models.generators.builder import GENERATORS +from ppgan.models.discriminators.discriminator_styleganv2 import ConvLayer from ppgan.modules.equalized import EqualLinear_gpen as EqualLinear from ppgan.modules.fused_act import FusedLeakyReLU from ppgan.modules.upfirdn2d import Upfirdn2dUpsample, Upfirdn2dBlur @@ -451,3 +454,85 @@ class StyleGANv2Generator(nn.Layer): else: return image, None + +@GENERATORS.register() +class GPENGenerator(nn.Layer): + + def __init__( + self, + size, + style_dim, + n_mlp, + channel_multiplier=2, + narrow=1, + blur_kernel=[1, 3, 3, 1], + lr_mlp=0.01, + is_concat=True, + ): + super(GPENGenerator, self).__init__() + channels = { + 4: int(512 * narrow), + 8: int(512 * narrow), + 16: int(512 * narrow), + 32: int(512 * narrow), + 64: int(256 * channel_multiplier * narrow), + 128: int(128 * channel_multiplier * narrow), + 256: int(64 * channel_multiplier * narrow), + 512: int(32 * channel_multiplier * narrow), + 1024: int(16 * channel_multiplier * narrow), + 2048: int(8 * channel_multiplier * narrow) + } + self.log_size = int(math.log(size, 2)) + self.generator = StyleGANv2Generator( + size, + style_dim, + n_mlp, + channel_multiplier=channel_multiplier, + narrow=narrow, + blur_kernel=blur_kernel, + lr_mlp=lr_mlp, + is_concat=is_concat) + + conv = [ConvLayer(3, channels[size], 1)] + self.ecd0 = nn.Sequential(*conv) + in_channel = channels[size] + + self.names = ['ecd%d' % i for i in range(self.log_size - 1)] + for i in range(self.log_size, 2, -1): + out_channel = channels[2**(i - 1)] + conv = [ConvLayer(in_channel, out_channel, 3, downsample=True)] + setattr(self, self.names[self.log_size - i + 1], + nn.Sequential(*conv)) + in_channel = out_channel + self.final_linear = nn.Sequential( + EqualLinear(channels[4] * 4 * 4, + style_dim, + activation='fused_lrelu')) + + def forward( + self, + inputs, + return_latents=False, + inject_index=None, + truncation=1, + truncation_latent=None, + input_is_latent=False, + ): + noise = [] + for i in range(self.log_size - 1): + ecd = getattr(self, self.names[i]) + inputs = ecd(inputs) + noise.append(inputs) + inputs = inputs.reshape([inputs.shape[0], -1]) + outs = self.final_linear(inputs) + noise = list( + itertools.chain.from_iterable( + itertools.repeat(x, 2) for x in noise))[::-1] + outs = self.generator([outs], + return_latents, + inject_index, + truncation, + truncation_latent, + input_is_latent, + noise=noise[1:]) + return outs diff --git a/ppgan/models/generators/gpen.py b/ppgan/models/generators/gpen.py index 34a7913a4e7fd142770419c5c5d36b68dfb8c8e6..bcf34d4c2938fdf80e3bd811a6178ce6a63320de 100644 --- a/ppgan/models/generators/gpen.py +++ b/ppgan/models/generators/gpen.py @@ -14,16 +14,14 @@ # code was heavily based on code was heavily based on https://github.com/yangxy/GPEN -import itertools -import paddle.nn as nn import math +import paddle +import paddle.nn as nn from ppgan.models.generators.builder import GENERATORS -from ppgan.modules.equalized import EqualLinear_gpen as EqualLinear -from ppgan.models.generators.generator_gpen import StyleGANv2Generator +from ppgan.models.generators import StyleGANv2Generator from ppgan.models.discriminators.discriminator_styleganv2 import ConvLayer +from ppgan.modules.equalized import EqualLinear - -@GENERATORS.register() class GPEN(nn.Layer): def __init__( @@ -32,23 +30,21 @@ class GPEN(nn.Layer): style_dim, n_mlp, channel_multiplier=2, - narrow=1, blur_kernel=[1, 3, 3, 1], lr_mlp=0.01, is_concat=True, ): super(GPEN, self).__init__() channels = { - 4: int(512 * narrow), - 8: int(512 * narrow), - 16: int(512 * narrow), - 32: int(512 * narrow), - 64: int(256 * channel_multiplier * narrow), - 128: int(128 * channel_multiplier * narrow), - 256: int(64 * channel_multiplier * narrow), - 512: int(32 * channel_multiplier * narrow), - 1024: int(16 * channel_multiplier * narrow), - 2048: int(8 * channel_multiplier * narrow) + 4: 512, + 8: 512, + 16: 512, + 32: 512, + 64: 256 * channel_multiplier, + 128: 128 * channel_multiplier, + 256: 64 * channel_multiplier, + 512: 32 * channel_multiplier, + 1024: 16 * channel_multiplier, } self.log_size = int(math.log(size, 2)) self.generator = StyleGANv2Generator( @@ -56,7 +52,6 @@ class GPEN(nn.Layer): style_dim, n_mlp, channel_multiplier=channel_multiplier, - narrow=narrow, blur_kernel=blur_kernel, lr_mlp=lr_mlp, is_concat=is_concat) @@ -85,7 +80,7 @@ class GPEN(nn.Layer): truncation=1, truncation_latent=None, input_is_latent=False, - ): + ): noise = [] for i in range(self.log_size - 1): ecd = getattr(self, self.names[i]) @@ -93,14 +88,7 @@ class GPEN(nn.Layer): noise.append(inputs) inputs = inputs.reshape([inputs.shape[0], -1]) outs = self.final_linear(inputs) - noise = list( - itertools.chain.from_iterable( - itertools.repeat(x, 2) for x in noise))[::-1] - outs = self.generator([outs], - return_latents, - inject_index, - truncation, - truncation_latent, - input_is_latent, - noise=noise[1:]) + outs = self.generator([outs], return_latents, inject_index, truncation, + truncation_latent, input_is_latent, + noise=noise[::-1]) return outs