inference_realesrgan.py

import argparse
import cv2
import glob
import os
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url

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


def main(**args):
    """Inference demo for Real-ESRGAN.
    """
    print(args)

    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
    netscale = 4
    file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']


    # determine model paths
    model_path = os.path.join('weights', 'RealESRGAN_x4plus.pth')

    # use dni to control the denoise strength
    dni_weight = None

    # restorer
    upsampler = RealESRGANer(
        scale=netscale,
        model_path=model_path,
        dni_weight=dni_weight,
        model=model,
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=not args['fp32'],
        gpu_id=None)

    if args['face_enhance'] == 'True':  # Use GFPGAN for face enhancement
        from gfpgan import GFPGANer
        face_enhancer = GFPGANer(
            model_path=os.path.join('weights', 'GFPGANv1.3.pth'),
            upscale=args['outscale'],
            arch='clean',
            channel_multiplier=2,
            bg_upsampler=upsampler)

    # os.makedirs(args.output, exist_ok=True)

    path = args['input']

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

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

    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:
        if args['ext'] == 'auto':
            extension = extension[1:]
        else:
            extension = args['ext']

        if img_mode == 'RGBA':  # RGBA images should be saved in png format
            extension = 'png'
        # if args.suffix == '':
        save_path = os.path.join(args['output'], f'{imgname}.{extension}')
        # else:
        #     save_path = os.path.join(args.output, f'{imgname}_{args.suffix}.{extension}')
        cv2.imwrite(save_path, output)
        return save_path
            # return output

# if __name__ == '__main__':
#     main()