lucasrufkahr/Real-ESRGAN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support outscale; Add RealESRGANx2 model; Version 0.2.1

Browse Source
This commit is contained in:
Xintao
2021-08-08 21:30:51 +08:00
parent 5745599813
commit 64ad194dda
4 changed files with 36 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
inference_realesrgan.py
View File

@@ -15,7 +15,8 @@ def main():
default='experiments/pretrained_models/RealESRGAN_x4plus.pth',
help='Path to the pre-trained model')
parser.add_argument('--output', type=str, default='results', help='Output folder')
parser.add_argument('--scale', type=int, default=4, help='Upsample scale factor')
parser.add_argument('--netscale', type=int, default=4, help='Upsample scale factor of the network')
parser.add_argument('--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('--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')
@@ -34,7 +35,7 @@ def main():
args = parser.parse_args()
upsampler = RealESRGANer(
scale=args.scale,
scale=args.netscale,
model_path=args.model_path,
tile=args.tile,
tile_pad=args.tile_pad,
@@ -51,15 +52,25 @@ def main():
print('Testing', idx, imgname)
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
output, img_mode = upsampler.enhance(img)
if args.ext == 'auto':
extension = extension[1:]
h, w = img.shape[0:2]
if max(h, w) > 1000 and args.netscale == 4:
print('WARNING: The input image is large, try X2 model for better performace.')
if max(h, w) < 500 and args.netscale == 2:
print('WARNING: The input image is small, try X4 model for better performace.')
try:
output, img_mode = upsampler.enhance(img, outscale=args.outscale)
except Exception as error:
print('Error', error)
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}')
cv2.imwrite(save_path, output)
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'
save_path = os.path.join(args.output, f'{imgname}_{args.suffix}.{extension}')
cv2.imwrite(save_path, output)
if __name__ == '__main__':