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()

[((15, 13, 15, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((71, 17, 71, 88), 'os.path.join', 'os.path.join', ({(71, 30, 71, 61): '"""experiments/pretrained_models"""', (71, 63, 71, 87): "args.model_name + '.pth'"}, {}), "('experiments/pretrained_models', args.model_name + '.pth')", False, 'import os\n'), ((78, 16, 85, 23), 'realesrgan.RealESRGANer', 'RealESRGANer', (), '', False, 'from realesrgan import RealESRGANer\n'), ((95, 4, 95, 43), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((97, 7, 97, 33), 'os.path.isfile', 'os.path.isfile', ({(97, 22, 97, 32): 'args.input'}, {}), '(args.input)', False, 'import os\n'), ((48, 16, 48, 102), 'basicsr.archs.rrdbnet_arch.RRDBNet', 'RRDBNet', (), '', False, 'from basicsr.archs.rrdbnet_arch import RRDBNet\n'), ((72, 11, 72, 37), 'os.path.isfile', 'os.path.isfile', ({(72, 26, 72, 36): 'model_path'}, {}), '(model_path)', False, 'import os\n'), ((73, 21, 73, 81), 'os.path.join', 'os.path.join', ({(73, 34, 73, 54): '"""realesrgan/weights"""', (73, 56, 73, 80): "args.model_name + '.pth'"}, {}), "('realesrgan/weights', args.model_name + '.pth')", False, 'import os\n'), ((74, 11, 74, 37), 'os.path.isfile', 'os.path.isfile', ({(74, 26, 74, 36): 'model_path'}, {}), '(model_path)', False, 'import os\n'), ((89, 24, 94, 35), 'gfpgan.GFPGANer', 'GFPGANer', (), '', False, 'from gfpgan import GFPGANer\n'), ((103, 20, 103, 39), 'time.perf_counter', 'time.perf_counter', ({}, {}), '()', False, 'import time\n'), ((107, 14, 107, 52), 'cv2.imread', 'cv2.imread', ({(107, 25, 107, 29): 'path', (107, 31, 107, 51): 'cv2.IMREAD_UNCHANGED'}, {}), '(path, cv2.IMREAD_UNCHANGED)', False, 'import cv2\n'), ((119, 20, 119, 85), 'os.path.join', 'os.path.join', ({(119, 33, 119, 44): 'args.output', (119, 46, 119, 84): 'f"""{imgname}-{args.suffix}.{extension}"""'}, {}), "(args.output, f'{imgname}-{args.suffix}.{extension}')", False, 'import os\n'), ((120, 11, 120, 36), 'os.path.exists', 'os.path.exists', ({(120, 26, 120, 35): 'save_path'}, {}), '(save_path)', False, 'import os\n'), ((51, 16, 51, 101), 'basicsr.archs.rrdbnet_arch.RRDBNet', 'RRDBNet', (), '', False, 'from basicsr.archs.rrdbnet_arch import RRDBNet\n'), ((104, 46, 104, 68), 'os.path.basename', 'os.path.basename', ({(104, 63, 104, 67): 'path'}, {}), '(path)', False, 'import os\n'), ((132, 12, 132, 42), 'cv2.imwrite', 'cv2.imwrite', ({(132, 24, 132, 33): 'save_path', (132, 35, 132, 41): 'output'}, {}), '(save_path, output)', False, 'import cv2\n'), ((54, 16, 54, 102), 'basicsr.archs.rrdbnet_arch.RRDBNet', 'RRDBNet', (), '', False, 'from basicsr.archs.rrdbnet_arch import RRDBNet\n'), ((100, 33, 100, 62), 'os.path.join', 'os.path.join', ({(100, 46, 100, 56): 'args.input', (100, 58, 100, 61): '"""*"""'}, {}), "(args.input, '*')", False, 'import os\n'), ((59, 16, 59, 113), 'realesrgan.archs.srvgg_arch.SRVGGNetCompact', 'SRVGGNetCompact', (), '', False, 'from realesrgan.archs.srvgg_arch import SRVGGNetCompact\n'), ((64, 16, 64, 113), 'realesrgan.archs.srvgg_arch.SRVGGNetCompact', 'SRVGGNetCompact', (), '', False, 'from realesrgan.archs.srvgg_arch import SRVGGNetCompact\n'), ((67, 16, 67, 101), 'basicsr.archs.rrdbnet_arch.RRDBNet', 'RRDBNet', (), '', False, 'from basicsr.archs.rrdbnet_arch import RRDBNet\n'), ((133, 62, 133, 81), 'time.perf_counter', 'time.perf_counter', ({}, {}), '()', False, 'import time\n')]