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

"""
Face detection (YuNet) + recognition (SFace) using OpenCV DNN.

YuNet is loaded via cv2.FaceDetectorYN which handles the multi-scale
anchor decoding and NMS internally. SFace recognition uses raw ONNX
Runtime for the 128-d embedding.

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

import numpy as np
import cv2
import onnxruntime as ort

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

logger = logging.getLogger(__name__)


def _align_face(image: np.ndarray, landmarks: np.ndarray) -> np.ndarray:
    """Align and crop a 112x112 face patch using 5-point landmarks."""
    left_eye = landmarks[0]
    right_eye = landmarks[1]

    dx = right_eye[0] - left_eye[0]
    dy = right_eye[1] - left_eye[1]
    angle = np.degrees(np.arctan2(dy, dx))
    eye_center = ((left_eye[0] + right_eye[0]) / 2, (left_eye[1] + right_eye[1]) / 2)
    eye_dist = np.sqrt(dx * dx + dy * dy)

    M = cv2.getRotationMatrix2D(eye_center, angle, 1.0)
    rotated = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))

    # Crop around face center
    scale = 64.0 / max(eye_dist, 1e-6)
    cx, cy = eye_center
    half = 56.0 / scale
    x1 = max(0, int(cx - half))
    y1 = max(0, int(cy - half * 0.8))
    x2 = min(rotated.shape[1], int(cx + half))
    y2 = min(rotated.shape[0], int(cy + half * 1.2))
    crop = rotated[y1:y2, x1:x2]

    if crop.size == 0:
        return np.zeros((112, 112, 3), dtype=np.float32)

    return cv2.resize(crop, (112, 112)).astype(np.float32)


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

        # YuNet via OpenCV's FaceDetectorYN — handles anchor decoding + NMS
        self._detector = cv2.FaceDetectorYN.create(
            yunet_path,
            "",
            (640, 640),
            settings.faces.recognition_threshold,
            0.3,  # NMS threshold
            5000,  # top_k
        )
        logger.info("YuNet face detector loaded via OpenCV")

        # SFace via ONNX Runtime
        from app.services.vision.providers import create_session
        ort.set_default_logger_severity(3)
        self._recognizer = create_session(sface_path, configured_providers=settings.execution_providers)
        logger.info("SFace recognizer loaded via ONNX Runtime")

        self._min_face_size = settings.faces.min_face_size

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

        # Convert RGB → BGR for OpenCV
        bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        # Set input size to actual image dimensions
        self._detector.setInputSize((orig_w, orig_h))

        # Detect faces
        _, faces_raw = self._detector.detect(bgr)

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

        results = []
        for face in faces_raw:
            # face: [x, y, w, h, right_eye_x, right_eye_y, left_eye_x, left_eye_y,
            #        nose_x, nose_y, right_mouth_x, right_mouth_y, left_mouth_x, left_mouth_y, score]
            x, y, w, h = int(face[0]), int(face[1]), int(face[2]), int(face[3])
            score = float(face[14])

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

            # Normalized bbox
            bbox = [
                max(0, x / orig_w),
                max(0, y / orig_h),
                min(1, (x + w) / orig_w),
                min(1, (y + h) / orig_h),
            ]

            # Extract 5-point landmarks for alignment
            landmarks = np.array([
                [face[4], face[5]],   # right eye
                [face[6], face[7]],   # left eye
                [face[8], face[9]],   # nose
                [face[10], face[11]], # right mouth
                [face[12], face[13]], # left mouth
            ], dtype=np.float32)

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

            # SFace expects (1, 3, 112, 112) float32, BGR
            face_bgr = cv2.cvtColor(face_crop.astype(np.uint8), cv2.COLOR_RGB2BGR)
            face_blob = (face_bgr.astype(np.float32) / 255.0).transpose(2, 0, 1)[np.newaxis]

            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