face_utils.py

from typing import Optional

import numpy as np
from facenet_pytorch import MTCNN
from PIL import Image


_mtcnn = MTCNN(keep_all=True, device="cpu")


def _clip_box(box: np.ndarray, width: int, height: int) -> tuple[int, int, int, int]:
    x1, y1, x2, y2 = box.astype(int).tolist()
    x1 = max(0, min(x1, width - 1))
    y1 = max(0, min(y1, height - 1))
    x2 = max(0, min(x2, width - 1))
    y2 = max(0, min(y2, height - 1))
    return x1, y1, x2, y2


def detect_largest_face(image: Image.Image) -> Optional[Image.Image]:
    """
    Detect and return the largest face crop from a PIL image.
    Returns None if no valid face is detected.
    """
    rgb_image = image.convert("RGB")
    img_array = np.array(rgb_image)
    boxes, _ = _mtcnn.detect(rgb_image)

    if boxes is None or len(boxes) == 0:
        return None

    widths = boxes[:, 2] - boxes[:, 0]
    heights = boxes[:, 3] - boxes[:, 1]
    areas = widths * heights
    largest_idx = int(np.argmax(areas))
    x1, y1, x2, y2 = _clip_box(boxes[largest_idx], img_array.shape[1], img_array.shape[0])

    if x2 <= x1 or y2 <= y1:
        return None

    face_crop = img_array[y1:y2, x1:x2]
    if face_crop.size == 0:
        return None
    return Image.fromarray(face_crop)