From e5adc0dd06aedbdce712eabd9a42755e28b0665f Mon Sep 17 00:00:00 2001 From: Xintao Date: Fri, 10 Sep 2021 18:46:45 +0800 Subject: [PATCH] add ext option --- inference_gfpgan.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/inference_gfpgan.py b/inference_gfpgan.py index a62708f..148923d 100644 --- a/inference_gfpgan.py +++ b/inference_gfpgan.py @@ -24,6 +24,12 @@ def main(): parser.add_argument('--aligned', action='store_true') parser.add_argument('--paste_back', action='store_false') parser.add_argument('--save_root', type=str, default='results') + 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() args = parser.parse_args() if args.test_path.endswith('/'): @@ -85,10 +91,16 @@ def main(): # save restored img if restored_img is not None: - if args.suffix is not None: - save_restore_path = os.path.join(args.save_root, 'restored_imgs', f'{basename}_{args.suffix}{ext}') + if args.ext == 'auto': + extension = ext[1:] else: - save_restore_path = os.path.join(args.save_root, 'restored_imgs', img_name) + extension = args.ext + + if args.suffix is not None: + save_restore_path = os.path.join(args.save_root, 'restored_imgs', + f'{basename}_{args.suffix}.{extension}') + else: + save_restore_path = os.path.join(args.save_root, 'restored_imgs', f'{basename}.{extension}') imwrite(restored_img, save_restore_path) print(f'Results are in the [{args.save_root}] folder.')