diff --git a/src/facebias/detectors/__init__.py b/src/facebias/detectors/__init__.py index e659dac..b25d35e 100644 --- a/src/facebias/detectors/__init__.py +++ b/src/facebias/detectors/__init__.py @@ -1,8 +1,13 @@ # -*- coding: utf-8 -*- +import logging from abc import ABC, abstractmethod +from collections import OrderedDict import numpy as np +from facebias import FaceBox + +logger = logging.getLogger("facebias") class NoFaceDetectedError(Exception): @@ -20,7 +25,38 @@ class BaseFaceDetector(ABC): and perform face detection, whenever a call to `predict` is performed.""" @abstractmethod - def detect(self, im) -> np.ndarray: + def detect(self, im: np.ndarray) -> np.ndarray: """Detects faces in image `im`. Throws `NoFaceDetectedError` if `im` has no valid face(s).""" pass + + +def get_face_boxes( + ims: dict[str, np.ndarray], + detector: BaseFaceDetector, + num_faces: int=1 +) -> OrderedDict[str, list[FaceBox]]: + """ + Parameters + ---------- + ims: dict[str, np.ndarray], + detector: BaseDetector, + num_faces: int | 1 + + Returns + ------- + all_boxes: OrderedDict[str, list[FaceBox]] + + """ + all_boxes: OrderedDict[str, list[FaceBox]] = OrderedDict() + for p, im in ims.items(): + boxes = detector.detect(im) + if len(boxes) > num_faces: + name = "face" if num_faces == 1 else "faces" + logger.warning( + f"Should only have {num_faces} {name} per image." + " Ignoring all others." + ) + + all_boxes[p] = boxes[:num_faces] + return all_boxes