From 85cf413cb81eb8e253f439d3da89ce3516c125ec Mon Sep 17 00:00:00 2001 From: yakhyo Date: Wed, 2 Apr 2025 16:31:54 +0900 Subject: [PATCH] feat: Add face attribute detection model --- uniface/attribute/__init__.py | 0 uniface/attribute/models.py | 100 ++++++++++++++++++++++++++++++++++ uniface/face_utils.py | 45 +++++++++++++++ uniface/recognition/base.py | 2 +- 4 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 uniface/attribute/__init__.py create mode 100644 uniface/attribute/models.py diff --git a/uniface/attribute/__init__.py b/uniface/attribute/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/uniface/attribute/models.py b/uniface/attribute/models.py new file mode 100644 index 0000000..c2379b4 --- /dev/null +++ b/uniface/attribute/models.py @@ -0,0 +1,100 @@ +import cv2 +import numpy as np +import onnxruntime +from typing import Tuple + +from uniface.face_utils import bbox_center_alignment + +__all__ = ["Attribute"] + + +class Attribute: + """ + Age and Gender Prediction Model. + """ + def __init__(self, model_path: str) -> None: + """ + Initializes the Attribute model for inference. + + Args: + model_path (str): Path to the ONNX file. + """ + self.model_path = model_path + + self.input_std = 1.0 + self.input_mean = 0.0 + + self._initialize_model(model_path=model_path) + + def _initialize_model(self, model_path: str): + """Initialize the model from the given path. + + Args: + model_path (str): Path to .onnx model. + """ + try: + self.session = onnxruntime.InferenceSession( + model_path, + providers=["CUDAExecutionProvider", "CPUExecutionProvider"] + + ) + + # Get model info + metadata = self.session.get_inputs()[0] + input_shape = metadata.shape + self.input_size = tuple(input_shape[2:4][::-1]) + + self.input_names = [x.name for x in self.session.get_inputs()] + self.output_names = [x.name for x in self.session.get_outputs()] + + except Exception as e: + print(f"Failed to load the model: {e}") + raise + + def preprocess(self, image: np.ndarray, bbox: np.ndarray): + """Preprocessing + + Args: + image (np.ndarray): Numpy image + bbox (np.ndarray): Bounding box coordinates: [x1, y1, x2, y2] + + Returns: + np.ndarray: Transformed image + """ + width, height = bbox[2] - bbox[0], bbox[3] - bbox[1] + center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2 + scale = self.input_size[0] / (max(width, height)*1.5) + rotation = 0.0 + + transformed_image, M = bbox_center_alignment(image, center, self.input_size[0], scale, rotation) + + input_size = tuple(transformed_image.shape[0:2][::-1]) + + blob = cv2.dnn.blobFromImage( + transformed_image, + 1.0/self.input_std, + input_size, + (self.input_mean, self.input_mean, self.input_mean), + swapRB=True + ) + return blob + + def postprocess(self, predictions: np.ndarray) -> Tuple[np.int64, int]: + """Postprocessing + + Args: + predictions (np.ndarray): Model predictions, shape: [1, 3] + + Returns: + Tuple[np.int64, int]: Gender and Age values + """ + gender = np.argmax(predictions[:2]) + age = int(np.round(predictions[2]*100)) + return gender, age + + def get(self, image: np.ndarray, bbox: np.ndarray) -> Tuple[np.int64, int]: + blob = self.preprocess(image, bbox) + predictions = self.session.run(self.output_names, {self.input_names[0]: blob})[0][0] + gender, age = self.postprocess(predictions) + + return gender, age diff --git a/uniface/face_utils.py b/uniface/face_utils.py index 100ae36..a3229b4 100644 --- a/uniface/face_utils.py +++ b/uniface/face_utils.py @@ -96,3 +96,48 @@ def compute_similarity(feat1: np.ndarray, feat2: np.ndarray) -> np.float32: feat2 = feat2.ravel() similarity = np.dot(feat1, feat2) / (np.linalg.norm(feat1) * np.linalg.norm(feat2)) return similarity + + +def bbox_center_alignment(image, center, output_size, scale, rotation): + """ + Apply center-based alignment, scaling, and rotation to an image. + + Args: + image (np.ndarray): Input image. + center (Tuple[float, float]): Center point (e.g., face center from bbox). + output_size (int): Desired output image size (square). + scale (float): Scaling factor to zoom in/out. + rotation (float): Rotation angle in degrees (clockwise). + + Returns: + cropped (np.ndarray): Aligned and cropped image. + M (np.ndarray): 2x3 affine transform matrix used. + """ + + # Convert rotation from degrees to radians + rot = float(rotation) * np.pi / 180.0 + + # Scale the image + t1 = trans.SimilarityTransform(scale=scale) + + # Translate the center point to the origin (after scaling) + cx = center[0] * scale + cy = center[1] * scale + t2 = trans.SimilarityTransform(translation=(-1 * cx, -1 * cy)) + + # Apply rotation around origin (center of face) + t3 = trans.SimilarityTransform(rotation=rot) + + # Translate origin to center of output image + t4 = trans.SimilarityTransform(translation=(output_size / 2, output_size / 2)) + + # Combine all transformations in order: scale → center shift → rotate → recentralize + t = t1 + t2 + t3 + t4 + + # Extract 2x3 affine matrix + M = t.params[0:2] + + # Warp the image using OpenCV + cropped = cv2.warpAffine(image, M, (output_size, output_size), borderValue=0.0) + + return cropped, M diff --git a/uniface/recognition/base.py b/uniface/recognition/base.py index 8371e9b..5be708e 100644 --- a/uniface/recognition/base.py +++ b/uniface/recognition/base.py @@ -32,7 +32,7 @@ class BaseFaceEncoder: Initializes the FaceEncoder model for inference. Args: - model_path (str): Path to the ONNX model file. + model_path (Optional[SphereFaceWeights | MobileFaceWeights]): Path to the ONNX model file. """ self.input_mean = 127.5 self.input_std = 127.5