mule-image/backend/app/services/vision/faces.py

"""
Face detection (YuNet) + recognition (SFace) using ONNX Runtime.

Both models are from opencv_zoo (Apache 2.0 license).
Expects {models_dir}/face/:
  - yunet.onnx     (~75 KB)
  - sface.onnx     (~37 MB, 128-d embeddings)
"""
import logging
from pathlib import Path

import numpy as np
import onnxruntime as ort

from app.config import VisionSettings
from app.services.vision.base import FaceProcessor, FaceDetection

logger = logging.getLogger(__name__)

_YUNET_INPUT_SIZE = 640


def _align_face(image: np.ndarray, landmarks: np.ndarray) -> np.ndarray:
    """Align and crop a 112x112 face patch using 5-point landmarks.
    landmarks shape: (5, 2) — left_eye, right_eye, nose, left_mouth, right_mouth."""
    from PIL import Image
    import math

    left_eye = landmarks[0]
    right_eye = landmarks[1]

    dx = right_eye[0] - left_eye[0]
    dy = right_eye[1] - left_eye[1]
    angle = math.degrees(math.atan2(dy, dx))
    eye_center = ((left_eye[0] + right_eye[0]) / 2, (left_eye[1] + right_eye[1]) / 2)
    eye_dist = math.sqrt(dx * dx + dy * dy)
    scale = 64.0 / max(eye_dist, 1e-6)  # target: eyes at ~64px apart in 112x112

    img = Image.fromarray(image)
    img = img.rotate(-angle, center=eye_center, resample=Image.BICUBIC)

    cx, cy = eye_center
    half = 56.0 / scale
    crop = img.crop((int(cx - half), int(cy - half * 0.8), int(cx + half), int(cy + half * 1.2)))
    crop = crop.resize((112, 112), Image.BICUBIC)
    return np.array(crop, dtype=np.float32)


class YuNetSFaceProcessor(FaceProcessor):
    def __init__(self, settings: VisionSettings):
        face_dir = Path(settings.models_dir) / "face"
        yunet_path = face_dir / "yunet.onnx"
        sface_path = face_dir / "sface.onnx"

        opts = ort.SessionOptions()
        opts.inter_op_num_threads = 2
        opts.intra_op_num_threads = 2

        logger.info("Loading YuNet from %s", yunet_path)
        self._detector = ort.InferenceSession(str(yunet_path), opts, providers=["CPUExecutionProvider"])

        logger.info("Loading SFace from %s", sface_path)
        self._recognizer = ort.InferenceSession(str(sface_path), opts, providers=["CPUExecutionProvider"])

        self._min_face_size = settings.faces.min_face_size
        self._score_threshold = settings.faces.recognition_threshold

    def process(self, image: np.ndarray) -> list[FaceDetection]:
        orig_h, orig_w = image.shape[:2]

        # Scale + letterbox to exactly 640x640 (YuNet fixed input)
        scale = min(_YUNET_INPUT_SIZE / orig_w, _YUNET_INPUT_SIZE / orig_h)
        new_w = int(orig_w * scale)
        new_h = int(orig_h * scale)

        from PIL import Image as PILImage
        resized = np.array(
            PILImage.fromarray(image).resize((new_w, new_h), PILImage.BICUBIC),
            dtype=np.uint8,
        )

        # Letterbox pad to 640x640
        canvas = np.full((_YUNET_INPUT_SIZE, _YUNET_INPUT_SIZE, 3), 128, dtype=np.uint8)
        pad_y = (_YUNET_INPUT_SIZE - new_h) // 2
        pad_x = (_YUNET_INPUT_SIZE - new_w) // 2
        canvas[pad_y:pad_y + new_h, pad_x:pad_x + new_w] = resized

        # YuNet expects BGR
        bgr = canvas[:, :, ::-1].copy()

        # Run detection
        det_input = self._detector.get_inputs()[0]
        blob = bgr.astype(np.float32)[np.newaxis]  # (1, H, W, 3)
        # Some YuNet ONNX exports expect (1, 3, H, W)
        if det_input.shape and len(det_input.shape) == 4 and det_input.shape[1] == 3:
            blob = blob.transpose(0, 3, 1, 2)

        detections_raw = self._detector.run(None, {det_input.name: blob})
        dets = detections_raw[0]  # (N, 15): x,y,w,h,score, 5x landmark pairs

        if dets is None or len(dets) == 0:
            return []

        results = []
        for det in dets:
            score = float(det[4]) if len(det) > 4 else float(det[-1])
            if score < self._score_threshold:
                continue

            x, y, w, h = det[0], det[1], det[2], det[3]

            # Filter small faces
            face_size = max(w, h) / scale
            if face_size < self._min_face_size:
                continue

            # Remove letterbox padding and rescale to original image coords
            x1 = (x - pad_x) / scale
            y1 = (y - pad_y) / scale
            x2 = (x + w - pad_x) / scale
            y2 = (y + h - pad_y) / scale

            bbox = [
                max(0, x1 / orig_w),
                max(0, y1 / orig_h),
                min(1, x2 / orig_w),
                min(1, y2 / orig_h),
            ]

            # Extract landmarks (5 points) for alignment
            if len(det) >= 15:
                landmarks = det[5:15].reshape(5, 2)
                landmarks[:, 0] = (landmarks[:, 0] - pad_x) / scale
                landmarks[:, 1] = (landmarks[:, 1] - pad_y) / scale
            else:
                # Fallback: no landmarks, skip recognition
                continue

            # Align face for recognition
            face_crop = _align_face(image, landmarks)

            # SFace expects (1, 3, 112, 112) float32, BGR, normalized
            face_bgr = face_crop[:, :, ::-1].copy()
            face_blob = (face_bgr / 255.0).transpose(2, 0, 1)[np.newaxis].astype(np.float32)

            rec_input = self._recognizer.get_inputs()[0].name
            embedding = self._recognizer.run(None, {rec_input: face_blob})[0][0]
            embedding = embedding / np.linalg.norm(embedding)

            results.append(FaceDetection(
                bbox=bbox,
                embedding=embedding.astype(np.float32),
                quality=score,
            ))

        return results

    @property
    def embedding_dim(self) -> int:
        return 128