medusa.recognize.retinaface

A face recognition model based on Insightface’s Retinaface model, but implemented in PyTorch (but parts of it run with ONNX model), so can be fully run on GPU (no numpy necessary).

Module Contents

class medusa.recognize.retinaface.RetinafaceRecognitionModel(model_path=None, device=DEVICE)

Face recognition model based on Insightface’s Retinaface model (trained using partial FC).

Parameters:

• model_path (str, Path) – Path to the ONNX model file; if None, will use the default model

• device (str) – Either ‘cuda’ (GPU) or ‘cpu’

forward(imgs)

Runs the recognition model on a set of images.

Parameters:

imgs (list, str, Path, torch.tensor) – Either a list of images, a path to a directory containing images, or an already loaded torch.tensor of shape N (batch) x C x H x W

Returns:

X_emb – Face embeddings of shape N x 512 (where N is the number of detections, not necessarily the number of input images)

Return type:

torch.tensor

class medusa.recognize.retinaface.RetinafaceGenderAgeModel(model_path=None, device=DEVICE)

Gender and age prediction model based on Insightface’s model.

Parameters:

• model_path (str, Path) – Path to the ONNX model file; if None, will use the default model

• device (str) – Either ‘cuda’ (GPU) or ‘cpu’

forward(imgs)

Runs the gender+age prediction model on a set of images.

Parameters:

imgs (torch.tensor) – A torch.tensor of shape N (batch) x C x H x W

Returns:

attr – Dictionary with the following keys: ‘gender’ (0 = female, 1 = male), ‘age’ (age in years), both torch.tensors with the first dimension being the number of faces detected in the input images

Return type:

dict