From 667e34e7cab86a0ff7f74ae49676a471b701054f Mon Sep 17 00:00:00 2001 From: Xintao Date: Sun, 22 Aug 2021 18:09:28 +0800 Subject: [PATCH] support face enhance --- README.md | 3 ++- VERSION | 2 +- inference_realesrgan.py | 24 ++++++++++++++++++++++-- realesrgan/utils.py | 2 +- requirements.txt | 4 +++- 5 files changed, 29 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 58c93d8..f8a231b 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ We extend the powerful ESRGAN to a practical restoration application (namely, Re :question: Frequently Asked Questions can be found in [FAQ.md](FAQ.md). :triangular_flag_on_post: **Updates** +- :white_check_mark: Integrate [GFPGAN](https://github.com/TencentARC/GFPGAN) to support face enhancement. - :white_check_mark: Integrated to [Huggingface Spaces](https://huggingface.co/spaces) with [Gradio](https://github.com/gradio-app/gradio). See [Gradio Web Demo](https://huggingface.co/spaces/akhaliq/Real-ESRGAN). - :white_check_mark: Support arbitrary scale with `--outscale` (It actually further resizes outputs with `LANCZOS4`). Add *RealESRGAN_x2plus.pth* model. - :white_check_mark: [The inference code](inference_realesrgan.py) supports: 1) **tile** options; 2) images with **alpha channel**; 3) **gray** images; 4) **16-bit** images. @@ -123,7 +124,7 @@ wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_ Inference! ```bash -python inference_realesrgan.py --model_path experiments/pretrained_models/RealESRGAN_x4plus.pth --input inputs +python inference_realesrgan.py --model_path experiments/pretrained_models/RealESRGAN_x4plus.pth --input inputs --face_enhance ``` Results are in the `results` folder diff --git a/VERSION b/VERSION index 0c62199..68cc747 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.2.1 +0.2.2.0 diff --git a/inference_realesrgan.py b/inference_realesrgan.py index 207eb83..784c817 100644 --- a/inference_realesrgan.py +++ b/inference_realesrgan.py @@ -18,9 +18,10 @@ def main(): 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', type=int, default=800, 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('--half', action='store_true', help='Use half precision during inference') parser.add_argument( '--alpha_upsampler', @@ -41,7 +42,17 @@ def main(): tile_pad=args.tile_pad, pre_pad=args.pre_pad, half=args.half) + + if args.face_enhance: + from gfpgan import GFPGANer + face_enhancer = GFPGANer( + model_path='https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth', + upscale=args.outscale, + arch='clean', + channel_multiplier=2, + bg_upsampler=upsampler) os.makedirs(args.output, exist_ok=True) + if os.path.isfile(args.input): paths = [args.input] else: @@ -52,6 +63,11 @@ def main(): 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 + h, w = img.shape[0:2] if max(h, w) > 1000 and args.netscale == 4: import warnings @@ -61,9 +77,13 @@ def main(): warnings.warn('The input image is small, try X4 model for better performace.') try: - output, img_mode = upsampler.enhance(img, outscale=args.outscale) + 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 Exception 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:] diff --git a/realesrgan/utils.py b/realesrgan/utils.py index 2ec3de1..15f1957 100644 --- a/realesrgan/utils.py +++ b/realesrgan/utils.py @@ -143,7 +143,7 @@ class RealESRGANer(): h_input, w_input = img.shape[0:2] # img: numpy img = img.astype(np.float32) - if np.max(img) > 255: # 16-bit image + if np.max(img) > 256: # 16-bit image max_range = 65535 print('\tInput is a 16-bit image') else: diff --git a/requirements.txt b/requirements.txt index d265bab..788adf0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,6 @@ -basicsr>=1.3.3.10 +basicsr>=1.3.3.11 +facexlib>=0.2.0.3 +gfpgan>=0.2.1 numpy opencv-python torch>=1.7