mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-09 04:36:50 -04:00
490 lines
22 KiB
Python
490 lines
22 KiB
Python
#!/usr/bin/env python3
|
|
""" Neural Network Blocks for faceswap.py. """
|
|
|
|
import logging
|
|
|
|
from keras.layers import Add, Concatenate, SeparableConv2D, UpSampling2D
|
|
from keras.layers.advanced_activations import LeakyReLU
|
|
from keras.layers.convolutional import Conv2D
|
|
from keras.layers.core import Activation
|
|
from keras.initializers import he_uniform, VarianceScaling
|
|
from .initializers import ICNR, ConvolutionAware
|
|
from .layers import PixelShuffler, ReflectionPadding2D
|
|
from .normalization import InstanceNormalization
|
|
|
|
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
class NNBlocks():
|
|
""" Blocks that are often used for multiple models are stored here for easy access.
|
|
|
|
This class is always brought in as ``self.blocks`` in all model plugins so that all models
|
|
have access to them.
|
|
|
|
The parameters passed into this class should ultimately originate from the user's training
|
|
configuration file, rather than being hard-coded at the plugin level.
|
|
|
|
Parameters
|
|
----------
|
|
use_icnr_init: bool, Optional
|
|
``True`` if ICNR initialization should be used rather than the default. Default: ``False``
|
|
use_convaware_init: bool, Optional
|
|
``True`` if Convolutional Aware initialization should be used rather than the default.
|
|
Default: ``False``
|
|
use_reflect_padding: bool, Optional
|
|
``True`` if Reflect Padding initialization should be used rather than the padding.
|
|
Default: ``False``
|
|
first_run: bool, Optional
|
|
``True`` if a model is being created for the first time, ``False`` if a model is being
|
|
resumed. Used to prevent Convolutional Aware weights from being calculated when a model
|
|
is being reloaded. Default: ``True``
|
|
"""
|
|
def __init__(self,
|
|
use_icnr_init=False,
|
|
use_convaware_init=False,
|
|
use_reflect_padding=False,
|
|
first_run=True):
|
|
logger.debug("Initializing %s: (use_icnr_init: %s, use_convaware_init: %s, "
|
|
"use_reflect_padding: %s, first_run: %s)",
|
|
self.__class__.__name__, use_icnr_init, use_convaware_init,
|
|
use_reflect_padding, first_run)
|
|
self.names = dict()
|
|
self.first_run = first_run
|
|
self.use_icnr_init = use_icnr_init
|
|
self.use_convaware_init = use_convaware_init
|
|
self.use_reflect_padding = use_reflect_padding
|
|
if self.use_convaware_init and self.first_run:
|
|
logger.info("Using Convolutional Aware Initialization. Model generation will take a "
|
|
"few minutes...")
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
def _get_name(self, name):
|
|
""" Return unique layer name for requested block.
|
|
|
|
As blocks can be used multiple times, auto appends an integer to the end of the requested
|
|
name to keep all block names unique
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
The requested name for the layer
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
The unique name for this layer
|
|
"""
|
|
self.names[name] = self.names.setdefault(name, -1) + 1
|
|
name = "{}_{}".format(name, self.names[name])
|
|
logger.debug("Generating block name: %s", name)
|
|
return name
|
|
|
|
def _set_default_initializer(self, kwargs):
|
|
""" Sets the default initializer for convolution 2D and Seperable convolution 2D layers
|
|
to Convolutional Aware or he_uniform.
|
|
|
|
if a specific initializer has been passed in from the model plugin, then the specified
|
|
initializer will be used rather than the default.
|
|
|
|
Parameters
|
|
----------
|
|
kwargs: dict
|
|
The keyword arguments for the current layer
|
|
|
|
Returns
|
|
-------
|
|
dict
|
|
The keyword arguments for the current layer with the initializer updated to
|
|
the select default value
|
|
"""
|
|
if "kernel_initializer" in kwargs:
|
|
logger.debug("Using model specified initializer: %s", kwargs["kernel_initializer"])
|
|
return kwargs
|
|
if self.use_convaware_init:
|
|
default = ConvolutionAware()
|
|
if self.first_run:
|
|
# Indicate the Convolutional Aware should be calculated on first run
|
|
default._init = True # pylint:disable=protected-access
|
|
else:
|
|
default = he_uniform()
|
|
if kwargs.get("kernel_initializer", None) != default:
|
|
kwargs["kernel_initializer"] = default
|
|
logger.debug("Set default kernel_initializer to: %s", kwargs["kernel_initializer"])
|
|
return kwargs
|
|
|
|
@staticmethod
|
|
def _switch_kernel_initializer(kwargs, initializer):
|
|
""" Switch the initializer in the given kwargs to the given initializer and return the
|
|
previous initializer to caller.
|
|
|
|
For residual blocks and up-scaling, user selected initializer methods should replace those
|
|
set by the model. This method updates the initializer for the layer, and returns the
|
|
original initializer so that it can be set back to the layer's key word arguments for
|
|
subsequent layers where the initializer should not be switched.
|
|
|
|
Parameters
|
|
----------
|
|
kwargs: dict
|
|
The keyword arguments for the current layer
|
|
initializer: keras or faceswap initializer class
|
|
The initializer that should replace the current initializer that exists in keyword
|
|
arguments
|
|
|
|
Returns
|
|
-------
|
|
keras or faceswap initializer class
|
|
The original initializer that existed in the given keyword arguments
|
|
"""
|
|
original = kwargs.get("kernel_initializer", None)
|
|
kwargs["kernel_initializer"] = initializer
|
|
logger.debug("Switched kernel_initializer from %s to %s", original, initializer)
|
|
return original
|
|
|
|
def conv2d(self, input_tensor, filters, kernel_size, strides=(1, 1), padding="same", **kwargs):
|
|
""" A standard Convolution 2D layer with correct initialization.
|
|
|
|
This layer creates a convolution kernel that is convolved with the layer input to produce
|
|
a tensor of outputs.
|
|
|
|
Parameters
|
|
----------
|
|
input_tensor: tensor
|
|
The input tensor to the layer
|
|
filters: int
|
|
The dimensionality of the output space (i.e. the number of output filters in the
|
|
convolution)
|
|
kernel_size: int
|
|
An integer or tuple/list of 2 integers, specifying the height and width of the 2D
|
|
convolution window. Can be a single integer to specify the same value for all spatial
|
|
dimensions
|
|
strides: tuple, optional
|
|
An integer or tuple/list of 2 integers, specifying the strides of the convolution along
|
|
the height and width. Can be a single integer to specify the same value for all spatial
|
|
dimensions. Default: `(1, 1)`
|
|
padding: ["valid", "same"], optional
|
|
The padding to use. Default: `"same"`
|
|
kwargs: dict
|
|
Any additional Keras standard layer keyword arguments
|
|
|
|
Returns
|
|
-------
|
|
tensor
|
|
The output tensor from the Convolution 2D Layer
|
|
"""
|
|
logger.debug("input_tensor: %s, filters: %s, kernel_size: %s, strides: %s, padding: %s, "
|
|
"kwargs: %s)", input_tensor, filters, kernel_size, strides, padding, kwargs)
|
|
if kwargs.get("name", None) is None:
|
|
kwargs["name"] = self._get_name("conv2d_{}".format(input_tensor.shape[1]))
|
|
kwargs = self._set_default_initializer(kwargs)
|
|
var_x = Conv2D(filters, kernel_size,
|
|
strides=strides,
|
|
padding=padding,
|
|
**kwargs)(input_tensor)
|
|
return var_x
|
|
|
|
# <<< Original Model Blocks >>> #
|
|
def conv(self, input_tensor, filters, kernel_size=5, strides=2, padding="same",
|
|
use_instance_norm=False, res_block_follows=False, **kwargs):
|
|
""" A standard Convolution 2D layer which applies user specified configuration to the
|
|
layer.
|
|
|
|
Adds reflection padding if it has been selected by the user, and other post-processing
|
|
if requested by the plugin.
|
|
|
|
Parameters
|
|
----------
|
|
input_tensor: tensor
|
|
The input tensor to the layer
|
|
filters: int
|
|
The dimensionality of the output space (i.e. the number of output filters in the
|
|
convolution)
|
|
kernel_size: int, optional
|
|
An integer or tuple/list of 2 integers, specifying the height and width of the 2D
|
|
convolution window. Can be a single integer to specify the same value for all spatial
|
|
dimensions. Default: 5
|
|
strides: tuple or int, optional
|
|
An integer or tuple/list of 2 integers, specifying the strides of the convolution along
|
|
the height and width. Can be a single integer to specify the same value for all spatial
|
|
dimensions. Default: `2`
|
|
padding: ["valid", "same"], optional
|
|
The padding to use. Default: `"same"`
|
|
use_instance_norm: bool, optional
|
|
``True`` if instance normalization should be applied after the convolutional layer.
|
|
Default: ``False``
|
|
res_block_follows: bool, optional
|
|
If a residual block will follow this layer, then this should be set to `True` to add
|
|
a leaky ReLu after the convolutional layer. Default: ``False``
|
|
kwargs: dict
|
|
Any additional Keras standard layer keyword arguments
|
|
|
|
Returns
|
|
-------
|
|
tensor
|
|
The output tensor from the Convolution 2D Layer
|
|
"""
|
|
logger.debug("input_tensor: %s, filters: %s, kernel_size: %s, strides: %s, "
|
|
"use_instance_norm: %s, kwargs: %s)", input_tensor, filters, kernel_size,
|
|
strides, use_instance_norm, kwargs)
|
|
name = self._get_name("conv_{}".format(input_tensor.shape[1]))
|
|
if self.use_reflect_padding:
|
|
input_tensor = ReflectionPadding2D(
|
|
stride=strides,
|
|
kernel_size=kernel_size,
|
|
name="{}_reflectionpadding2d".format(name))(input_tensor)
|
|
padding = "valid"
|
|
var_x = self.conv2d(input_tensor, filters,
|
|
kernel_size=kernel_size,
|
|
strides=strides,
|
|
padding=padding,
|
|
name="{}_conv2d".format(name),
|
|
**kwargs)
|
|
if use_instance_norm:
|
|
var_x = InstanceNormalization(name="{}_instancenorm".format(name))(var_x)
|
|
if not res_block_follows:
|
|
var_x = LeakyReLU(0.1, name="{}_leakyrelu".format(name))(var_x)
|
|
return var_x
|
|
|
|
def upscale(self, input_tensor, filters, kernel_size=3, padding="same",
|
|
use_instance_norm=False, res_block_follows=False, scale_factor=2, **kwargs):
|
|
""" An upscale layer for sub-pixel up-scaling.
|
|
|
|
Adds reflection padding if it has been selected by the user, and other post-processing
|
|
if requested by the plugin.
|
|
|
|
Parameters
|
|
----------
|
|
input_tensor: tensor
|
|
The input tensor to the layer
|
|
filters: int
|
|
The dimensionality of the output space (i.e. the number of output filters in the
|
|
convolution)
|
|
kernel_size: int, optional
|
|
An integer or tuple/list of 2 integers, specifying the height and width of the 2D
|
|
convolution window. Can be a single integer to specify the same value for all spatial
|
|
dimensions. Default: 3
|
|
padding: ["valid", "same"], optional
|
|
The padding to use. Default: `"same"`
|
|
use_instance_norm: bool, optional
|
|
``True`` if instance normalization should be applied after the convolutional layer.
|
|
Default: ``False``
|
|
res_block_follows: bool, optional
|
|
If a residual block will follow this layer, then this should be set to `True` to add
|
|
a leaky ReLu after the convolutional layer. Default: ``False``
|
|
scale_factor: int, optional
|
|
The amount to upscale the image. Default: `2`
|
|
kwargs: dict
|
|
Any additional Keras standard layer keyword arguments
|
|
|
|
Returns
|
|
-------
|
|
tensor
|
|
The output tensor from the Upscale layer
|
|
"""
|
|
logger.debug("input_tensor: %s, filters: %s, kernel_size: %s, use_instance_norm: %s, "
|
|
"kwargs: %s)", input_tensor, filters, kernel_size, use_instance_norm, kwargs)
|
|
name = self._get_name("upscale_{}".format(input_tensor.shape[1]))
|
|
if self.use_reflect_padding:
|
|
input_tensor = ReflectionPadding2D(
|
|
stride=1,
|
|
kernel_size=kernel_size,
|
|
name="{}_reflectionpadding2d".format(name))(input_tensor)
|
|
padding = "valid"
|
|
kwargs = self._set_default_initializer(kwargs)
|
|
if self.use_icnr_init:
|
|
original_init = self._switch_kernel_initializer(
|
|
kwargs,
|
|
ICNR(initializer=kwargs["kernel_initializer"]))
|
|
var_x = self.conv2d(input_tensor, filters * scale_factor * scale_factor,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
name="{}_conv2d".format(name),
|
|
**kwargs)
|
|
if self.use_icnr_init:
|
|
self._switch_kernel_initializer(kwargs, original_init)
|
|
if use_instance_norm:
|
|
var_x = InstanceNormalization(name="{}_instancenorm".format(name))(var_x)
|
|
if not res_block_follows:
|
|
var_x = LeakyReLU(0.1, name="{}_leakyrelu".format(name))(var_x)
|
|
var_x = PixelShuffler(name="{}_pixelshuffler".format(name), size=scale_factor)(var_x)
|
|
return var_x
|
|
|
|
# <<< DLight Model Blocks >>> #
|
|
def upscale2x(self, input_tensor, filters,
|
|
kernel_size=3, padding="same", interpolation="bilinear", res_block_follows=False,
|
|
sr_ratio=0.5, scale_factor=2, fast=False, **kwargs):
|
|
""" Custom hybrid upscale layer for sub-pixel up-scaling.
|
|
|
|
Most of up-scaling is approximating lighting gradients which can be accurately achieved
|
|
using linear fitting. This layer attempts to improve memory consumption by splitting
|
|
with bilinear and convolutional layers so that the sub-pixel update will get details
|
|
whilst the bilinear filter will get lighting.
|
|
|
|
Adds reflection padding if it has been selected by the user, and other post-processing
|
|
if requested by the plugin.
|
|
|
|
Parameters
|
|
----------
|
|
input_tensor: tensor
|
|
The input tensor to the layer
|
|
filters: int
|
|
The dimensionality of the output space (i.e. the number of output filters in the
|
|
convolution)
|
|
kernel_size: int, optional
|
|
An integer or tuple/list of 2 integers, specifying the height and width of the 2D
|
|
convolution window. Can be a single integer to specify the same value for all spatial
|
|
dimensions. Default: 3
|
|
padding: ["valid", "same"], optional
|
|
The padding to use. Default: `"same"`
|
|
interpolation: ["nearest", "bilinear"], optional
|
|
Interpolation to use for up-sampling. Default: `"bilinear"`
|
|
res_block_follows: bool, optional
|
|
If a residual block will follow this layer, then this should be set to `True` to add
|
|
a leaky ReLu after the convolutional layer. Default: ``False``
|
|
scale_factor: int, optional
|
|
The amount to upscale the image. Default: `2`
|
|
sr_ratio: float, optional
|
|
The proportion of super resolution (pixel shuffler) filters to use. Non-fast mode only.
|
|
Default: `0.5`
|
|
kwargs: dict
|
|
Any additional Keras standard layer keyword arguments
|
|
fast: bool, optional
|
|
Use a faster up-scaling method that may appear more rugged. Default: ``False``
|
|
|
|
Returns
|
|
-------
|
|
tensor
|
|
The output tensor from the Upscale layer
|
|
"""
|
|
name = self._get_name("upscale2x_{}".format("fast" if fast else "hyb"))
|
|
var_x = input_tensor
|
|
if not fast:
|
|
sr_filters = int(filters * sr_ratio)
|
|
filters = filters - sr_filters
|
|
var_x_sr = self.upscale(var_x, filters,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
scale_factor=scale_factor,
|
|
res_block_follows=res_block_follows,
|
|
**kwargs)
|
|
|
|
if fast or (not fast and filters > 0):
|
|
var_x2 = self.conv2d(var_x, filters,
|
|
kernel_size=3,
|
|
padding=padding,
|
|
name="{}_conv2d".format(name),
|
|
**kwargs)
|
|
var_x2 = UpSampling2D(size=(scale_factor, scale_factor),
|
|
interpolation=interpolation,
|
|
name="{}_upsampling2D".format(name))(var_x2)
|
|
if fast:
|
|
var_x1 = self.upscale(var_x, filters,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
scale_factor=scale_factor,
|
|
res_block_follows=res_block_follows, **kwargs)
|
|
var_x = Add()([var_x2, var_x1])
|
|
else:
|
|
var_x = Concatenate(name="{}_concatenate".format(name))([var_x_sr, var_x2])
|
|
else:
|
|
var_x = var_x_sr
|
|
return var_x
|
|
|
|
# <<< DFaker Model Blocks >>> #
|
|
def res_block(self, input_tensor, filters, kernel_size=3, padding="same", **kwargs):
|
|
""" Residual block.
|
|
|
|
Parameters
|
|
----------
|
|
input_tensor: tensor
|
|
The input tensor to the layer
|
|
filters: int
|
|
The dimensionality of the output space (i.e. the number of output filters in the
|
|
convolution)
|
|
kernel_size: int, optional
|
|
An integer or tuple/list of 2 integers, specifying the height and width of the 2D
|
|
convolution window. Can be a single integer to specify the same value for all spatial
|
|
dimensions. Default: 3
|
|
padding: ["valid", "same"], optional
|
|
The padding to use. Default: `"same"`
|
|
kwargs: dict
|
|
Any additional Keras standard layer keyword arguments
|
|
|
|
Returns
|
|
-------
|
|
tensor
|
|
The output tensor from the Upscale layer
|
|
"""
|
|
logger.debug("input_tensor: %s, filters: %s, kernel_size: %s, kwargs: %s)",
|
|
input_tensor, filters, kernel_size, kwargs)
|
|
name = self._get_name("residual_{}".format(input_tensor.shape[1]))
|
|
var_x = LeakyReLU(alpha=0.2, name="{}_leakyrelu_0".format(name))(input_tensor)
|
|
if self.use_reflect_padding:
|
|
var_x = ReflectionPadding2D(stride=1,
|
|
kernel_size=kernel_size,
|
|
name="{}_reflectionpadding2d_0".format(name))(var_x)
|
|
padding = "valid"
|
|
var_x = self.conv2d(var_x, filters,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
name="{}_conv2d_0".format(name),
|
|
**kwargs)
|
|
var_x = LeakyReLU(alpha=0.2, name="{}_leakyrelu_1".format(name))(var_x)
|
|
if self.use_reflect_padding:
|
|
var_x = ReflectionPadding2D(stride=1,
|
|
kernel_size=kernel_size,
|
|
name="{}_reflectionpadding2d_1".format(name))(var_x)
|
|
padding = "valid"
|
|
if not self.use_convaware_init:
|
|
original_init = self._switch_kernel_initializer(kwargs, VarianceScaling(
|
|
scale=0.2,
|
|
mode="fan_in",
|
|
distribution="uniform"))
|
|
var_x = self.conv2d(var_x, filters,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
**kwargs)
|
|
if not self.use_convaware_init:
|
|
self._switch_kernel_initializer(kwargs, original_init)
|
|
var_x = Add()([var_x, input_tensor])
|
|
var_x = LeakyReLU(alpha=0.2, name="{}_leakyrelu_3".format(name))(var_x)
|
|
return var_x
|
|
|
|
# <<< Unbalanced Model Blocks >>> #
|
|
def conv_sep(self, input_tensor, filters, kernel_size=5, strides=2, **kwargs):
|
|
""" Seperable Convolution Layer.
|
|
|
|
Parameters
|
|
----------
|
|
input_tensor: tensor
|
|
The input tensor to the layer
|
|
filters: int
|
|
The dimensionality of the output space (i.e. the number of output filters in the
|
|
convolution)
|
|
kernel_size: int, optional
|
|
An integer or tuple/list of 2 integers, specifying the height and width of the 2D
|
|
convolution window. Can be a single integer to specify the same value for all spatial
|
|
dimensions. Default: 5
|
|
strides: tuple or int, optional
|
|
An integer or tuple/list of 2 integers, specifying the strides of the convolution along
|
|
the height and width. Can be a single integer to specify the same value for all spatial
|
|
dimensions. Default: `2`
|
|
kwargs: dict
|
|
Any additional Keras standard layer keyword arguments
|
|
|
|
Returns
|
|
-------
|
|
tensor
|
|
The output tensor from the Upscale layer
|
|
"""
|
|
logger.debug("input_tensor: %s, filters: %s, kernel_size: %s, strides: %s, kwargs: %s)",
|
|
input_tensor, filters, kernel_size, strides, kwargs)
|
|
name = self._get_name("separableconv2d_{}".format(input_tensor.shape[1]))
|
|
kwargs = self._set_default_initializer(kwargs)
|
|
var_x = SeparableConv2D(filters,
|
|
kernel_size=kernel_size,
|
|
strides=strides,
|
|
padding="same",
|
|
name="{}_seperableconv2d".format(name),
|
|
**kwargs)(input_tensor)
|
|
var_x = Activation("relu", name="{}_relu".format(name))(var_x)
|
|
return var_x
|