File size: 1,099 Bytes
22b8701 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
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)
|