support face enhance

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