93 lines
4.1 KiB
Python
93 lines
4.1 KiB
Python
import argparse
|
|
import os
|
|
import struct
|
|
from basicsr.archs.rrdbnet_arch import RRDBNet
|
|
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='../TestData3', help='Input image or folder')
|
|
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 | '
|
|
'realesr-animevideov3'))
|
|
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='out', 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(
|
|
'--fp32', action='store_true', help='Use fp32 precision during inference. Default: fp16 (half precision).')
|
|
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 ['realesr-animevideov3']: # 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
|
|
|
|
# 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.fp32)
|
|
|
|
if os.path.isfile('real-esrgan.wts'):
|
|
print('Already, real-esrgan.wts file exists.')
|
|
else:
|
|
print('making real-esrgan.wts file ...')
|
|
f = open("real-esrgan.wts", 'w')
|
|
f.write("{}\n".format(len(upsampler.model.state_dict().keys())))
|
|
for k, v in upsampler.model.state_dict().items():
|
|
print('key: ', k)
|
|
print('value: ', v.shape)
|
|
vr = v.reshape(-1).cpu().numpy()
|
|
f.write("{} {}".format(k, len(vr)))
|
|
for vv in vr:
|
|
f.write(" ")
|
|
f.write(struct.pack(">f", float(vv)).hex())
|
|
f.write("\n")
|
|
print('Completed real-esrgan.wts file!')
|
|
|
|
if __name__ == '__main__':
|
|
main()
|