medusa.landmark.retinaface#

A landmark detection 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.landmark.retinaface.RetinafaceLandmarkModel(model_name='2d106det', model_path=None, device=DEVICE)[source]#

Landmark detection model based on Insightface’s Retinaface model.

Parameters:

  • model_name (str) – Name of the landmark model from Insightface that should be used; options are ‘2d106det’ (106 landmarks, 2D) or ‘1k3d68’ (68 landmarks, 3D)

  • detector (BaseDetector) – Which detector to use; options are SCRFDetector or YunetDetector

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

forward(imgs)[source]#

Runs the landmark 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:

out_lms – Dictionary with the following keys: ‘lms’ (landmarks), ‘conf’ (confidence), ‘img_idx’ (image index), ‘bbox’ (bounding box)

Return type:

dict