UNetDFL

class plugins.extract.mask.unet_dfl.UNetDFL

Bases: FacePlugin

Neural network to process face image into a segmentation mask of the face

Attributes Summary

device

The selected device to run torch ops on

Methods Summary

from_torch(batch)

Run inference on a PyTorch model.

load_model()

Initialize the UNet-DFL Model

load_torch_model(model, weights_path[, ...])

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

post_process(batch)

Override to perform post-processing

pre_process(batch)

Override to perform pre-processing

process(batch)

Get the masks from the model

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() → UnetDFL

Initialize the UNet-DFL Model

Return type:

The loaded UnetDFL 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[tuple[Any, ...], dtype[float32]]

Override to perform post-processing

Parameters:

batch (ndarray) – This will be the output from the previous ‘process’ step

Returns:

  • For detect plugins this must be an (N, M, left, top, right, bottom) bounding boxes for

  • detected faces scaled to model input size as float32. N is the batch size, M is the number

  • of detections per batch

  • For align plugins this must be an (N, 68, 2) float32 array for each (x, y) landmark point

  • for each face in the batch. co-ordinates should be normalized to 0.0 to 1.0 range

  • For mask plugins this must be an (N, size, size) float32 image in range 0. - 1.0 for each

  • face in the batch

  • For identity plugins this must be an (N, M) float32 identity embedding

Return type:

ndarray[tuple[Any, …], dtype[float32]]

pre_process(batch: ndarray) → ndarray

Override to perform pre-processing

Parameters:

batch (ndarray) –

For detection plugins, this will be a batch of square, padded, images at model input size in the plugin’s color order, image format and data range.

For align plugins this will be a face detection ROI bounding box (batch size, left, top, right, bottom) as INT32.

For all other plugins this will be a batch of aligned face images at model input size in the plugin’s color order, image format and data range

Returns:

  • For align plugins, this should be an adjustment of the detected face’s bounding box to cut

  • a square out of the original image for feeding the model. Out of bounds values are allowed,

  • as these will be handled. This bounding box will be used to prepare the image at the

  • correct size for feeding the model.

  • For all other plugins, any pre-processing (eg normalization) should be applied ready for

  • feeding the model.

Return type:

ndarray

process(batch: ndarray) → ndarray

Get the masks from the model

Parameters:

batch (ndarray) – The batch to feed into the masker

Return type:

The predicted masks from the plugin

load_model() → UnetDFL

Initialize the UNet-DFL Model

Return type:

The loaded UnetDFL model

model: UnetDFL

The loaded model for the plugin

process(batch: ndarray) → ndarray

Get the masks from the model

Parameters:

batch (ndarray) – The batch to feed into the masker

Return type:

The predicted masks from the plugin