MTCNN

class plugins.extract.detect.mtcnn.MTCNN

Bases: ExtractPlugin

MTCNN detector for face recognition.

Attributes Summary

device

The selected device to run torch ops on

Methods Summary

`from_torch`(batch)	Run inference on a PyTorch model.
`load_model`()	Load the model
`load_torch_model`(model, weights_path[, ...])	Load a PyTorch model, apply the weights and pass a warmup batch through
`post_process`(batch)	Remove confidences from output
`pre_process`(batch)	Compile the detection image(s) for prediction.
`process`(batch)	Run model to get predictions

Attributes Documentation

device: The selected device to run torch ops on

Methods Documentation

from_torch(batch: ndarray) → ndarray

Run inference on a PyTorch model.

This function does not need to be used, however it handles torch backend for better throughput, so it is recommended. Must have used self.load_torch_model to load the Torch model to use this function.

Parameters:: batch (ndarray) – The batch array to feed to the PyTorch model
Return type:: The result from the PyTorch model

load_model() → MTCNNModel

Load the model

Return type:: The loaded MTCNN model

load_torch_model(model: Module, weights_path: str, return_indices: list[int] | None = None) → Module

Load a PyTorch model, apply the weights and pass a warmup batch through

This function does not need to be used, but some default Faceswap optimizations are performed here, so without using this function you will either need to apply them yourself or not have them applied

Parameters:

model (Module) – The Torch model to load
weights_path (str) – Full path to the weights file to load
return_indices (list[int] | None) – If the model outputs multiple items, but you only require some of them, the indices of the required items can be placed here so that when calling from_torch any extra data is not copied from the GPU. Default: None (return all data)

Return type:

The loaded model ready for inference

post_process(batch: ndarray) → ndarray

Remove confidences from output

Parameters:: batch (ndarray) – The detection results for the model
Return type:: The processed detection bounding box from the model at model input size

pre_process(batch: ndarray) → ndarray

Compile the detection image(s) for prediction. No further pre-processing required for MTCNN

Parameters:: batch (ndarray) – The input batch of images at model input size in the correct color order
Return type:: The batch of images ready for feeding the model

process(batch: ndarray) → ndarray

Run model to get predictions

Parameters:: batch (ndarray) – A batch of images ready to feed the model
Return type:: The batch of detection results from the model