moving to self trained face segmentation model

app.py

@@ -1,67 +1,67 @@
 import cv2
 import gradio as gr
-from PIL import Image
+from typing import Union, Tuple
+from PIL import Image, ImageOps
 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))
-    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)
+model = torch.jit.load('./model/model.pt').eval()
+
+def resize_with_padding(img: Image.Image, expected_size: Tuple[int, int]) -> Image.Image:
+    img.thumbnail((expected_size[0], expected_size[1]))
+    delta_width = expected_size[0] - img.size[0]
+    delta_height = expected_size[1] - img.size[1]
+    pad_width = delta_width // 2
+    pad_height = delta_height // 2
+    padding = (pad_width, pad_height, delta_width - pad_width, delta_height - pad_height)
+    return ImageOps.expand(img, padding), padding
+
+def preprocess_image(img: Image.Image, size: int = 512) -> Tuple[Image.Image, torch.tensor, Tuple[int]]:    
+    pil_img, padding = resize_with_padding(img, (size, size))
+    
+    img = (np.array(pil_img).astype(np.float32) / 255) - np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
+    img = img / np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)
+    img = np.transpose(img, (2, 0, 1))
+    
+    return pil_img, torch.tensor(img[None]), padding
+
+def soft_blur_with_mask(image: Image.Image, mask: torch.tensor, padding: Tuple[int]) -> Image.Image:
+    image = np.array(image)
+    # Create a blurred copy of the original image.
+    blurred_image = cv2.GaussianBlur(image, (221, 221), sigmaX=20, sigmaY=20)
+    image_height, image_width = image.shape[:2]
+    mask = cv2.resize(mask.astype(np.uint8), (image_width, image_height), interpolation=cv2.INTER_NEAREST)
+    # Blurring the mask itself to get a softer mask with no firm edges
+    mask = cv2.GaussianBlur(mask.astype(np.float32), (11, 11), 10, 10)[:, :, None]
+
+    # Take the blurred image where the mask it positive, and the original image where the image is original
+    image = (mask * blurred_image + (1.0 - mask) * image)
+    pad_w, pad_h, _, _ = padding
+    img_w, img_h, _ = image.shape
+    image = image[(pad_h):(img_h-pad_h), (pad_w):(img_w-pad_w), :]
+    return Image.fromarray(image.astype(np.uint8))
+
+def run(image, size):
+    pil_image, torch_image, padding = preprocess_image(image, size=size)
+
+    with torch.inference_mode():
+        mask = model(torch_image)
+    mask = mask.argmax(dim=1).numpy().squeeze()
+
+    return soft_blur_with_mask(pil_image, mask, padding)
 
 content_image_input = gr.inputs.Image(label="Content Image", type="pil")
+model_image_size = gr.inputs.Radio([256, 384, 512, 1024], type="value", default=512, label="Inference size")
 
 description="Privacy first! Upload an image of a groupf of people and blur their faces automatically."
 article="""
-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.
+Demo built on top of a face segmentation model trained from scratch with IceVision on the 
+<a href='https://github.com/microsoft/FaceSynthetics' target='_blank'>FaceSynthetics</a> dataset.
 """
-examples = [["./images/family.jpeg"], ["./images/crowd.jpeg"], ["./images/crowd1.jpeg"]]
+examples = [["./images/girls.jpeg", 384], ["./images/kid.jpeg", 256], ["./images/family.jpeg", 512], ["./images/crowd1.jpeg", 1024], ["./images/crowd2.jpeg", 1024]]
 
 app_interface = gr.Interface(fn=run,
-                             inputs=[content_image_input],
+                             inputs=[content_image_input, model_image_size],
                              outputs="image",
                              title="Blurry Faces",
                              description=description,

model/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e310f25944aaf0a35a334798e72aca4494dd19f3785225042017743ecd37757
+size 165321408