import argparse
import cv2
import glob
import os
from basicsr.archs.rrdbnet_arch import RRDBNet
import time

from realesrgan import RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact


def main():
    """Inference demo for Real-ESRGAN.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', type=str, default='inputs', help='Input image or folder')
    parser.add_argument(
        '-n',
        '--model_name',
        type=str,
        default='RealESRGAN_x4plus',
        help=('Model names: RealESRGAN_x4plus | RealESRNet_x4plus | RealESRGAN_x4plus_anime_6B | RealESRGAN_x2plus'
              'RealESRGANv2-anime-xsx2 | RealESRGANv2-animevideo-xsx2-nousm | RealESRGANv2-animevideo-xsx2'
              'RealESRGANv2-anime-xsx4 | RealESRGANv2-animevideo-xsx4-nousm | RealESRGANv2-animevideo-xsx4'))
    parser.add_argument('-o', '--output', type=str, default='results', help='Output folder')
    parser.add_argument('-s', '--outscale', type=float, default=4, help='The final upsampling scale of the image')
    parser.add_argument('--suffix', type=str, default='Realesrgan-4x', help='Suffix of the restored image')
    parser.add_argument('-t', '--tile', type=int, default=0, help='Tile size, 0 for no tile during testing')
    parser.add_argument('--tile_pad', type=int, default=10, help='Tile padding')
    parser.add_argument('--pre_pad', type=int, default=0, help='Pre padding size at each border')
    parser.add_argument('--face_enhance', action='store_true', help='Use GFPGAN to enhance face')
    parser.add_argument('--half', action='store_true', help='Use half precision during inference')
    parser.add_argument(
        '--alpha_upsampler',
        type=str,
        default='realesrgan',
        help='The upsampler for the alpha channels. Options: realesrgan | bicubic')
    parser.add_argument(
        '--ext',
        type=str,
        default='auto',
        help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs')
    args = parser.parse_args()

    # determine models according to model names
    args.model_name = args.model_name.split('.')[0]
    if args.model_name in ['RealESRGAN_x4plus', 'RealESRNet_x4plus']:  # x4 RRDBNet model
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
    elif args.model_name in ['RealESRGAN_x4plus_anime_6B']:  # x4 RRDBNet model with 6 blocks
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
        netscale = 4
    elif args.model_name in ['RealESRGAN_x2plus']:  # x2 RRDBNet model
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
        netscale = 2
    elif args.model_name in [
            'RealESRGANv2-anime-xsx2', 'RealESRGANv2-animevideo-xsx2-nousm', 'RealESRGANv2-animevideo-xsx2'
    ]:  # x2 VGG-style model (XS size)
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')
        netscale = 2
    elif args.model_name in [
            'RealESRGANv2-anime-xsx4', 'RealESRGANv2-animevideo-xsx4-nousm', 'RealESRGANv2-animevideo-xsx4'
    ]:  # x4 VGG-style model (XS size)
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
        netscale = 4
    else:
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
        netscale = 4

    # determine model paths
    model_path = os.path.join('experiments/pretrained_models', args.model_name + '.pth')
    if not os.path.isfile(model_path):
        model_path = os.path.join('realesrgan/weights', args.model_name + '.pth')
    if not os.path.isfile(model_path):
        raise ValueError(f'Model {args.model_name} does not exist.')

    # restorer
    upsampler = RealESRGANer(
        scale=netscale,
        model_path=model_path,
        model=model,
        tile=args.tile,
        tile_pad=args.tile_pad,
        pre_pad=args.pre_pad,
        half=args.half)

    if args.face_enhance:  # Use GFPGAN for face enhancement
        from gfpgan import GFPGANer
        face_enhancer = GFPGANer(
            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth',
            upscale=args.outscale,
            arch='clean',
            channel_multiplier=2,
            bg_upsampler=upsampler)
    os.makedirs(args.output, exist_ok=True)

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

    for idx, path in enumerate(paths):
        startTime = time.perf_counter()
        imgname, extension = os.path.splitext(os.path.basename(path))
        

        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
        
        if len(img.shape) == 3 and img.shape[2] == 4:
            img_mode = 'RGBA'
        else:
            img_mode = None
        if args.ext == 'auto':
            extension = "png"
        else:
            extension = args.ext
        if img_mode == 'RGBA':  # RGBA images should be saved in png format
            extension = 'png'
        save_path = os.path.join(args.output, f'{imgname}-{args.suffix}.{extension}')
        if os.path.exists(save_path):
            continue
        
        try:
            if args.face_enhance:
                _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
            else:
                output, _ = upsampler.enhance(img, outscale=args.outscale)
        except RuntimeError as error:
            print('Error', error)
            print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
        else:
            cv2.imwrite(save_path, output)
            print(f'NO.{idx}, {imgname} is done, used {round((time.perf_counter() - startTime), 4)} seconds')


if __name__ == '__main__':
    main()

[((275, 300), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (298, 300), False, 'import argparse\n'), ((3590, 3661), 'os.path.join', 'os.path.join', (['"""experiments/pretrained_models"""', "(args.model_name + '.pth')"], {}), "('experiments/pretrained_models', args.model_name + '.pth')\n", (3602, 3661), False, 'import os\n'), ((3923, 4070), 'realesrgan.RealESRGANer', 'RealESRGANer', ([], {'scale': 'netscale', 'model_path': 'model_path', 'model': 'model', 'tile': 'args.tile', 'tile_pad': 'args.tile_pad', 'pre_pad': 'args.pre_pad', 'half': 'args.half'}), '(scale=netscale, model_path=model_path, model=model, tile=args.\n    tile, tile_pad=args.tile_pad, pre_pad=args.pre_pad, half=args.half)\n', (3935, 4070), False, 'from realesrgan import RealESRGANer\n'), ((4504, 4543), 'os.makedirs', 'os.makedirs', (['args.output'], {'exist_ok': '(True)'}), '(args.output, exist_ok=True)\n', (4515, 4543), False, 'import os\n'), ((4552, 4578), 'os.path.isfile', 'os.path.isfile', (['args.input'], {}), '(args.input)\n', (4566, 4578), False, 'import os\n'), ((2267, 2358), 'basicsr.archs.rrdbnet_arch.RRDBNet', 'RRDBNet', ([], {'num_in_ch': '(3)', 'num_out_ch': '(3)', 'num_feat': '(64)', 'num_block': '(23)', 'num_grow_ch': '(32)', 'scale': '(4)'}), '(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=\n    32, scale=4)\n', (2274, 2358), False, 'from basicsr.archs.rrdbnet_arch import RRDBNet\n'), ((3673, 3699), 'os.path.isfile', 'os.path.isfile', (['model_path'], {}), '(model_path)\n', (3687, 3699), False, 'import os\n'), ((3722, 3782), 'os.path.join', 'os.path.join', (['"""realesrgan/weights"""', "(args.model_name + '.pth')"], {}), "('realesrgan/weights', args.model_name + '.pth')\n", (3734, 3782), False, 'import os\n'), ((3794, 3820), 'os.path.isfile', 'os.path.isfile', (['model_path'], {}), '(model_path)\n', (3808, 3820), False, 'import os\n'), ((4245, 4452), 'gfpgan.GFPGANer', 'GFPGANer', ([], {'model_path': '"""https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth"""', 'upscale': 'args.outscale', 'arch': '"""clean"""', 'channel_multiplier': '(2)', 'bg_upsampler': 'upsampler'}), "(model_path=\n    'https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth'\n    , upscale=args.outscale, arch='clean', channel_multiplier=2,\n    bg_upsampler=upsampler)\n", (4253, 4452), False, 'from gfpgan import GFPGANer\n'), ((4744, 4763), 'time.perf_counter', 'time.perf_counter', ([], {}), '()\n', (4761, 4763), False, 'import time\n'), ((4858, 4896), 'cv2.imread', 'cv2.imread', (['path', 'cv2.IMREAD_UNCHANGED'], {}), '(path, cv2.IMREAD_UNCHANGED)\n', (4868, 4896), False, 'import cv2\n'), ((5266, 5331), 'os.path.join', 'os.path.join', (['args.output', 'f"""{imgname}-{args.suffix}.{extension}"""'], {}), "(args.output, f'{imgname}-{args.suffix}.{extension}')\n", (5278, 5331), False, 'import os\n'), ((5343, 5368), 'os.path.exists', 'os.path.exists', (['save_path'], {}), '(save_path)\n', (5357, 5368), False, 'import os\n'), ((2485, 2574), 'basicsr.archs.rrdbnet_arch.RRDBNet', 'RRDBNet', ([], {'num_in_ch': '(3)', 'num_out_ch': '(3)', 'num_feat': '(64)', 'num_block': '(6)', 'num_grow_ch': '(32)', 'scale': '(4)'}), '(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32,\n    scale=4)\n', (2492, 2574), False, 'from basicsr.archs.rrdbnet_arch import RRDBNet\n'), ((4810, 4832), 'os.path.basename', 'os.path.basename', (['path'], {}), '(path)\n', (4826, 4832), False, 'import os\n'), ((5855, 5885), 'cv2.imwrite', 'cv2.imwrite', (['save_path', 'output'], {}), '(save_path, output)\n', (5866, 5885), False, 'import cv2\n'), ((2679, 2770), 'basicsr.archs.rrdbnet_arch.RRDBNet', 'RRDBNet', ([], {'num_in_ch': '(3)', 'num_out_ch': '(3)', 'num_feat': '(64)', 'num_block': '(23)', 'num_grow_ch': '(32)', 'scale': '(2)'}), '(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=\n    32, scale=2)\n', (2686, 2770), False, 'from basicsr.archs.rrdbnet_arch import RRDBNet\n'), ((4652, 4681), 'os.path.join', 'os.path.join', (['args.input', '"""*"""'], {}), "(args.input, '*')\n", (4664, 4681), False, 'import os\n'), ((2980, 3081), 'realesrgan.archs.srvgg_arch.SRVGGNetCompact', 'SRVGGNetCompact', ([], {'num_in_ch': '(3)', 'num_out_ch': '(3)', 'num_feat': '(64)', 'num_conv': '(16)', 'upscale': '(2)', 'act_type': '"""prelu"""'}), "(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16,\n    upscale=2, act_type='prelu')\n", (2995, 3081), False, 'from realesrgan.archs.srvgg_arch import SRVGGNetCompact\n'), ((3292, 3393), 'realesrgan.archs.srvgg_arch.SRVGGNetCompact', 'SRVGGNetCompact', ([], {'num_in_ch': '(3)', 'num_out_ch': '(3)', 'num_feat': '(64)', 'num_conv': '(16)', 'upscale': '(4)', 'act_type': '"""prelu"""'}), "(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16,\n    upscale=4, act_type='prelu')\n", (3307, 3393), False, 'from realesrgan.archs.srvgg_arch import SRVGGNetCompact\n'), ((3437, 3526), 'basicsr.archs.rrdbnet_arch.RRDBNet', 'RRDBNet', ([], {'num_in_ch': '(3)', 'num_out_ch': '(3)', 'num_feat': '(64)', 'num_block': '(6)', 'num_grow_ch': '(32)', 'scale': '(4)'}), '(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32,\n    scale=4)\n', (3444, 3526), False, 'from basicsr.archs.rrdbnet_arch import RRDBNet\n'), ((5948, 5967), 'time.perf_counter', 'time.perf_counter', ([], {}), '()\n', (5965, 5967), False, 'import time\n')]