fix queue jam.

CodeFormer/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):

CodeFormer/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
+from facelib.utils.misc import img2tensor, imwrite, is_gray, bgr2gray
 
 
 def get_largest_face(det_faces, h, w):
@@ -125,6 +125,9 @@ class FaceRestoreHelper(object):
             img = img[:, :, 0:3]
 
         self.input_img = img
+        self.is_gray = is_gray(img, threshold=5)
+        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])
@@ -296,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)
 
 

CodeFormer/facelib/utils/misc.py

@@ -1,6 +1,8 @@
 import cv2
 import os
 import os.path as osp
+import numpy as np
+from PIL import Image
 import torch
 from torch.hub import download_url_to_file, get_dir
 from urllib.parse import urlparse
@@ -139,3 +141,34 @@ def scandir(dir_path, suffix=None, recursive=False, full_path=False):
                     continue
 
     return _scandir(dir_path, suffix=suffix, recursive=recursive)
+
+
+def is_gray(img, threshold=10):
+    img = Image.fromarray(img)
+    if len(img.getbands()) == 1:
+        return True
+    img1 = np.asarray(img.getchannel(channel=0), dtype=np.int16)
+    img2 = np.asarray(img.getchannel(channel=1), dtype=np.int16)
+    img3 = np.asarray(img.getchannel(channel=2), dtype=np.int16)
+    diff1 = (img1 - img2).var()
+    diff2 = (img2 - img3).var()
+    diff3 = (img3 - img1).var()
+    diff_sum = (diff1 + diff2 + diff3) / 3.0
+    if diff_sum <= threshold:
+        return True
+    else:
+        return False
+
+def rgb2gray(img, out_channel=3):
+    r, g, b = img[:,:,0], img[:,:,1], img[:,:,2]
+    gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
+    if out_channel == 3:
+        gray = gray[:,:,np.newaxis].repeat(3, axis=2)
+    return gray
+
+def bgr2gray(img, out_channel=3):
+    b, g, r = img[:,:,0], img[:,:,1], img[:,:,2]
+    gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
+    if out_channel == 3:
+        gray = gray[:,:,np.newaxis].repeat(3, axis=2)
+    return gray

app.py

@@ -16,6 +16,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
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils.realesrgan_utils import RealESRGANer
 
@@ -108,6 +109,7 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide
     draw_box = False
     detection_model = "retinaface_resnet50"
 
+    upscale = int(upscale) # covert type to int
     face_helper = FaceRestoreHelper(
         upscale,
         face_size=512,
@@ -125,6 +127,9 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide
     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('Grayscale input: True')
         face_helper.cropped_faces = [img]
     else:
         face_helper.read_image(img)
@@ -228,7 +233,7 @@ If you have any questions, please feel free to reach me out at <b>shangchenzhou@
 ![visitors](https://visitor-badge.glitch.me/badge?page_id=sczhou/CodeFormer)
 """
 
-gr.Interface(
+demo = gr.Interface(
     inference, [
         gr.inputs.Image(type="filepath", label="Input"),
         gr.inputs.Checkbox(default=True, label="Background_Enhance"),
@@ -249,3 +254,6 @@ gr.Interface(
         ['05.jpg', True, True, 2, 0.1]
       ]
     ).launch()
+
+demo.queue(concurrency_count=4)
+demo.launch()