src/FaceDetector/face_detector.py

from typing import NamedTuple, Optional, Tuple

from insightface.model_zoo import model_zoo
import numpy as np
from pathlib import Path


class Detection(NamedTuple):
    bbox: Optional[np.ndarray]
    score: Optional[np.ndarray]
    key_points: Optional[np.ndarray]


class FaceDetector:
    def __init__(
        self,
        model_path: Path,
        det_thresh: float = 0.5,
        det_size: Tuple[int, int] = (640, 640),
        mode: str = "None",
        device: str = "cpu",
    ):
        self.det_thresh = det_thresh
        self.mode = mode
        self.device = device
        self.handler = model_zoo.get_model(str(model_path))
        ctx_id = -1 if device == "cpu" else 0
        self.handler.prepare(ctx_id, input_size=det_size)

    def __call__(self, img: np.ndarray, max_num: int = 0) -> Detection:
        bboxes, kpss = self.handler.detect(
            img, threshold=self.det_thresh, max_num=max_num, metric="default"
        )
        if bboxes.shape[0] == 0:
            return Detection(None, None, None)

        return Detection(bboxes[..., :-1], bboxes[..., -1], kpss)