Real-ESRGAN

1
import argparse
2
import cv2
3
import glob
4
import os
5
from basicsr.archs.rrdbnet_arch import RRDBNet
6
from basicsr.utils.download_util import load_file_from_url
7

8
from realesrgan import RealESRGANer
9
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
10

11

12
def main():
13
    """Inference demo for Real-ESRGAN.
14
    """
15
    parser = argparse.ArgumentParser()
16
    parser.add_argument('-i', '--input', type=str, default='inputs', help='Input image or folder')
17
    parser.add_argument(
18
        '-n',
19
        '--model_name',
20
        type=str,
21
        default='RealESRGAN_x4plus',
22
        help=('Model names: RealESRGAN_x4plus | RealESRNet_x4plus | RealESRGAN_x4plus_anime_6B | RealESRGAN_x2plus | '
23
              'realesr-animevideov3 | realesr-general-x4v3'))
24
    parser.add_argument('-o', '--output', type=str, default='results', help='Output folder')
25
    parser.add_argument(
26
        '-dn',
27
        '--denoise_strength',
28
        type=float,
29
        default=0.5,
30
        help=('Denoise strength. 0 for weak denoise (keep noise), 1 for strong denoise ability. '
31
              'Only used for the realesr-general-x4v3 model'))
32
    parser.add_argument('-s', '--outscale', type=float, default=4, help='The final upsampling scale of the image')
33
    parser.add_argument(
34
        '--model_path', type=str, default=None, help='[Option] Model path. Usually, you do not need to specify it')
35
    parser.add_argument('--suffix', type=str, default='out', help='Suffix of the restored image')
36
    parser.add_argument('-t', '--tile', type=int, default=0, help='Tile size, 0 for no tile during testing')
37
    parser.add_argument('--tile_pad', type=int, default=10, help='Tile padding')
38
    parser.add_argument('--pre_pad', type=int, default=0, help='Pre padding size at each border')
39
    parser.add_argument('--face_enhance', action='store_true', help='Use GFPGAN to enhance face')
40
    parser.add_argument(
41
        '--fp32', action='store_true', help='Use fp32 precision during inference. Default: fp16 (half precision).')
42
    parser.add_argument(
43
        '--alpha_upsampler',
44
        type=str,
45
        default='realesrgan',
46
        help='The upsampler for the alpha channels. Options: realesrgan | bicubic')
47
    parser.add_argument(
48
        '--ext',
49
        type=str,
50
        default='auto',
51
        help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs')
52
    parser.add_argument(
53
        '-g', '--gpu-id', type=int, default=None, help='gpu device to use (default=None) can be 0,1,2 for multi-gpu')
54

55
    args = parser.parse_args()
56

57
    # determine models according to model names
58
    args.model_name = args.model_name.split('.')[0]
59
    if args.model_name == 'RealESRGAN_x4plus':  # x4 RRDBNet model
60
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
61
        netscale = 4
62
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
63
    elif args.model_name == 'RealESRNet_x4plus':  # x4 RRDBNet model
64
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
65
        netscale = 4
66
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
67
    elif args.model_name == 'RealESRGAN_x4plus_anime_6B':  # x4 RRDBNet model with 6 blocks
68
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
69
        netscale = 4
70
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
71
    elif args.model_name == 'RealESRGAN_x2plus':  # x2 RRDBNet model
72
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
73
        netscale = 2
74
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
75
    elif args.model_name == 'realesr-animevideov3':  # x4 VGG-style model (XS size)
76
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
77
        netscale = 4
78
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth']
79
    elif args.model_name == 'realesr-general-x4v3':  # x4 VGG-style model (S size)
80
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
81
        netscale = 4
82
        file_url = [
83
            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
84
            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
85
        ]
86

87
    # determine model paths
88
    if args.model_path is not None:
89
        model_path = args.model_path
90
    else:
91
        model_path = os.path.join('weights', args.model_name + '.pth')
92
        if not os.path.isfile(model_path):
93
            ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
94
            for url in file_url:
95
                # model_path will be updated
96
                model_path = load_file_from_url(
97
                    url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
98

99
    # use dni to control the denoise strength
100
    dni_weight = None
101
    if args.model_name == 'realesr-general-x4v3' and args.denoise_strength != 1:
102
        wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
103
        model_path = [model_path, wdn_model_path]
104
        dni_weight = [args.denoise_strength, 1 - args.denoise_strength]
105

106
    # restorer
107
    upsampler = RealESRGANer(
108
        scale=netscale,
109
        model_path=model_path,
110
        dni_weight=dni_weight,
111
        model=model,
112
        tile=args.tile,
113
        tile_pad=args.tile_pad,
114
        pre_pad=args.pre_pad,
115
        half=not args.fp32,
116
        gpu_id=args.gpu_id)
117

118
    if args.face_enhance:  # Use GFPGAN for face enhancement
119
        from gfpgan import GFPGANer
120
        face_enhancer = GFPGANer(
121
            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
122
            upscale=args.outscale,
123
            arch='clean',
124
            channel_multiplier=2,
125
            bg_upsampler=upsampler)
126
    os.makedirs(args.output, exist_ok=True)
127

128
    if os.path.isfile(args.input):
129
        paths = [args.input]
130
    else:
131
        paths = sorted(glob.glob(os.path.join(args.input, '*')))
132

133
    for idx, path in enumerate(paths):
134
        imgname, extension = os.path.splitext(os.path.basename(path))
135
        print('Testing', idx, imgname)
136

137
        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
138
        if len(img.shape) == 3 and img.shape[2] == 4:
139
            img_mode = 'RGBA'
140
        else:
141
            img_mode = None
142

143
        try:
144
            if args.face_enhance:
145
                _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
146
            else:
147
                output, _ = upsampler.enhance(img, outscale=args.outscale)
148
        except RuntimeError as error:
149
            print('Error', error)
150
            print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
151
        else:
152
            if args.ext == 'auto':
153
                extension = extension[1:]
154
            else:
155
                extension = args.ext
156
            if img_mode == 'RGBA':  # RGBA images should be saved in png format
157
                extension = 'png'
158
            if args.suffix == '':
159
                save_path = os.path.join(args.output, f'{imgname}.{extension}')
160
            else:
161
                save_path = os.path.join(args.output, f'{imgname}_{args.suffix}.{extension}')
162
            cv2.imwrite(save_path, output)
163

164

165
if __name__ == '__main__':
166
    main()
167