some modifications on detection and fusion.

README.md

@@ -20,6 +20,7 @@ S-Lab, Nanyang Technological University
 
 ### Updates
 
+- **2022.08.23**: Some modifications on face detection and fusion for better AI-created face enhancement.
 - **2022.08.07**: Integrate Real-ESRGAN to support background image enhancement.
 - **2022.07.29**: New face detector with supporting `['YOLOv5', 'RetinaFace']`. 
 - **2022.07.17**: The Colab demo of CodeFormer is available now. <a href="https://colab.research.google.com/drive/1m52PNveE4PBhYrecj34cnpEeiHcC5LTb?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a>
@@ -110,7 +111,7 @@ If our work is useful for your research, please consider citing:
 
 ### Acknowledgement
 
-This project is based on [BasicSR](https://github.com/XPixelGroup/BasicSR). We also borrow some codes from [Unleashing Transformers](https://github.com/samb-t/unleashing-transformers), [YOLOv5-face](https://github.com/deepcam-cn/yolov5-face), and [FaceXLib](https://github.com/xinntao/facexlib).
+This project is based on [BasicSR](https://github.com/XPixelGroup/BasicSR). We also borrow some codes from [Unleashing Transformers](https://github.com/samb-t/unleashing-transformers), [YOLOv5-face](https://github.com/deepcam-cn/yolov5-face), and [FaceXLib](https://github.com/xinntao/facexlib). Thanks for their awesome works.
 
 ### Contact
 If you have any question, please feel free to reach me out at `shangchenzhou@gmail.com`.

facelib/utils/face_restoration_helper.py

@@ -321,6 +321,60 @@ class FaceRestoreHelper(object):
             inverse_affine[:, 2] += extra_offset
             inv_restored = cv2.warpAffine(restored_face, inverse_affine, (w_up, h_up))
 
+            # if draw_box or not self.use_parse:  # use square parse maps
+            #     mask = np.ones(self.face_size, dtype=np.float32)
+            #     inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up))
+            #     # remove the black borders
+            #     inv_mask_erosion = cv2.erode(
+            #         inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8))
+            #     pasted_face = inv_mask_erosion[:, :, None] * inv_restored
+            #     total_face_area = np.sum(inv_mask_erosion)  # // 3
+            #     # add border
+            #     if draw_box:
+            #         h, w = self.face_size
+            #         mask_border = np.ones((h, w, 3), dtype=np.float32)
+            #         border = int(1400/np.sqrt(total_face_area))
+            #         mask_border[border:h-border, border:w-border,:] = 0
+            #         inv_mask_border = cv2.warpAffine(mask_border, inverse_affine, (w_up, h_up))
+            #         inv_mask_borders.append(inv_mask_border)
+            #     if not self.use_parse:
+            #         # compute the fusion edge based on the area of face
+            #         w_edge = int(total_face_area**0.5) // 20
+            #         erosion_radius = w_edge * 2
+            #         inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
+            #         blur_size = w_edge * 2
+            #         inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
+            #         if len(upsample_img.shape) == 2:  # upsample_img is gray image
+            #             upsample_img = upsample_img[:, :, None]
+            #         inv_soft_mask = inv_soft_mask[:, :, None]
+
+            # always use square mask
+            mask = np.ones(self.face_size, dtype=np.float32)
+            inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up))
+            # remove the black borders
+            inv_mask_erosion = cv2.erode(
+                inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8))
+            pasted_face = inv_mask_erosion[:, :, None] * inv_restored
+            total_face_area = np.sum(inv_mask_erosion)  # // 3
+            # add border
+            if draw_box:
+                h, w = self.face_size
+                mask_border = np.ones((h, w, 3), dtype=np.float32)
+                border = int(1400/np.sqrt(total_face_area))
+                mask_border[border:h-border, border:w-border,:] = 0
+                inv_mask_border = cv2.warpAffine(mask_border, inverse_affine, (w_up, h_up))
+                inv_mask_borders.append(inv_mask_border)
+            # compute the fusion edge based on the area of face
+            w_edge = int(total_face_area**0.5) // 20
+            erosion_radius = w_edge * 2
+            inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
+            blur_size = w_edge * 2
+            inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
+            if len(upsample_img.shape) == 2:  # upsample_img is gray image
+                upsample_img = upsample_img[:, :, None]
+            inv_soft_mask = inv_soft_mask[:, :, None]
+
+            # parse mask
             if self.use_parse:
                 # inference
                 face_input = cv2.resize(restored_face, (512, 512), interpolation=cv2.INTER_LINEAR)
@@ -331,52 +385,27 @@ class FaceRestoreHelper(object):
                     out = self.face_parse(face_input)[0]
                 out = out.argmax(dim=1).squeeze().cpu().numpy()
 
-                mask = np.zeros(out.shape)
+                parse_mask = np.zeros(out.shape)
                 MASK_COLORMAP = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 255, 0, 0, 0]
                 for idx, color in enumerate(MASK_COLORMAP):
-                    mask[out == idx] = color
+                    parse_mask[out == idx] = color
                 #  blur the mask
-                mask = cv2.GaussianBlur(mask, (101, 101), 11)
-                mask = cv2.GaussianBlur(mask, (101, 101), 11)
+                parse_mask = cv2.GaussianBlur(parse_mask, (101, 101), 11)
+                parse_mask = cv2.GaussianBlur(parse_mask, (101, 101), 11)
                 # remove the black borders
                 thres = 10
-                mask[:thres, :] = 0
-                mask[-thres:, :] = 0
-                mask[:, :thres] = 0
-                mask[:, -thres:] = 0
-                mask = mask / 255.
-
-                mask = cv2.resize(mask, restored_face.shape[:2])
-                mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up), flags=3)
-                inv_soft_mask = mask[:, :, None]
-                pasted_face = inv_restored
-
-            if draw_box or not self.use_parse:  # use square parse maps
-                mask = np.ones(self.face_size, dtype=np.float32)
-                inv_mask = cv2.warpAffine(mask, inverse_affine, (w_up, h_up))
-                # remove the black borders
-                inv_mask_erosion = cv2.erode(
-                    inv_mask, np.ones((int(2 * self.upscale_factor), int(2 * self.upscale_factor)), np.uint8))
-                pasted_face = inv_mask_erosion[:, :, None] * inv_restored
-                total_face_area = np.sum(inv_mask_erosion)  # // 3
-                # add border
-                if draw_box:
-                    h, w = self.face_size
-                    mask_border = np.ones((h, w, 3), dtype=np.float32)
-                    border = int(1400/np.sqrt(total_face_area))
-                    mask_border[border:h-border, border:w-border,:] = 0
-                    inv_mask_border = cv2.warpAffine(mask_border, inverse_affine, (w_up, h_up))
-                    inv_mask_borders.append(inv_mask_border)
-                if not self.use_parse:
-                    # compute the fusion edge based on the area of face
-                    w_edge = int(total_face_area**0.5) // 20
-                    erosion_radius = w_edge * 2
-                    inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
-                    blur_size = w_edge * 2
-                    inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
-                    if len(upsample_img.shape) == 2:  # upsample_img is gray image
-                        upsample_img = upsample_img[:, :, None]
-                    inv_soft_mask = inv_soft_mask[:, :, None]
+                parse_mask[:thres, :] = 0
+                parse_mask[-thres:, :] = 0
+                parse_mask[:, :thres] = 0
+                parse_mask[:, -thres:] = 0
+                parse_mask = parse_mask / 255.
+
+                parse_mask = cv2.resize(parse_mask, restored_face.shape[:2])
+                parse_mask = cv2.warpAffine(parse_mask, inverse_affine, (w_up, h_up), flags=3)
+                inv_soft_parse_mask = parse_mask[:, :, None]
+                # pasted_face = inv_restored
+                fuse_mask = (inv_soft_parse_mask<inv_soft_mask).astype('int')
+                inv_soft_mask = inv_soft_parse_mask*fuse_mask + inv_soft_mask*(1-fuse_mask)
 
             if len(upsample_img.shape) == 3 and upsample_img.shape[2] == 4:  # alpha channel
                 alpha = upsample_img[:, :, 3:]
@@ -390,7 +419,7 @@ class FaceRestoreHelper(object):
         else:
             upsample_img = upsample_img.astype(np.uint8)
 
-        # add border
+        # draw bounding box
         if draw_box:
             # upsample_input_img = cv2.resize(input_img, (w_up, h_up))
             img_color = np.ones([*upsample_img.shape], dtype=np.float32)
@@ -414,4 +443,4 @@ class FaceRestoreHelper(object):
         self.cropped_faces = []
         self.inverse_affine_matrices = []
         self.det_faces = []
-        self.pad_input_imgs = []
+        self.pad_input_imgs = []

inference_codeformer.py

@@ -25,6 +25,9 @@ if __name__ == '__main__':
     parser.add_argument('--test_path', type=str, default='./inputs/cropped_faces')
     parser.add_argument('--has_aligned', action='store_true', help='Input are cropped and aligned faces')
     parser.add_argument('--only_center_face', action='store_true', help='Only restore the center face')
+    # large det_model: 'YOLOv5l', 'retinaface_resnet50'
+    # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'
+    parser.add_argument('--detection_model', type=str, default='retinaface_resnet50')
     parser.add_argument('--draw_box', action='store_true')
     parser.add_argument('--bg_upsampler', type=str, default='None', help='background upsampler. Optional: realesrgan')
     parser.add_argument('--bg_tile', type=int, default=400, help='Tile size for background sampler. Default: 400')
@@ -55,7 +58,7 @@ if __name__ == '__main__':
                 model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
                 model=model,
                 tile=args.bg_tile,
-                tile_pad=10,
+                tile_pad=40,
                 pre_pad=0,
                 half=True)  # need to set False in CPU mode
     else:
@@ -75,11 +78,13 @@ if __name__ == '__main__':
     # ------------------ set up FaceRestoreHelper -------------------
     # large det_model: 'YOLOv5l', 'retinaface_resnet50'
     # small det_model: 'YOLOv5n', 'retinaface_mobile0.25'
+    if not args.has_aligned: 
+        print(f'Using [{args.detection_model}] for face detection network.')
     face_helper = FaceRestoreHelper(
         args.upscale,
         face_size=512,
         crop_ratio=(1, 1),
-        det_model = 'YOLOv5l',
+        det_model = args.detection_model,
         save_ext='png',
         use_parse=True,
         device=device)
@@ -89,7 +94,7 @@ if __name__ == '__main__':
     for img_path in sorted(glob.glob(os.path.join(args.test_path, '*.[jp][pn]g'))):
         # clean all the intermediate results to process the next image
         face_helper.clean_all()
-
+        
         img_name = os.path.basename(img_path)
         print(f'Processing: {img_name}')
         basename, ext = os.path.splitext(img_name)

scripts/download_pretrained_models.py

@@ -27,7 +27,8 @@ if __name__ == '__main__':
             'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
         },
         'facelib': {
-            'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',
+            # 'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',
+            'detection_Resnet50_Final.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',
             'parsing_parsenet.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
         }
     }