"""
This module provides functionality for gender and age detection in facial images using ONNX models.
It includes functions for loading attribute detection models and performing gender/age inference on faces.
The module integrates with Hugging Face's model hub and provides caching capabilities for model loading.
"""
from typing import Union, Tuple
import numpy as np
from huggingface_hub import hf_hub_download
from imgutils.data import ImageTyping, load_image
from imgutils.utils import ts_lru_cache, open_onnx_model
from .base import _REPO_ID, _DEFAULT_MODEL, _affine_transform, Face
@ts_lru_cache()
def _open_attribute_model(model_name: str):
"""
Load and cache an ONNX model for facial attribute detection.
:param model_name: Name of the model to load
:type model_name: str
:return: A tuple containing:
- ONNX session object
- Input tensor name
- Input shape tuple (height, width)
- List of output tensor names
:rtype: tuple[onnxruntime.InferenceSession, str, tuple[int, int], list[str]]
"""
session = open_onnx_model(hf_hub_download(
repo_id=_REPO_ID,
repo_type='model',
filename=f'{model_name}/genderage.onnx'
))
input_cfg = session.get_inputs()[0]
input_name = input_cfg.name
input_shape = tuple(input_cfg.shape[2:4][::-1])
output_names = [o.name for o in session.get_outputs()]
return session, input_name, input_shape, output_names
[docs]def isf_genderage(image: ImageTyping, face: Union[Face, Tuple[float, float, float, float]],
model_name: str = _DEFAULT_MODEL, no_write: bool = False):
"""
Detect gender and age from a facial image.
This function performs gender and age detection on a given face in an image. It can work with
either a Face object or raw bounding box coordinates.
:param image: Input image (can be path, URL, PIL Image, or numpy array)
:type image: ImageTyping
:param face: Face object or bounding box coordinates (x1, y1, x2, y2)
:type face: Union[Face, Tuple[float, float, float, float]]
:param model_name: Name of the model to use for detection, defaults to _DEFAULT_MODEL
:type model_name: str
:param no_write: If True, don't update the Face object with results (if face is Face object)
:type no_write: bool
:return: Tuple of (gender, age) where gender is 'F' or 'M' and age is an integer
:rtype: tuple[str, int]
"""
if isinstance(face, Face):
bbox = face.bbox
else:
bbox = face
image = load_image(image, force_background='white', mode='RGB')
w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
center = ((bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2)
session, input_name, input_shape, output_names = _open_attribute_model(model_name=model_name)
assert input_shape[0] == input_shape[1], f'Input shape is not a square - {input_shape!r}.'
input_size = input_shape[0]
scale = input_size / (max(w, h) * 1.5)
aimg, _ = _affine_transform(np.array(image), center, input_size, scale, 0)
blob = aimg.astype(np.float32)
blob = blob.transpose(2, 0, 1)[np.newaxis] # NCHW format
pred = session.run(None, {session.get_inputs()[0].name: blob})[0][0]
gender = ['F', 'M'][np.argmax(pred[:2]).item()]
age = int(np.round(pred[2] * 100))
if isinstance(face, Face) and not no_write:
face.age = age
face.gender = gender
return gender, age