Spaces:
Running
Running
File size: 2,695 Bytes
8875fed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
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()
|