initial

.gitignore

@@ -0,0 +1,4 @@
+venv
+**/__pycache__/
+**/.DS_Store
+**/*.pt

app.py

@@ -0,0 +1,25 @@
+import gradio as gr
+from detect_faces import MAX_OUTPUT, check_and_download_model, detect_faces
+
+
+def process_image(input_image, box_margin):
+    if input_image is None:
+        return INITIAL_OUTPUTS
+    output_images = detect_faces(input_image, box_margin=box_margin)
+    output_images = [gr.Image(image, visible=True) for image in output_images]
+    output_images += INITIAL_OUTPUTS[: 100 - len(output_images)]
+    return output_images
+
+
+INITIAL_OUTPUTS = [gr.Image(visible=False) for _ in range(MAX_OUTPUT)]
+outputs = INITIAL_OUTPUTS.copy()
+
+iface = gr.Interface(
+    fn=process_image,
+    inputs=[gr.Image(type="filepath"), gr.Slider(0, 40)],
+    outputs=outputs,
+    allow_flagging="never",
+    live=True,
+)
+check_and_download_model()
+iface.launch()

detect_faces.py

@@ -0,0 +1,49 @@
+from ultralytics import YOLO
+import os
+import requests
+from PIL import Image
+from functools import lru_cache
+
+MODEL_LINK = (
+    "https://github.com/akanametov/yolov8-face/releases/download/v0.0.0/yolov8l-face.pt"
+)
+MODEL_PATH = "yolov8l-face.pt"
+MAX_OUTPUT = 10
+
+model = YOLO(MODEL_PATH)
+
+
+def check_and_download_model():
+    if os.path.exists(MODEL_PATH):
+        return
+    response = requests.get(MODEL_LINK, timeout=60)
+    with open(MODEL_PATH, "wb") as f:
+        f.write(response.content)
+
+
+def adjust_box(b, image_shape, margin=10):
+    x1, y1, x2, y2 = b
+    x1 = max(0, x1 - margin)
+    y1 = max(0, y1 - margin)
+    x2 = min(image_shape[0], x2 + margin)
+    y2 = min(image_shape[1], y2 + margin)
+    return (int(x1), int(y1), int(x2), int(y2))
+
+
+@lru_cache(maxsize=None)
+def predict(path):
+    results = model.predict(path)
+    orimg = Image.open(path)
+    return results, orimg
+
+
+def detect_faces(image, box_margin=0):
+    faces = []
+    results, orimg = predict(image)
+    for box in results[0].boxes:
+        coords = box.xyxy[0]
+        coords = adjust_box(coords, orimg.size, margin=box_margin)
+        crop = orimg.crop(coords)
+        crop = crop.resize((512, 512))
+        faces.append(crop)
+    return faces

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+ultralytics
+dill