face_engine.py

"""Thin wrapper around InsightFace's FaceAnalysis (RetinaFace detector + ArcFace embedder).

Intentionally duplicated from ../app/face_engine.py: this folder is a separate,
self-contained deployable (Hugging Face Spaces web demo) with its own
requirements.txt (headless OpenCV, no GUI deps) and its own tuning defaults for
a shared/free CPU, independent of the NUC/local app in ../app.
"""
from typing import List, Optional, Tuple

import numpy as np
import onnxruntime
from insightface.app import FaceAnalysis
from insightface.app.common import Face

# We only ever read .bbox, .det_score, .normed_embedding, so the landmark_3d_68,
# landmark_2d_106 and genderage models are pure overhead: excluding them via
# allowed_modules skips loading AND running 3 extra ONNX forward passes per face.
_REQUIRED_MODULES = ["detection", "recognition"]


class FaceEngine:
    """Detects faces and produces L2-normalized 512-d ArcFace embeddings."""

    def __init__(
        self,
        det_size: Tuple[int, int] = (320, 320),
        ctx_id: int = -1,
        model_name: str = "buffalo_s",
        num_threads: Optional[int] = None,
    ) -> None:
        # ctx_id=-1 -> CPUExecutionProvider only; ctx_id>=0 -> CUDAExecutionProvider device id.
        providers = ["CPUExecutionProvider"] if ctx_id < 0 else ["CUDAExecutionProvider", "CPUExecutionProvider"]

        sess_options = None
        if num_threads is not None and num_threads > 0:
            sess_options = onnxruntime.SessionOptions()
            sess_options.intra_op_num_threads = num_threads
            sess_options.inter_op_num_threads = 1

        self.app = FaceAnalysis(
            name=model_name,
            providers=providers,
            allowed_modules=_REQUIRED_MODULES,
            sess_options=sess_options,
        )
        self.app.prepare(ctx_id=ctx_id, det_size=det_size)

    def detect(self, frame_bgr: np.ndarray) -> List[Face]:
        """Returns InsightFace Face objects with .bbox, .det_score, .normed_embedding (512-d, L2-normed)."""
        return self.app.get(frame_bgr)