ExtractPlugin

class plugins.extract.base.ExtractPlugin(input_size: int, batch_size: int = 1, is_rgb: bool = False, dtype: str = 'float32', scale: tuple[int, int] = (0, 1), force_cpu: bool = False)

Bases: ABC

Base extract plugin that all plugins must inherit from.

Parameters:

  • input_size (int) – The size of the input required by the plugin. The input will always be square at these dimensions

  • batch_size (int) – The batch size that the plugin processes data at. Note: Only the process method is guaranteed to receive data at this batch size (or less). The other processes may receive higher batch sizes for re-processing reasons. Do not rely on this when processing data. Default: 1

  • is_rgb (bool) – True if the plugin expects input images to be RGB rather than BGR. Default: False

  • dtype (str) – A valid datatype that the plugin expects to receive the image at. Default: “float32”

  • scale (tuple[int, int]) – The scale that the plugin expects to receive the image at eg: (0, 255) for uint8 images. Default: (0, 1)

  • force_cpu (bool) – For Torch models, force running on the CPU, rather than the accelerated device. Sets the torch.device to device. Default: False

Attributes Summary

device

The selected device to run torch ops on

Methods Summary

from_torch(batch)

Run inference on a PyTorch model.

load_model()

Override to perform any model initialization code

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)

Override to perform processing.

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

abstractmethod load_model() → torch.nn.Module | cv2.dnn.Net | T.Any

Override to perform any model initialization code

Return type:

The loaded model that will be accessible from 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

abstractmethod process(batch: ndarray) → ndarray

Override to perform processing. This is where the model should be called

Parameters:

batch (ndarray) –

For detection plugins, this will be a batch of square, padded, images at model input size in the correct format for feeding the model

For align, mask and identity plugins this will be a batch of square face patches at model input size in the correct format for feeding the model

Returns:

  • This can return any numpy array, but it must be a numpy array. For detect plugins that can

  • return several results, usually in a list, then this must be an object array

Return type:

ndarray

batch_size

The maximum batch size that this plugin’s ‘process’ method will receive

property device: device

The selected device to run torch ops on

dtype

The datatype that the plugin expects images at

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

input_size

The size of the plugin’s input in pixels

is_rgb

True if the plugin expects RGB images. False for BGR

abstractmethod load_model() → torch.nn.Module | cv2.dnn.Net | T.Any

Override to perform any model initialization code

Return type:

The loaded model that will be accessible from 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

model: torch.nn.Module | cv2.dnn.Net | T.Any

The loaded model for the plugin

name

The name of the plugin. Derived from the module name

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

abstractmethod process(batch: ndarray) → ndarray

Override to perform processing. This is where the model should be called

Parameters:

batch (ndarray) –

For detection plugins, this will be a batch of square, padded, images at model input size in the correct format for feeding the model

For align, mask and identity plugins this will be a batch of square face patches at model input size in the correct format for feeding the model

Returns:

  • This can return any numpy array, but it must be a numpy array. For detect plugins that can

  • return several results, usually in a list, then this must be an object array

Return type:

ndarray

scale

The numeric range that the plugin expects images to be in