Ensures adain_npy returned pixel value is within the range of bit depth

facelib/utils/face_restoration_helper.py

@@ -6,7 +6,7 @@
 
 from facelib.detection import init_detection_model
 from facelib.parsing import init_parsing_model
-from facelib.utils.misc import img2tensor, imwrite, is_gray, bgr2gray, adain_npy
+from facelib.utils.misc import img2tensor, imwrite, color_diff, bgr2gray, adain_npy
 from basicsr.utils.download_util import load_file_from_url
 from basicsr.utils.misc import get_device
 
@@ -141,14 +141,24 @@ def read_image(self, img):
             img = img[:, :, 0:3]
 
         self.input_img = img
-        self.is_gray = is_gray(img, threshold=10)
+        self.color_diff = color_diff(img)
+        self.is_gray = (self.color_diff <= 10)
         if self.is_gray:
             print('Grayscale input: True')
 
         if min(self.input_img.shape[:2])<512:
             f = 512.0/min(self.input_img.shape[:2])
             self.input_img = cv2.resize(self.input_img, (0,0), fx=f, fy=f, interpolation=cv2.INTER_LINEAR)
 
+    def read_aligned(self, img, gray_threshold = 10):
+        # the input faces are already cropped and aligned
+        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
+        self.color_diff = color_diff(img)
+        self.is_gray = (self.color_diff <= gray_threshold)
+        if self.is_gray:
+            print('Grayscale input: True')
+        self.cropped_faces = [img]
+
     def init_dlib(self, detection_path, landmark5_path):
         """Initialize the dlib detectors and predictors."""
         try:
@@ -364,7 +374,7 @@ def get_inverse_affine(self, save_inverse_affine_path=None):
     def add_restored_face(self, restored_face, input_face=None):
         if self.is_gray:
             restored_face = bgr2gray(restored_face) # convert img into grayscale
-            if input_face is not None:
+            if self.color_diff >= 1 and input_face is not None:
                 restored_face = adain_npy(restored_face, input_face) # transfer the color
         self.restored_faces.append(restored_face)
 
@@ -522,4 +532,4 @@ def clean_all(self):
         self.cropped_faces = []
         self.inverse_affine_matrices = []
         self.det_faces = []
-        self.pad_input_imgs = []
+        self.pad_input_imgs = []

facelib/utils/misc.py

@@ -143,7 +143,7 @@ def _scandir(dir_path, suffix, recursive):
     return _scandir(dir_path, suffix=suffix, recursive=recursive)
 
 
-def is_gray(img, threshold=10):
+def color_diff(img):
     img = Image.fromarray(img)
     if len(img.getbands()) == 1:
         return True
@@ -153,8 +153,10 @@ def is_gray(img, threshold=10):
     diff1 = (img1 - img2).var()
     diff2 = (img2 - img3).var()
     diff3 = (img3 - img1).var()
-    diff_sum = (diff1 + diff2 + diff3) / 3.0
-    if diff_sum <= threshold:
+    return (diff1 + diff2 + diff3) / 3.0
+
+def is_gray(img, threshold=10):
+    if color_diff(img) <= threshold:
         return True
     else:
         return False
@@ -199,4 +201,18 @@ def adain_npy(content_feat, style_feat):
     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)
+    result_feat = normalized_feat * np.broadcast_to(style_std, size) + np.broadcast_to(style_mean, size)
+
+    # Ensure values are within the range of source image bit depth
+    bit_range = 256 if np.max(content_feat) < 256 else 65536 # determine 8 bit or 16 bit.
+    a_min, a_max = np.min(result_feat, axis=(0,1)), np.max(result_feat, axis=(0,1)) # min/max of each color
+    i_min, i_max = np.argmin(a_min), np.argmax(a_max) # find the color index of global min and max
+    v_min, v_max = a_min[i_min], a_max[i_max]         # global min and max
+    if v_max > bit_range or v_min < 0:                # pixel value is out of the range of bit depth
+        # reduce the style_std to clamp values in range.
+        mean_min, mean_max = style_mean[0][0][i_min], style_mean[0][0][i_max] # mean of color for min/max
+        ratio = min(mean_min / (mean_min - v_min), (bit_range - 1e-12 - mean_max) / (v_max - mean_max))
+        style_std = style_std * np.broadcast_to([ratio, ratio, ratio], style_std.shape)
+        result_feat = normalized_feat * np.broadcast_to(style_std, size) + np.broadcast_to(style_mean, size)
+
+    return result_feat

inference_codeformer.py

@@ -8,7 +8,6 @@
 from basicsr.utils.download_util import load_file_from_url
 from basicsr.utils.misc import gpu_is_available, get_device
 from facelib.utils.face_restoration_helper import FaceRestoreHelper
-from facelib.utils.misc import is_gray
 
 from basicsr.utils.registry import ARCH_REGISTRY
 
@@ -178,12 +177,7 @@ def set_realesrgan():
             img = img_path
 
         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=10)
-            if face_helper.is_gray:
-                print('Grayscale input: True')
-            face_helper.cropped_faces = [img]
+            face_helper.read_aligned(img)
         else:
             face_helper.read_image(img)
             # get face landmarks for each face

web-demos/hugging_face/app.py

@@ -21,7 +21,6 @@
 from basicsr.utils.registry import ARCH_REGISTRY
 
 from facelib.utils.face_restoration_helper import FaceRestoreHelper
-from facelib.utils.misc import is_gray
 
 
 os.system("pip freeze")
@@ -140,12 +139,7 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide
         face_upsampler = upsampler if face_upsample else None
 
         if 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('\tgrayscale input: True')
-            face_helper.cropped_faces = [img]
+            face_helper.read_aligned(img, gray_threshold=5)
         else:
             face_helper.read_image(img)
             # get face landmarks for each face
@@ -280,4 +274,4 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide
     )
 
 demo.queue(concurrency_count=2)
-demo.launch()
+demo.launch()