fast-face-python 0.1.1__py3-none-any.whl

fast_face/__init__.py

@@ -0,0 +1,6 @@
+from .models.factory import FaceModelFactory
+from .models.yunet import YuNet
+from .models.retinaface import RetinaFace
+from .models.adaface import AdaFace
+
+__all__ = ["FaceModelFactory", "YuNet", "RetinaFace", "AdaFace"]

fast_face/models/__init__.py

@@ -0,0 +1,11 @@
+from .base import BaseFaceModel
+from .factory import FaceModelFactory
+from .session import ONNXSession
+from .yunet import YuNet
+
+__all__ = [
+    "BaseFaceModel",
+    "FaceModelFactory",
+    "ONNXSession",
+    "YuNet",
+]

fast_face/models/adaface.py

@@ -0,0 +1,91 @@
+from typing import Any
+
+import cv2
+import numpy as np
+
+from ..schema import ProviderType
+from ..tools import align_face
+from .base_recognition import BaseRecognitionModel
+
+
+class AdaFace(BaseRecognitionModel):
+    """AdaFace face recognition model for embedding extraction.
+
+    Supports IR-18, IR-50, and IR-101 backbones.
+    Expects RGB input images; converts to BGR internally as required by AdaFace.
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        align_size: tuple[int, int] = (112, 112),
+        providers: list[ProviderType] | None = None,
+        **kwargs,
+    ):
+        if providers is None:
+            providers = ["CPUExecutionProvider"]
+        super().__init__(
+            model_path=model_path,
+            providers=providers,
+            sess_options=kwargs.get("sess_options"),
+        )
+        self.align_size = align_size
+
+    def preprocess(self, imgs: list[np.ndarray]) -> np.ndarray:
+        """Preprocess aligned face images: RGB->BGR, normalize to [-1, 1], transpose to NCHW."""
+        processed = []
+        for img in imgs:
+            bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+            h, w = bgr.shape[:2]
+            if (w, h) != self.align_size:
+                bgr = cv2.resize(bgr, self.align_size)
+
+            processed.append(bgr)
+
+        batch = np.stack(processed, axis=0).astype(np.float32)
+
+        batch = (batch / 255.0 - 0.5) / 0.5
+
+        batch = np.transpose(batch, (0, 3, 1, 2))
+        return batch
+
+    def extract(
+        self,
+        imgs: list[np.ndarray] | np.ndarray,
+        landmarks: list[np.ndarray | dict[str, Any]] | None = None,
+    ) -> np.ndarray:
+        """Extract L2-normalized face embeddings from images.
+
+        Args:
+            imgs: List of face images in RGB, or a single image array.
+                  If landmarks are provided, images should be the full/cropped
+                  frames from which faces will be aligned.
+            landmarks: Optional list of landmarks for face alignment.
+                       Each element can be an np.ndarray of shape (5, 2) or (10,),
+                       or a dictionary with landmark keys.
+
+        Returns:
+            np.ndarray: L2-normalized embeddings of shape (N, embedding_dim).
+        """
+        if isinstance(imgs, np.ndarray) and imgs.ndim == 3:
+            imgs = [imgs]
+
+        if landmarks is not None:
+            aligned_imgs = []
+            for img, lmk in zip(imgs, landmarks):
+                aligned = align_face(img, lmk, align_size=self.align_size)
+                aligned_imgs.append(aligned)
+        else:
+            aligned_imgs = imgs
+
+        batch = self.preprocess(aligned_imgs)
+        outputs = self.session(batch)
+
+        embeddings = outputs[0]
+
+        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+        norms = np.maximum(norms, 1e-10)
+        embeddings = embeddings / norms
+
+        return embeddings

fast_face/models/base.py

@@ -0,0 +1,108 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+from ..schema import ProviderType
+from ..tools import nms, parse_det
+from .session import ONNXSession
+
+
+class BaseFaceModel(ABC):
+    def __init__(
+        self,
+        model_path: str,
+        conf_threshold: float = 0.6,
+        nms_threshold: float = 0.4,
+        top_k: int = 5000,
+        keep_top_k: int = 1000,
+        providers: list[ProviderType] | None = None,
+        sess_options: ort.SessionOptions | None = None,
+    ):
+        if providers is None:
+            providers = ["CPUExecutionProvider"]
+        self.conf_threshold = conf_threshold
+        self.nms_threshold = nms_threshold
+        self.top_k = top_k
+        self.keep_top_k = keep_top_k
+
+        self.session = ONNXSession(
+            model_path=model_path, providers=providers, sess_options=sess_options
+        )
+        self.input_name = self.session.session.get_inputs()[0].name
+
+    @abstractmethod
+    def preprocess(self, imgs: np.ndarray) -> np.ndarray:
+        """Preprocess images for the specific model."""
+
+    @abstractmethod
+    def post_process(
+        self,
+        outputs: list[np.ndarray],
+        original_shapes: list[tuple[int, int]],
+        preprocessed_shape: tuple[int, int],
+    ) -> list[np.ndarray]:
+        """Post-process ONNX outputs to bounding boxes, confidence, and landmarks."""
+
+    def detect(
+        self,
+        imgs: str | list[str] | np.ndarray | list[np.ndarray],
+        return_dict: bool = False,
+    ) -> list[np.ndarray | list[dict[str, Any]]]:
+        """Run face detection inference on input images."""
+        original_shapes = []
+        if isinstance(imgs, str):
+            imgs = [imgs]
+
+        if isinstance(imgs, list):
+            raw_imgs = []
+            for item in imgs:
+                if isinstance(item, str):
+                    img = cv2.imread(item)
+                    if img is None:
+                        raise ValueError(f"Could not load image from {item}")
+                    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                else:
+                    img = item
+                original_shapes.append(img.shape)
+                raw_imgs.append(img)
+            preprocessed_imgs = self.preprocess(raw_imgs)
+        else:
+            if len(imgs.shape) == 3:  # (H, W, C)
+                original_shapes.append(imgs.shape)
+                preprocessed_imgs = self.preprocess([imgs])
+            else:
+                for idx in range(imgs.shape[0]):
+                    original_shapes.append(imgs[idx].shape)
+                preprocessed_imgs = self.preprocess(
+                    [imgs[idx] for idx in range(imgs.shape[0])]
+                )
+
+        _, _, h, w = preprocessed_imgs.shape
+        preprocessed_shape = (h, w)
+
+        outputs = self.session(preprocessed_imgs)
+
+        batch_dets = self.post_process(outputs, original_shapes, preprocessed_shape)
+
+        results = []
+        for i in range(len(batch_dets)):
+            dets = batch_dets[i]
+            if dets.shape[0] == 0:
+                results.append(
+                    [] if return_dict else np.empty((0, 15), dtype=np.float32)
+                )
+                continue
+
+            keep = nms(dets, self.nms_threshold)
+            dets = dets[keep, :]
+
+            dets = dets[: self.keep_top_k, :]
+
+            if return_dict:
+                dets = [parse_det(x) for x in dets]
+            results.append(dets)
+
+        return results

fast_face/models/base_recognition.py

@@ -0,0 +1,46 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+import numpy as np
+import onnxruntime as ort
+
+from ..schema import ProviderType
+from .session import ONNXSession
+
+
+class BaseRecognitionModel(ABC):
+    """Base class for face recognition models that extract embeddings."""
+
+    def __init__(
+        self,
+        model_path: str,
+        providers: list[ProviderType] | None = None,
+        sess_options: ort.SessionOptions | None = None,
+    ):
+        if providers is None:
+            providers = ["CPUExecutionProvider"]
+
+        self.session = ONNXSession(
+            model_path=model_path, providers=providers, sess_options=sess_options
+        )
+        self.input_name = self.session.session.get_inputs()[0].name
+
+    @abstractmethod
+    def preprocess(self, imgs: list[np.ndarray]) -> np.ndarray:
+        """Preprocess aligned face images for the specific model."""
+
+    @abstractmethod
+    def extract(
+        self,
+        imgs: list[np.ndarray],
+        landmarks: list[np.ndarray | dict[str, Any]] | None = None,
+    ) -> np.ndarray:
+        """Extract face embeddings from images.
+
+        Args:
+            imgs: List of face images (RGB).
+            landmarks: Optional list of landmarks for alignment.
+
+        Returns:
+            np.ndarray: L2-normalized embeddings of shape (N, embedding_dim).
+        """

fast_face/models/downloader.py

@@ -0,0 +1,46 @@
+import os
+import urllib.request
+import logging
+import shutil
+
+logger = logging.getLogger(__name__)
+
+BASE_URL = "https://mwmaulana.my.id/fast_face_onnx"
+
+
+def download_model(model_filename: str, target_path: str):
+    """Download the model from the remote server if it doesn't exist.
+
+    Args:
+        model_filename (str): The name of the file to download (e.g. 'yunet.onnx').
+        target_path (str): The local path where the file should be saved.
+    """
+    if os.path.exists(target_path):
+        return
+
+    url = f"{BASE_URL}/{model_filename}"
+
+    os.makedirs(os.path.dirname(target_path), exist_ok=True)
+
+    logger.info(f"Model not found locally. Downloading {model_filename} from {url}...")
+    print(f"Downloading {model_filename} from {url}...")
+
+    try:
+        req = urllib.request.Request(
+            url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}
+        )
+        with (
+            urllib.request.urlopen(req) as response,
+            open(target_path, "wb") as out_file,
+        ):
+            shutil.copyfileobj(response, out_file)
+
+        logger.info(f"Successfully downloaded to {target_path}")
+        print(f"Successfully downloaded to {target_path}")
+    except Exception as e:
+        logger.error(f"Failed to download model {model_filename}: {e}")
+        if os.path.exists(target_path):
+            os.remove(target_path)
+        raise RuntimeError(
+            f"Failed to download {model_filename} from {url}. Error: {e}"
+        ) from e

fast_face/models/factory.py

@@ -0,0 +1,69 @@
+import os
+from typing import Union
+
+from .base import BaseFaceModel
+from .base_recognition import BaseRecognitionModel
+from .yunet import YuNet
+from .retinaface import RetinaFace
+from .adaface import AdaFace
+from .downloader import download_model
+from ..schema import MODEL_FILENAMES
+
+
+class FaceModelFactory:
+    """Factory to instantiate face detection models."""
+
+    _models = {
+        "YUNET": YuNet,
+        "RETINAFACE_MOBILENET": RetinaFace,
+        "RETINAFACE_RESNET50": RetinaFace,
+        "ADAFACE_IR101": AdaFace,
+        "ADAFACE_IR50": AdaFace,
+        "ADAFACE_IR18": AdaFace,
+        # "SCRFD": SCRFD, # To be implemented
+    }
+
+    @classmethod
+    def get_model(
+        cls, model_type: str, **kwargs
+    ) -> Union[BaseFaceModel, BaseRecognitionModel]:
+        """Get an instance of a face detection or recognition model.
+
+        Args:
+            model_type (str): The type of model to instantiate (e.g., "YUNET", "ADAFACE_IR50").
+            **kwargs: Arguments to pass to the model's constructor.
+
+        Returns:
+            Union[BaseFaceModel, BaseRecognitionModel]: An instance of the requested model.
+        """
+        model_type = model_type.upper()
+        if model_type not in cls._models:
+            raise ValueError(
+                f"Model type '{model_type}' is not supported. Supported models: {list(cls._models.keys())}"
+            )
+
+        canonical_filename = MODEL_FILENAMES.get(
+            model_type, f"{model_type.lower()}.onnx"
+        )
+        if "model_path" not in kwargs:
+            kwargs["model_path"] = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), canonical_filename
+            )
+
+        model_path = kwargs["model_path"]
+        if not os.path.exists(model_path):
+            os.makedirs(os.path.dirname(model_path), exist_ok=True)
+            download_model(canonical_filename, model_path)
+
+        model_class = cls._models[model_type]
+        return model_class(**kwargs)
+
+    @classmethod
+    def register_model(cls, model_type: str, model_class: type[BaseFaceModel]):
+        """Register a custom model with the factory.
+
+        Args:
+            model_type (str): The identifier for the model.
+            model_class (Type[BaseFaceModel]): The model class.
+        """
+        cls._models[model_type.upper()] = model_class

fast_face/models/retinaface.py

@@ -0,0 +1,109 @@
+import cv2
+import numpy as np
+
+from ..schema import ProviderType
+from ..tools import decode, decode_landmark, get_priorbox
+from .base import BaseFaceModel
+
+
+class RetinaFace(BaseFaceModel):
+    """RetinaFace detector supporting MobileNet and ResNet50 backbones."""
+
+    def __init__(
+        self,
+        model_path: str,
+        input_size: int = 640,
+        conf_threshold: float = 0.6,
+        nms_threshold: float = 0.2,
+        top_k: int = 10000,
+        keep_top_k: int = 1000,
+        variance: tuple[float, float] = (0.1, 0.2),
+        providers: list[ProviderType] | None = None,
+        **kwargs,
+    ):
+        if providers is None:
+            providers = ["CPUExecutionProvider"]
+        super().__init__(
+            model_path=model_path,
+            conf_threshold=conf_threshold,
+            nms_threshold=nms_threshold,
+            top_k=top_k,
+            keep_top_k=keep_top_k,
+            providers=providers,
+            sess_options=kwargs.get("sess_options"),
+        )
+        self.input_size = input_size
+        self.variance = variance
+
+    def _resize_single(self, image: np.ndarray) -> np.ndarray:
+        """Letterbox-resize a single image to (input_size x input_size)."""
+        h, w = image.shape[:2]
+        scale = self.input_size / max(h, w)
+        new_w, new_h = int(w * scale), int(h * scale)
+        resized = cv2.resize(image, (new_w, new_h))
+        pad_w = self.input_size - new_w
+        pad_h = self.input_size - new_h
+        return cv2.copyMakeBorder(
+            resized, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=(0, 0, 0)
+        )
+
+    def preprocess(self, imgs: list[np.ndarray]) -> np.ndarray:
+        """Preprocess images: resize, normalize, transpose to NCHW."""
+        processed = []
+        for img in imgs:
+            resized = self._resize_single(img)
+            processed.append(resized)
+
+        batch = np.stack(processed, axis=0).astype(np.float32)
+
+        # RetinaFace normalization: subtract mean, no std division
+        mean = np.array([104.0, 117.0, 123.0], dtype=np.float32)
+        batch -= mean
+
+        # HWC -> CHW
+        batch = np.transpose(batch, (0, 3, 1, 2))
+        return batch
+
+    def post_process(
+        self,
+        outputs: list[np.ndarray],
+        original_shapes: list[tuple[int, int]],
+        preprocessed_shape: tuple[int, int],
+    ) -> list[np.ndarray]:
+        """Decode RetinaFace outputs into [x1, y1, x2, y2, score, landmarks...] arrays."""
+        loc, conf, landms = outputs
+        h, w = preprocessed_shape
+        prior_data = get_priorbox(image_size=(h, w))
+
+        results = []
+        for i in range(loc.shape[0]):
+            orig_h, orig_w = original_shapes[i][:2]
+            padded_size = max(orig_h, orig_w)
+            scale_box = np.array([padded_size] * 4, dtype=np.float32)
+            scale_lm = np.array([padded_size] * 10, dtype=np.float32)
+
+            boxes = decode(loc[i], prior_data, self.variance) * scale_box
+            landmarks = decode_landmark(landms[i], prior_data, self.variance) * scale_lm
+            scores = conf[i][:, 1]
+
+            inds = np.where(scores > self.conf_threshold)[0]
+            boxes = boxes[inds]
+            landmarks = landmarks[inds]
+            scores = scores[inds]
+
+            order = scores.argsort()[::-1][: self.top_k]
+            boxes = boxes[order]
+            landmarks = landmarks[order]
+            scores = scores[order]
+
+            if boxes.shape[0] == 0:
+                results.append(np.empty((0, 15), dtype=np.float32))
+                continue
+
+            dets = np.hstack((boxes, scores[:, np.newaxis], landmarks)).astype(
+                np.float32, copy=False
+            )
+
+            results.append(dets)
+
+        return results

fast_face/models/session.py

@@ -0,0 +1,29 @@
+import numpy as np
+import onnxruntime as ort
+
+
+from ..schema import ProviderType
+
+
+class ONNXSession:
+    def __init__(
+        self,
+        model_path: str,
+        providers: list[ProviderType] | None = None,
+        sess_options: ort.SessionOptions | None = None,
+    ):
+
+        if providers is None:
+            providers = ["CPUExecutionProvider"]
+        self.session = ort.InferenceSession(
+            model_path, providers=providers, sess_options=sess_options
+        )
+
+    def process(self, inputs: np.ndarray):
+        assert len(inputs.shape) == 4, "Inputs shape length != 4"
+        input_name = self.session.get_inputs()[0].name
+        outputs = self.session.run(None, {input_name: inputs})
+        return outputs
+
+    def __call__(self, inputs: np.ndarray):
+        return self.process(inputs=inputs)