support face enhance

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

VERSION

@@ -1 +1 @@
-0.2.1
+0.2.2.0

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:]

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:

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