detect_faces_compare.py

detect_faces_compare.py

#!/usr/bin/env python3
"""Run four face detectors (MediaPipe, RetinaFace, SCRFD/InsightFace, YOLOv11-face)
on all images in the current folder and save comparison outputs.

Output naming: ``{stem}_{modelname}{ext}`` — one file per model per input.
"""

import sys
import urllib.request
from pathlib import Path

import cv2


# --- Configuration -----------------------------------------------------------
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff", ".tif"}
BOX_COLOR = (0, 255, 0)  # green BGR
BOX_THICKNESS = 2

DETECTORS = []


def register(name):
    """Decorator to register a detector by *name*."""
    def deco(fn):
        DETECTORS.append((name, fn))
        return fn
    return deco


# ---------------------------------------------------------------------------
# 1. MediaPipe
# ---------------------------------------------------------------------------
@register("mediapipe")
def build_mediapipe():
    import mediapipe as mp
    fd = mp.solutions.face_detection.FaceDetection(
        model_selection=1, min_detection_confidence=0.5)

    def detect(img_bgr):
        rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        rgb.flags.writeable = False
        results = fd.process(rgb)
        faces = []
        if results.detections:
            h, w = img_bgr.shape[:2]
            for d in results.detections:
                bb = d.location_data.relative_bounding_box
                faces.append((
                    max(0, int(bb.xmin * w)),
                    max(0, int(bb.ymin * h)),
                    min(w - 1, int((bb.xmin + bb.width) * w)),
                    min(h - 1, int((bb.ymin + bb.height) * h)),
                    float(d.score[0]),
                ))
        return faces
    return detect


# ---------------------------------------------------------------------------
# 2. RetinaFace
# ---------------------------------------------------------------------------
@register("retinaface")
def build_retinaface():
    from retinaface import RetinaFace  # noqa: F811

    def detect(img_bgr):
        rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        h, w = img_bgr.shape[:2]
        result = RetinaFace.detect_faces(rgb)
        faces = []
        if isinstance(result, dict):
            for _key, val in result.items():
                if isinstance(val, dict) and "facial_area" in val:
                    x1, y1, x2, y2 = val["facial_area"]
                    score = float(val.get("score", 0.999))
                    faces.append((
                        max(0, int(x1)), max(0, int(y1)),
                        min(w - 1, int(x2)), min(h - 1, int(y2)),
                        score,
                    ))
        return faces
    return detect


# ---------------------------------------------------------------------------
# 3. SCRFD  (InsightFace — buffalo_l pack)
# ---------------------------------------------------------------------------
@register("scrfd")
def build_scrfd():
    from insightface.app import FaceAnalysis

    # buffalo_l includes SCRFD — detection only (no recognition/landmarks needed)
    app = FaceAnalysis(name="buffalo_l", allowed_modules=["detection"])
    app.prepare(ctx_id=-1)  # -1 = CPU, 0+ = GPU

    def detect(img_bgr):
        rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
        h, w = img_bgr.shape[:2]
        detections = app.get(rgb)
        faces = []
        for d in detections:
            x1, y1, x2, y2 = d.bbox.astype(int).tolist()
            conf = float(d.det_score)
            faces.append((
                max(0, x1), max(0, y1),
                min(w - 1, x2), min(h - 1, y2),
                conf,
            ))
        return faces
    return detect


# ---------------------------------------------------------------------------
# 4. YOLOv11-face
# ---------------------------------------------------------------------------
@register("yolo")
def build_yolo_face():
    from ultralytics import YOLO

    # Community face-detection weights from akanametov/yolo-face
    model_url = (
        "https://github.com/akanametov/yolo-face/releases/download/1.0.0/"
        "yolov11n-face.pt"
    )
    model_path = Path.home() / ".cache" / "face-detect" / "yolov11n-face.pt"
    if not model_path.exists():
        model_path.parent.mkdir(parents=True, exist_ok=True)
        print(f"  Downloading YOLOv11-face model to {model_path} ...")
        urllib.request.urlretrieve(model_url, str(model_path))
        print("  Download complete.")

    model = YOLO(str(model_path))

    def detect(img_bgr):
        h, w = img_bgr.shape[:2]
        results = model(img_bgr, verbose=False)
        faces = []
        for r in results:
            if r.boxes is None:
                continue
            for box in r.boxes:
                x1, y1, x2, y2 = box.xyxy[0].tolist()
                conf = float(box.conf[0])
                faces.append((
                    max(0, int(x1)), max(0, int(y1)),
                    min(w - 1, int(x2)), min(h - 1, int(y2)),
                    conf,
                ))
        return faces
    return detect


# ---------------------------------------------------------------------------
# Drawing helpers
# ---------------------------------------------------------------------------
def draw_boxes(img_bgr, faces, detector_name):
    """Draw bounding boxes + labels on a copy of *img_bgr*.  Returns the copy."""
    img = img_bgr.copy()
    for (x1, y1, x2, y2, conf) in faces:
        cv2.rectangle(img, (x1, y1), (x2, y2), BOX_COLOR, BOX_THICKNESS)
        label = f"{detector_name} {conf:.2f}"
        label_y = y1 - 8 if y1 > 20 else y1 + 20
        cv2.putText(img, label, (x1, label_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, BOX_COLOR, 2)
    return img


def find_image_files(folder):
    """Return original image files, skipping any output files from previous runs."""
    output_suffixes = (
        "_mediapipe", "_retinaface", "_scrfd", "_yolo",
        "_boundingbox",
    )
    images = []
    for entry in sorted(folder.iterdir()):
        if not entry.is_file():
            continue
        if entry.suffix.lower() not in IMAGE_EXTENSIONS:
            continue
        if any(s in entry.stem for s in output_suffixes):
            continue
        images.append(entry)
    return images


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    folder = Path.cwd()
    images = find_image_files(folder)
    if not images:
        print("No image files found.")
        sys.exit(0)

    print(f"Found {len(images)} image(s).\n")
    print(f"Detectors: {', '.join(name for name, _ in DETECTORS)}\n")

    # Build all detectors (they may download models on first call).
    # IMPORTANT: RetinaFace must initialise first — its Keras model conflicts
    # with MediaPipe's TF Lite backend if MediaPipe loads first.
    def init_order(name_buildfn):
        name, _ = name_buildfn
        return (0 if name == "retinaface" else 1, name)

    detectors = []
    for name, build_fn in sorted(DETECTORS, key=init_order):
        print(f"Initialising {name} ...")
        try:
            detectors.append((name, build_fn()))
            print(f"  {name} ready.")
        except Exception as exc:
            print(f"  [SKIP] {name} failed to initialise: {exc}")

    if not detectors:
        print("No detectors available. Exiting.")
        sys.exit(1)

    print()

    # Run every detector on every image
    for img_path in images:
        img_bgr = cv2.imread(str(img_path))
        if img_bgr is None:
            print(f"[SKIP] Could not read: {img_path.name}")
            continue
        if img_bgr.ndim >= 3 and img_bgr.shape[2] == 4:
            img_bgr = img_bgr[:, :, :3]

        print(f"Processing: {img_path.name}  ({img_bgr.shape[1]}x{img_bgr.shape[0]})")

        for det_name, detect_fn in detectors:
            try:
                faces = detect_fn(img_bgr)
                out_img = draw_boxes(img_bgr, faces, det_name)
                out_path = folder / f"{img_path.stem}_{det_name}{img_path.suffix}"
                cv2.imwrite(str(out_path), out_img)
                print(f"  {det_name:14s}: {len(faces)} face(s)  -> {out_path.name}")
            except Exception as exc:
                print(f"  {det_name:14s}: ERROR — {exc}")

    print("\nDone.  Compare the *_<modelname>.jpg outputs side by side.")


if __name__ == "__main__":
    main()