moving to kornia

app.py

@@ -3,27 +3,63 @@ import cv2
 import gradio as gr
 from PIL import Image
 import numpy as np
+import torch
+import kornia as K
+from kornia.contrib import FaceDetector, FaceDetectorResult
+
+device = torch.device('cpu')
+face_detection = FaceDetector().to(device)
+
+def scale_image(img: np.ndarray, size: int) -> np.ndarray:
+    h, w = img.shape[:2]
+    scale = 1. * size / w
+    return cv2.resize(img, (int(w * scale), int(h * scale)))
+
+
+def apply_blur_face(img: torch.Tensor, img_vis: np.ndarray, det: FaceDetectorResult):
+    # crop the face
+    x1, y1 = det.xmin.int(), det.ymin.int()
+    x2, y2 = det.xmax.int(), det.ymax.int()
+    roi = img[..., y1:y2, x1:x2]
+    if roi.shape[-1]==0 or roi.shape[-2]==0:
+        return
+
+    # apply blurring and put back to the visualisation image
+    roi = K.filters.gaussian_blur2d(roi, (21, 21), (100., 100.))
+    roi = K.color.rgb_to_bgr(roi)
+    img_vis[y1:y2, x1:x2] = K.tensor_to_image(roi)
+
 
 def run(image):
     image.thumbnail((1280, 1280))
-    image = np.array(image)
-    face_locations = face_recognition.face_locations(image, model="cnn")
+    img_raw = np.array(image)
+
+    # preprocess
+    img = K.image_to_tensor(img_raw, keepdim=False).to(device)
+    img = K.color.bgr_to_rgb(img.float())
+
+    with torch.no_grad():
+        dets = face_detection(img)
+    dets = [FaceDetectorResult(o) for o in dets]
+
+    img_vis = img_raw.copy()
+
+    for b in dets:
+        if b.score < 0.5:
+            continue
+
+        apply_blur_face(img, img_vis, b)
 
-    for top, right, bottom, left in face_locations:
-        face_image = image[top:bottom, left:right]
-        face_image = cv2.GaussianBlur(face_image, (99, 99), 30)
-        image[top:bottom, left:right] = face_image
-    
-    return Image.fromarray(image)
+    return Image.fromarray(img_vis)
 
 content_image_input = gr.inputs.Image(label="Content Image", type="pil")
 
 description="Privacy first! Upload an image of a groupf of people and blur their faces automatically."
 article="""
-Demo built with the face_recognition package and opencv, based on 
-<a href='https://github.com/ageitgey/face_recognition/blob/master/examples/blur_faces_on_webcam.py' target='_blank'>this</a> example.
+Demo built on top of kornia and opencv, based on 
+<a href='https://github.com/kornia/kornia/blob/master/examples/face_detection/main.py' target='_blank'>this</a> example.
 """
-examples = [["family.jpeg"], ["crowd.jpeg"], ["crowd1.jpeg"]]
+examples = [["./images/family.jpeg"], ["./images/crowd.jpeg"], ["./images/crowd1.jpeg"]]
 
 app_interface = gr.Interface(fn=run,
                              inputs=[content_image_input],

face_rec_benchmark.py

@@ -0,0 +1,27 @@
+import face_recognition
+import cv2
+import gradio as gr
+from PIL import Image
+import numpy as np
+import time
+
+def run(image):
+    image.thumbnail((1280, 1280))
+    image = np.array(image)
+    face_locations = face_recognition.face_locations(image, model="cnn")
+
+    for top, right, bottom, left in face_locations:
+        face_image = image[top:bottom, left:right]
+        face_image = cv2.GaussianBlur(face_image, (99, 99), 30)
+        image[top:bottom, left:right] = face_image
+    
+    return Image.fromarray(image)
+
+if __name__ == "__main__":
+
+    start = time.time()
+    for _ in range(100):
+        image = Image.open("./images/crowd.jpeg")
+        _ = run(image)
+
+    print('It took', (time.time()-start)/100, 'seconds.')

kornia_benchmark.py

@@ -0,0 +1,63 @@
+import cv2
+import gradio as gr
+from PIL import Image
+import numpy as np
+import torch
+import kornia as K
+from kornia.contrib import FaceDetector, FaceDetectorResult
+import time
+
+device = torch.device('cpu')
+face_detection = FaceDetector().to(device)
+
+def scale_image(img: np.ndarray, size: int) -> np.ndarray:
+    h, w = img.shape[:2]
+    scale = 1. * size / w
+    return cv2.resize(img, (int(w * scale), int(h * scale)))
+
+
+def apply_blur_face(img: torch.Tensor, img_vis: np.ndarray, det: FaceDetectorResult):
+    # crop the face
+    x1, y1 = det.xmin.int(), det.ymin.int()
+    x2, y2 = det.xmax.int(), det.ymax.int()
+    roi = img[..., y1:y2, x1:x2]
+    #print(roi.shape)
+    if roi.shape[-1]==0 or roi.shape[-2]==0:
+        return
+
+    # apply blurring and put back to the visualisation image
+    roi = K.filters.gaussian_blur2d(roi, (21, 21), (100., 100.))
+    roi = K.color.rgb_to_bgr(roi)
+    img_vis[y1:y2, x1:x2] = K.tensor_to_image(roi)
+
+
+def run(image):
+    image.thumbnail((1280, 1280))
+    img_raw = np.array(image)
+
+    # preprocess
+    img = K.image_to_tensor(img_raw, keepdim=False).to(device)
+    img = K.color.bgr_to_rgb(img.float())
+
+    with torch.no_grad():
+        dets = face_detection(img)
+    dets = [FaceDetectorResult(o) for o in dets]
+
+    img_vis = img_raw.copy()
+
+    for b in dets:
+        if b.score < 0.5:
+            continue
+
+        apply_blur_face(img, img_vis, b)
+
+    return Image.fromarray(img_vis)
+
+if __name__ == "__main__":
+
+    start = time.time()
+    for _ in range(100):
+        image = Image.open("./images/crowd.jpeg")
+        _ = run(image)
+
+    print('It took', (time.time()-start)/100, 'seconds.')

requirements.txt

@@ -1,2 +1,2 @@
 opencv-python==4.5.5.62
-face_recognition==1.2.3
+kornia==0.6.3