transfer color for grayscale inputs.

facelib/utils/face_restoration_helper.py

@@ -6,7 +6,7 @@ from torchvision.transforms.functional import normalize
 
 from facelib.detection import init_detection_model
 from facelib.parsing import init_parsing_model
-from facelib.utils.misc import img2tensor, imwrite, is_gray, bgr2gray
+from facelib.utils.misc import img2tensor, imwrite, is_gray, bgr2gray, adain_npy
 from basicsr.utils.misc import get_device
 
 
@@ -300,10 +300,12 @@ class FaceRestoreHelper(object):
                 torch.save(inverse_affine, save_path)
 
 
-    def add_restored_face(self, face):
+    def add_restored_face(self, restored_face, input_face=None):
         if self.is_gray:
-            face = bgr2gray(face) # convert img into grayscale
-        self.restored_faces.append(face)
+            restored_face = bgr2gray(restored_face) # convert img into grayscale
+            if input_face is not None:
+                restored_face = adain_npy(restored_face, input_face) # transfer the color
+        self.restored_faces.append(restored_face)
 
 
     def paste_faces_to_input_image(self, save_path=None, upsample_img=None, draw_box=False, face_upsampler=None):

facelib/utils/misc.py

@@ -172,3 +172,31 @@ def bgr2gray(img, out_channel=3):
     if out_channel == 3:
         gray = gray[:,:,np.newaxis].repeat(3, axis=2)
     return gray
+
+
+def calc_mean_std(feat, eps=1e-5):
+    """
+    Args:
+        feat (numpy): 3D [w h c]s
+    """
+    size = feat.shape
+    assert len(size) == 3, 'The input feature should be 3D tensor.'
+    c = size[2]
+    feat_var = feat.reshape(-1, c).var(axis=0) + eps
+    feat_std = np.sqrt(feat_var).reshape(1, 1, c)
+    feat_mean = feat.reshape(-1, c).mean(axis=0).reshape(1, 1, c)
+    return feat_mean, feat_std
+
+
+def adain_npy(content_feat, style_feat):
+    """Adaptive instance normalization for numpy.
+
+    Args:
+        content_feat (numpy): The input feature.
+        style_feat (numpy): The reference feature.
+    """
+    size = content_feat.shape
+    style_mean, style_std = calc_mean_std(style_feat)
+    content_mean, content_std = calc_mean_std(content_feat)
+    normalized_feat = (content_feat - np.broadcast_to(content_mean, size)) / np.broadcast_to(content_std, size)
+    return normalized_feat * np.broadcast_to(style_std, size) + np.broadcast_to(style_mean, size)

inference_codeformer.py

@@ -205,7 +205,7 @@ if __name__ == '__main__':
                 restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
 
             restored_face = restored_face.astype('uint8')
-            face_helper.add_restored_face(restored_face)
+            face_helper.add_restored_face(restored_face, cropped_face)
 
         # paste_back
         if not args.has_aligned: