add grayscale judgement (#32)

facelib/detection/yolov5face/face_detector.py

@@ -17,7 +17,7 @@ from facelib.detection.yolov5face.utils.general import (
     scale_coords_landmarks,
 )
 
-IS_HIGH_VERSION = tuple(map(int, torch.__version__.split('+')[0].split('.'))) >= (1, 9, 0)
+IS_HIGH_VERSION = tuple(map(int, torch.__version__.split('+')[0].split('.')[:3])) >= (1, 9, 0)
 
 
 def isListempty(inList):

facelib/utils/face_restoration_helper.py

@@ -299,6 +299,8 @@ class FaceRestoreHelper(object):
 
 
     def add_restored_face(self, face):
+        if self.is_gray:
+            face = bgr2gray(face) # convert img into grayscale
         self.restored_faces.append(face)
 
 
@@ -419,9 +421,6 @@ class FaceRestoreHelper(object):
                 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 self.is_gray:
-                pasted_face = bgr2gray(pasted_face) # convert img into grayscale
-
             if len(upsample_img.shape) == 3 and upsample_img.shape[2] == 4:  # alpha channel
                 alpha = upsample_img[:, :, 3:]
                 upsample_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * upsample_img[:, :, 0:3]

inference_codeformer.py

@@ -1,4 +1,3 @@
-# Modified by Shangchen Zhou from: https://github.com/TencentARC/GFPGAN/blob/master/inference_gfpgan.py
 import os
 import cv2
 import argparse
@@ -8,6 +7,7 @@ from torchvision.transforms.functional import normalize
 from basicsr.utils import imwrite, img2tensor, tensor2img
 from basicsr.utils.download_util import load_file_from_url
 from facelib.utils.face_restoration_helper import FaceRestoreHelper
+from facelib.utils.misc import is_gray
 import torch.nn.functional as F
 
 from basicsr.utils.registry import ARCH_REGISTRY
@@ -122,6 +122,9 @@ if __name__ == '__main__':
         if args.has_aligned: 
             # the input faces are already cropped and aligned
             img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
+            face_helper.is_gray = is_gray(img, threshold=5)
+            if face_helper.is_gray:
+                print('Grayscale input: True')
             face_helper.cropped_faces = [img]
         else:
             face_helper.read_image(img)