detect_face.py

import warnings
from typing import Optional

import cv2
import numpy as np
from PIL import Image


def _load_mtcnn(device='cpu'):
    """Attempt to import and return MTCNN detector. Returns None on failure."""
    try:
        from facenet_pytorch import MTCNN
        return MTCNN(keep_all=True, device=device, post_process=False)
    except ImportError:
        warnings.warn(
            "facenet-pytorch not installed — falling back to Haar cascade. "
            "Install with: pip install facenet-pytorch",
            RuntimeWarning,
            stacklevel=3
        )
        return None


def _pil_to_bgr(image: Image.Image) -> np.ndarray:
    rgb = np.array(image.convert('RGB'))
    return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)


def _bgr_to_pil(frame: np.ndarray) -> Image.Image:
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    return Image.fromarray(rgb)


def _crop_with_margin(image: Image.Image, x1: int, y1: int, x2: int, y2: int, margin: int) -> tuple:
    """Apply margin, clamp to image bounds, return (cropped_pil, bbox_xyxy)."""
    w, h = image.size
    x1 = max(0, x1 - margin)
    y1 = max(0, y1 - margin)
    x2 = min(w, x2 + margin)
    y2 = min(h, y2 + margin)
    return image.crop((x1, y1, x2, y2)), (x1, y1, x2 - x1, y2 - y1)  # bbox as (x,y,w,h)


def detect_and_crop_faces(
    image,
    method: str = 'mtcnn',
    margin: int = 20,
    device: str = 'cpu'
) -> list[tuple]:
    """
    Detect faces and return crops.

    Args:
        image: PIL Image, numpy array (H,W,3 BGR or RGB), or file path str.
        method: 'mtcnn' (with auto-fallback) or 'haar'.
        margin: pixel margin to add around each detected bbox.
        device: torch device string used by MTCNN.

    Returns:
        List of (face_crop_PIL, bbox_xywh) tuples.
        bbox_xywh is (x, y, w, h) in original image coords, or None if no face detected.
        If no face is detected the full image (resized to 224) is returned with bbox=None.
    """
    # Normalise to PIL RGB
    if isinstance(image, str):
        pil_img = Image.open(image).convert('RGB')
    elif isinstance(image, np.ndarray):
        if image.ndim == 3 and image.shape[2] == 3:
            # Assume BGR (OpenCV convention)
            pil_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        else:
            pil_img = Image.fromarray(image).convert('RGB')
    elif isinstance(image, Image.Image):
        pil_img = image.convert('RGB')
    else:
        raise TypeError(f"Unsupported image type: {type(image)}")

    faces = []

    if method == 'mtcnn':
        mtcnn = _load_mtcnn(device=device)
        if mtcnn is not None:
            faces = _detect_mtcnn(mtcnn, pil_img, margin)
        else:
            method = 'haar'

    if method == 'haar' or (method == 'mtcnn' and not faces):
        faces = _detect_haar(pil_img, margin)

    if not faces:
        warnings.warn(
            "No face detected — running inference on full image.",
            RuntimeWarning,
            stacklevel=2
        )
        fallback = pil_img.resize((224, 224))
        return [(fallback, None)]

    return faces


def _detect_mtcnn(mtcnn, pil_img: Image.Image, margin: int) -> list[tuple]:
    import torch
    boxes, probs = mtcnn.detect(pil_img)
    if boxes is None or len(boxes) == 0:
        return []

    results = []
    for box in boxes:
        x1, y1, x2, y2 = (int(v) for v in box)
        crop, bbox = _crop_with_margin(pil_img, x1, y1, x2, y2, margin)
        results.append((crop, bbox))
    return results


def _detect_haar(pil_img: Image.Image, margin: int) -> list[tuple]:
    bgr = _pil_to_bgr(pil_img)
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)

    cascade_path = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml'
    cascade = cv2.CascadeClassifier(cascade_path)

    detected = cascade.detectMultiScale(
        gray,
        scaleFactor=1.1,
        minNeighbors=5,
        minSize=(30, 30)
    )

    if len(detected) == 0:
        return []

    results = []
    for (x, y, w, h) in detected:
        crop, bbox = _crop_with_margin(pil_img, x, y, x + w, y + h, margin)
        results.append((crop, bbox))
    return results