feat: Add face recognition, rename and modify some files

uniface/__init__.py

@@ -19,7 +19,7 @@ __version__ = "0.1.8"
 from uniface.retinaface import RetinaFace
 from uniface.log import Logger
 from uniface.model_store import verify_model_weights
-from uniface.alignment import face_alignment
+from uniface.face_utils import face_alignment, compute_similarity
 from uniface.visualization import draw_detections
 
 __all__ = [
@@ -30,5 +30,6 @@ __all__ = [
     "Logger",
     "verify_model_weights",
     "draw_detections",
-    "face_alignment"
+    "face_alignment",
+    "compute_similarity",
 ]

uniface/constants.py

@@ -5,14 +5,25 @@
 from enum import Enum
 from typing import Dict
 
+# fmt: off
+class FaceEncoderWeights(str, Enum):
+    SPHERE20      = "sphere20"
+    SPHERE36      = "sphere36"
+    MNET_025      = "mobilenetv1_025"
+    MNET_V2       = "mobilenetv2"
+    MNET_V3_SMALL = "mobilenetv3_small"
+    MNET_V3_LARGE = "mobilenetv3_large"
+
 
 class RetinaFaceWeights(str, Enum):
-    MNET_025 = "retinaface_mnet025"
-    MNET_050 = "retinaface_mnet050"
-    MNET_V1  = "retinaface_mnet_v1"
-    MNET_V2  = "retinaface_mnet_v2"
-    RESNET18 = "retinaface_r18"
-    RESNET34 = "retinaface_r34"
+    MNET_025 =  "retinaface_mnet025"
+    MNET_050 =  "retinaface_mnet050"
+    MNET_V1  =  "retinaface_mnet_v1"
+    MNET_V2  =  "retinaface_mnet_v2"
+    RESNET18 =  "retinaface_r18"
+    RESNET34 =  "retinaface_r34"
+
+# fmt: on
 
 
 MODEL_URLS: Dict[RetinaFaceWeights, str] = {

uniface/face_utils.py

@@ -80,3 +80,19 @@ def face_alignment(image: np.ndarray, landmark: np.ndarray, image_size: int = 11
     warped = cv2.warpAffine(image, M, (image_size, image_size), borderValue=0.0)
 
     return warped, M_inv
+
+
+def compute_similarity(feat1: np.ndarray, feat2: np.ndarray) -> np.float32:
+    """Computing Similarity between two faces.
+
+    Args:
+        feat1 (np.ndarray): Face features.
+        feat2 (np.ndarray): Face features.
+
+    Returns:
+        np.float32: Cosine similarity between face features.
+    """
+    feat1 = feat1.ravel()
+    feat2 = feat2.ravel()
+    similarity = np.dot(feat1, feat2) / (np.linalg.norm(feat1) * np.linalg.norm(feat2))
+    return similarity

uniface/model_store.py

@@ -10,6 +10,9 @@ from uniface.log import Logger
 import uniface.constants as const
 
 
+__all__ = ['verify_model_weights']
+
+
 def verify_model_weights(model_name: str, root: str = '~/.uniface/models') -> str:
     """
     Ensures model weights are available by downloading if missing and verifying integrity with a SHA-256 hash.

uniface/recognition/encoder.py

@@ -0,0 +1,114 @@
+# Copyright 2025 Yakhyokhuja Valikhujaev
+# Author: Yakhyokhuja Valikhujaev
+# GitHub: https://github.com/yakhyo
+
+import os
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+from typing import Tuple, List, Optional, Literal
+
+from uniface.face_utils import compute_similarity, face_alignment
+from uniface.model_store import verify_model_weights
+from uniface.constants import FaceEncoderWeights
+from uniface.logger import Logger
+
+
+class FaceEncoder:
+    """
+    Face recognition model using ONNX Runtime for inference and OpenCV for image preprocessing,
+    utilizing an external face alignment function.
+    """
+
+    def __init__(
+        self,
+        model_path: Optional[FaceEncoderWeights] = FaceEncoderWeights.MNET_V2,
+    ) -> None:
+        """
+        Initializes the FaceEncoder model for inference.
+
+        Args:
+            model_path (str): Path to the ONNX model file.
+        """
+        self.input_mean = 127.5
+        self.input_std = 127.5
+
+        # Get path to model weights
+        self._model_path = verify_model_weights(model_path)
+        Logger.info(f"Verfied model weights located at: {self._model_path}")
+
+        self._initialize_model(self._model_path)
+
+    def _initialize_model(self, model_path: str) -> None:
+        """
+        Loads the ONNX model and prepares it for inference.
+
+        Args:
+            model_path (str): Path to the ONNX model file.
+
+        Raises:
+            RuntimeError: If the model fails to load or initialize.
+        """
+        try:
+            self.session = ort.InferenceSession(
+                model_path,
+                providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
+            )
+            self._setup_model()
+            Logger.info(f"Successfully initialized face encoder from {model_path}")
+        except Exception as e:
+            Logger.error(f"Failed to load face encoder model from '{model_path}'", exc_info=True)
+            raise RuntimeError(f"Failed to initialize model session for '{model_path}'") from e
+
+    def _setup_model(self) -> None:
+        """
+        Extracts input/output configuration from the ONNX model session.
+        """
+        input_cfg = self.session.get_inputs()[0]
+        input_shape = input_cfg.shape
+
+        self.input_name = input_cfg.name
+        self.input_size = tuple(input_shape[2:4][::-1])  # (width, height)
+
+        outputs = self.session.get_outputs()
+        self.output_names = [output.name for output in outputs]
+
+        assert len(self.output_names) == 1, "Expected only one output node."
+        self.output_shape = outputs[0].shape
+
+    def preprocess(self, image: np.ndarray) -> np.ndarray:
+        """
+        Preprocess the image: resize, normalize, and convert it to a blob.
+
+        Args:
+            image (np.ndarray): Input image in BGR format.
+
+        Returns:
+            np.ndarray: Preprocessed image as a NumPy array ready for inference.
+        """
+        image = cv2.resize(image, self.input_size)  # Resize to (112, 112)
+        blob = cv2.dnn.blobFromImage(
+            image,
+            scalefactor=1.0 / self.input_std,
+            size=self.input_size,
+            mean=(self.input_mean, self.input_mean, self.input_mean),
+            swapRB=True  # Convert BGR to RGB
+        )
+        return blob
+
+    def get_embedding(self, image: np.ndarray, landmarks: np.ndarray) -> np.ndarray:
+        """
+        Extracts face embedding from an aligned image.
+
+        Args:
+            image (np.ndarray): Input face image (BGR format).
+            landmarks (np.ndarray): Facial landmarks (5 points for alignment).
+
+        Returns:
+            np.ndarray: 512-dimensional face embedding.
+        """
+        aligned_face = face_alignment(image, landmarks)  # Use your function for alignment
+        blob = self.preprocess(image)  # Convert to blob
+        embedding = self.session.run(self.output_names, {self.input_name: blob})[0]
+        return embedding  # Return the 512-D feature vector

uniface/retinaface.py

@@ -96,11 +96,14 @@ class RetinaFace:
             RuntimeError: If the model fails to load, logs an error and raises an exception.
         """
         try:
-            self.session = ort.InferenceSession(model_path)
+            self.session = ort.InferenceSession(
+                model_path,
+                providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
+            )
             self.input_name = self.session.get_inputs()[0].name
             Logger.info(f"Successfully initialized the model from {model_path}")
         except Exception as e:
-            Logger.error(f"Failed to load model from '{model_path}': {e}")
+            Logger.error(f"Failed to load model from '{model_path}': {e}", exc_info=True)
             raise RuntimeError(f"Failed to initialize model session for '{model_path}'") from e
 
     def preprocess(self, image: np.ndarray) -> np.ndarray: