face_detect.py

import os
import cv2
import numpy as np

# =========================================================
# GenMake Face Detection (Reliable)
# - Tries OpenCV DNN model if present (best)
# - Falls back to Haar if model not present (always works)
#
# Supported model formats:
# 1) Caffe SSD:
#    - deploy.prototxt
#    - res10_300x300_ssd_iter_140000_fp16.caffemodel  (or fp32)
#
# Put files in:
#   GenMakeEngineV1/models/face/
# =========================================================


def _models_dir() -> str:
    here = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(here, "models", "face")


def _try_load_dnn():
    """
    Returns (net, kind) or (None, None)
    kind = 'caffe'
    """
    mdir = _models_dir()

    # Caffe SSD (common, stable)
    prototxt = os.path.join(mdir, "deploy.prototxt")
    caffemodel = os.path.join(mdir, "res10_300x300_ssd_iter_140000_fp16.caffemodel")
    if os.path.exists(prototxt) and os.path.exists(caffemodel):
        net = cv2.dnn.readNetFromCaffe(prototxt, caffemodel)
        return net, "caffe"

    # Allow fp32 name too (optional)
    caffemodel2 = os.path.join(mdir, "res10_300x300_ssd_iter_140000.caffemodel")
    if os.path.exists(prototxt) and os.path.exists(caffemodel2):
        net = cv2.dnn.readNetFromCaffe(prototxt, caffemodel2)
        return net, "caffe"

    return None, None


def _haar_detector():
    path = os.path.join(cv2.data.haarcascades, "haarcascade_frontalface_default.xml")
    return cv2.CascadeClassifier(path)


def _dnn_detect_faces(net, kind: str, bgr: np.ndarray, conf_thresh: float = 0.55):
    """
    OpenCV DNN face detection. Returns list of (x1,y1,x2,y2).
    """
    h, w = bgr.shape[:2]

    blob = cv2.dnn.blobFromImage(
        cv2.resize(bgr, (300, 300)),
        1.0,
        (300, 300),
        (104.0, 177.0, 123.0),
        swapRB=False,
        crop=False
    )

    net.setInput(blob)
    det = net.forward()

    boxes = []
    # det shape: [1,1,N,7] => [batch, class, i, (img_id, label, conf, x1,y1,x2,y2)]
    for i in range(det.shape[2]):
        conf = float(det[0, 0, i, 2])
        if conf < conf_thresh:
            continue
        x1 = int(det[0, 0, i, 3] * w)
        y1 = int(det[0, 0, i, 4] * h)
        x2 = int(det[0, 0, i, 5] * w)
        y2 = int(det[0, 0, i, 6] * h)

        # clamp
        x1 = max(0, min(w - 1, x1))
        y1 = max(0, min(h - 1, y1))
        x2 = max(0, min(w - 1, x2))
        y2 = max(0, min(h - 1, y2))

        # discard tiny
        if (x2 - x1) < 50 or (y2 - y1) < 50:
            continue

        boxes.append((x1, y1, x2, y2))

    # largest first
    boxes.sort(key=lambda b: (b[2] - b[0]) * (b[3] - b[1]), reverse=True)
    return boxes


def detect_faces(bgr: np.ndarray) -> list[tuple[int, int, int, int]]:
    """
    Returns list of face boxes (x1,y1,x2,y2), largest first.
    Uses DNN if models exist, else Haar fallback.
    """
    # 1) Try DNN
    net, kind = _try_load_dnn()
    if net is not None:
        try:
            return _dnn_detect_faces(net, kind, bgr, conf_thresh=0.55)
        except Exception:
            # If anything goes wrong, fall back safely
            pass

    # 2) Haar fallback (always works, less accurate)
    h, w = bgr.shape[:2]
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    det = _haar_detector()
    faces = det.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(70, 70))

    boxes = []
    for (x, y, fw, fh) in faces:
        x1, y1 = max(0, x), max(0, y)
        x2, y2 = min(w, x + fw), min(h, y + fh)
        boxes.append((x1, y1, x2, y2))

    boxes.sort(key=lambda b: (b[2] - b[0]) * (b[3] - b[1]), reverse=True)
    return boxes