mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-18 22:45:59 +08:00
torzdf
614 lines
25 KiB
Python
614 lines
25 KiB
Python
#!/usr/bin/env python3
|
|
"""Processes the augmentation of images for feeding into a Faceswap model."""
|
|
from __future__ import annotations
|
|
import logging
|
|
import typing as T
|
|
from dataclasses import dataclass
|
|
|
|
import cv2
|
|
import numexpr as ne
|
|
import numpy as np
|
|
from scipy.interpolate import griddata
|
|
|
|
from lib.align.aligned_utils import batch_create_matrices
|
|
from lib.image import batch_convert_color
|
|
from lib.logger import format_array, parse_class_init
|
|
from lib.utils import get_module_objects
|
|
from plugins.train.trainer import trainer_config as cfg
|
|
|
|
if T.TYPE_CHECKING:
|
|
import numpy.typing as npt
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ConstantsColor:
|
|
"""Dataclass for holding constants for enhancing an image (ie contrast/color adjustment)
|
|
|
|
Parameters
|
|
----------
|
|
clahe_base_contrast
|
|
The base number for Contrast Limited Adaptive Histogram Equalization
|
|
clahe_chance
|
|
Probability to perform Contrast Limited Adaptive Histogram Equalization
|
|
clahe_max_size
|
|
Maximum clahe window size
|
|
lab_adjust
|
|
Adjustment amounts for L*A*B augmentation
|
|
"""
|
|
clahe_base_contrast: int
|
|
"""The base number for Contrast Limited Adaptive Histogram Equalization"""
|
|
clahe_chance: float
|
|
"""Probability to perform Contrast Limited Adaptive Histogram Equalization"""
|
|
clahe_max_size: int
|
|
"""Maximum clahe window size"""
|
|
lab_adjust: np.ndarray
|
|
"""Adjustment amounts for L*A*B augmentation"""
|
|
|
|
|
|
@dataclass
|
|
class ConstantsTransform:
|
|
"""Dataclass for holding constants for transforming an image
|
|
|
|
Parameters
|
|
----------
|
|
rotation
|
|
Rotation range for transformations
|
|
zoom
|
|
Zoom range for transformations
|
|
shift
|
|
Shift range for transformations
|
|
"""
|
|
rotation: int
|
|
"""Rotation range for transformations"""
|
|
zoom: float
|
|
"""Zoom range for transformations"""
|
|
shift: float
|
|
"""Shift range for transformations"""
|
|
flip: float
|
|
"""The chance to flip an image"""
|
|
|
|
|
|
@dataclass
|
|
class ConstantsWarp:
|
|
"""Dataclass for holding constants for warping an image
|
|
|
|
Parameters
|
|
----------
|
|
maps
|
|
The stacked (x, y) mappings for image warping
|
|
pad
|
|
The padding to apply for image warping
|
|
slices
|
|
The slices for extracting a warped image
|
|
lm_edge_anchors
|
|
The edge anchors for landmark based warping
|
|
lm_grids
|
|
The grids for landmark based warping
|
|
"""
|
|
maps: np.ndarray
|
|
"""The stacked (x, y) mappings for image warping"""
|
|
pad: tuple[int, int]
|
|
"""The padding to apply for image warping"""
|
|
slices: slice
|
|
"""The slices for extracting a warped image"""
|
|
scale: float
|
|
"""The scaling to apply to standard warping"""
|
|
lm_edge_anchors: np.ndarray
|
|
"""The edge anchors for landmark based warping"""
|
|
lm_grids: np.ndarray
|
|
"""The grids for landmark based warping"""
|
|
lm_scale: float
|
|
"""The scaling to apply to landmark based warping"""
|
|
|
|
def __repr__(self) -> str:
|
|
"""Display shape/type information for arrays in __repr__"""
|
|
params = {k: f"array[shape: {v.shape}, dtype: {v.dtype}]"
|
|
if isinstance(v, np.ndarray) else v
|
|
for k, v in self.__dict__.items()}
|
|
str_params = ", ".join(f"{k}={v}" for k, v in params.items())
|
|
return f"{self.__class__.__name__}({str_params})"
|
|
|
|
|
|
@dataclass
|
|
class ConstantsAugmentation:
|
|
"""Dataclass for holding constants for Image Augmentation.
|
|
|
|
Attributes
|
|
----------
|
|
color
|
|
The constants for adjusting color/contrast in an image
|
|
transform
|
|
The constants for image transformation
|
|
warp
|
|
The constants for image warping
|
|
|
|
Dataclass should be initialized using its :func:`from_config` method:
|
|
|
|
Example
|
|
-------
|
|
>>> constants = ConstantsAugmentation.from_config(processing_size=256,
|
|
... batch_size=16)
|
|
"""
|
|
color: ConstantsColor
|
|
"""The constants for adjusting color/contrast in an image"""
|
|
transform: ConstantsTransform
|
|
"""The constants for image transformation"""
|
|
warp: ConstantsWarp
|
|
"""The constants for image warping"""
|
|
|
|
@classmethod
|
|
def _get_clahe(cls, size: int) -> tuple[int, float, int]:
|
|
"""Get the CLAHE constants from user config
|
|
|
|
Parameters
|
|
----------
|
|
size
|
|
The size of image to augment the data for
|
|
|
|
Returns
|
|
-------
|
|
clahe_base_contrast
|
|
The base number for Contrast Limited Adaptive Histogram Equalization
|
|
clahe_chance
|
|
Probability to perform Contrast Limited Adaptive Histogram Equalization
|
|
clahe_max_size
|
|
Maximum clahe window size
|
|
"""
|
|
clahe_base_contrast = max(2, size // 128)
|
|
clahe_chance = cfg.Augmentation.color_clahe_chance() / 100
|
|
clahe_max_size = cfg.Augmentation.color_clahe_max_size()
|
|
logger.debug("[AugConstants] clahe_base_contrast: %s, clahe_chance: %s, "
|
|
"clahe_max_size: %s", clahe_base_contrast, clahe_chance, clahe_max_size)
|
|
return clahe_base_contrast, clahe_chance, clahe_max_size
|
|
|
|
@classmethod
|
|
def _get_lab(cls) -> np.ndarray:
|
|
"""Load the random L*A*B augmentation constants
|
|
|
|
Returns
|
|
-------
|
|
Adjustment amounts for L*A*B augmentation
|
|
"""
|
|
amount_l = cfg.Augmentation.color_lightness() / 100.
|
|
amount_ab = cfg.Augmentation.color_ab() / 100.
|
|
|
|
lab_adjust = np.array([amount_l, amount_ab, amount_ab], dtype="float32")
|
|
logger.debug("[AugConstants] lab_adjust: %s", lab_adjust)
|
|
return lab_adjust
|
|
|
|
@classmethod
|
|
def _get_color(cls, size: int) -> ConstantsColor:
|
|
"""Get the image enhancements constants from user config
|
|
|
|
Parameters
|
|
----------
|
|
size
|
|
The size of image to augment the data for
|
|
|
|
Returns
|
|
-------
|
|
The constants for image enhancement
|
|
"""
|
|
clahe_base_contrast, clahe_chance, clahe_max_size = cls._get_clahe(size)
|
|
retval = ConstantsColor(clahe_base_contrast=clahe_base_contrast,
|
|
clahe_chance=clahe_chance,
|
|
clahe_max_size=clahe_max_size,
|
|
lab_adjust=cls._get_lab())
|
|
logger.debug("[AugConstants] color: %s", retval)
|
|
return retval
|
|
|
|
@classmethod
|
|
def _get_transform(cls, size: int) -> ConstantsTransform:
|
|
"""Load the random transform constants
|
|
|
|
Parameters
|
|
----------
|
|
size
|
|
The size of image to augment the data for
|
|
|
|
Returns
|
|
-------
|
|
The constants for image transformation
|
|
"""
|
|
retval = ConstantsTransform(rotation=cfg.Augmentation.rotation_range(),
|
|
zoom=cfg.Augmentation.zoom_amount() / 100.,
|
|
shift=(cfg.Augmentation.shift_range() / 100.) * size,
|
|
flip=cfg.Augmentation.flip_chance() / 100.)
|
|
logger.debug("[AugConstants] transform: %s", retval)
|
|
return retval
|
|
|
|
@classmethod
|
|
def _get_warp_to_landmarks(cls, size: int, batch_size: int) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Load the warp-to-landmarks augmentation constants
|
|
|
|
Parameters
|
|
----------
|
|
size
|
|
The size of image to augment the data for
|
|
batch_size
|
|
The batch size that augmented data is being prepared for
|
|
|
|
Returns
|
|
-------
|
|
edge_anchors
|
|
The edge anchors for landmark based warping
|
|
grids
|
|
The grids for landmark based warping
|
|
"""
|
|
p_mx = size - 1
|
|
p_hf = (size // 2) - 1
|
|
edge_anchors = np.array([(0, 0), (0, p_mx), (p_mx, p_mx), (p_mx, 0),
|
|
(p_hf, 0), (p_hf, p_mx), (p_mx, p_hf), (0, p_hf)]).astype("int32")
|
|
edge_anchors = np.broadcast_to(edge_anchors, (batch_size, 8, 2))
|
|
grids = np.mgrid[0: p_mx: complex(size), # type:ignore[misc] # pylint:disable=no-member
|
|
0: p_mx: complex(size)].astype("float32") # type:ignore[misc]
|
|
|
|
logger.debug("[AugConstants] edge_anchors: (%s, %s), grids: (%s, %s)",
|
|
edge_anchors.shape, edge_anchors.dtype,
|
|
grids.shape, grids.dtype) # pylint:disable=no-member
|
|
return edge_anchors, grids
|
|
|
|
@classmethod
|
|
def _get_warp(cls, size: int, batch_size: int) -> ConstantsWarp:
|
|
"""Load the warp augmentation constants
|
|
|
|
Parameters
|
|
----------
|
|
size
|
|
The size of image to augment the data for
|
|
batch_size
|
|
The batch size that augmented data is being prepared for
|
|
|
|
Returns
|
|
-------
|
|
The constants for image warping
|
|
"""
|
|
lm_edge_anchors, lm_grids = cls._get_warp_to_landmarks(size, batch_size)
|
|
|
|
warp_range = np.linspace(0, size, 5, dtype='float32')
|
|
warp_map_x = np.broadcast_to(warp_range, (batch_size, 5, 5)).astype("float32")
|
|
warp_map_y = np.broadcast_to(warp_map_x[0].T, (batch_size, 5, 5)).astype("float32")
|
|
warp_pad = int(1.25 * size)
|
|
|
|
retval = ConstantsWarp(maps=np.stack((warp_map_x, warp_map_y), axis=1),
|
|
pad=(warp_pad, warp_pad),
|
|
slices=slice(warp_pad // 10, -warp_pad // 10),
|
|
scale=5 / 256 * size, # Normal random variable scale
|
|
lm_edge_anchors=lm_edge_anchors,
|
|
lm_grids=lm_grids,
|
|
lm_scale=2 / 256 * size) # Normal random variable scale
|
|
logger.debug("[AugConstants] warp constants: %s", retval)
|
|
return retval
|
|
|
|
@classmethod
|
|
def from_config(cls,
|
|
processing_size: int,
|
|
batch_size: int) -> ConstantsAugmentation:
|
|
"""Create a new dataclass instance from user config
|
|
|
|
Parameters
|
|
----------
|
|
processing_size
|
|
The size of image to augment the data for
|
|
batch_size
|
|
The batch size that augmented data is being prepared for
|
|
"""
|
|
logger.debug("[AugConstants] Initializing %s(processing_size=%s, batch_size=%s)",
|
|
cls.__name__, processing_size, batch_size)
|
|
retval = cls(color=cls._get_color(processing_size),
|
|
transform=cls._get_transform(processing_size),
|
|
warp=cls._get_warp(processing_size, batch_size))
|
|
logger.debug(retval)
|
|
return retval
|
|
|
|
|
|
class ImageAugmentation():
|
|
"""Performs augmentation on batches of training images.
|
|
|
|
Parameters
|
|
----------
|
|
batch_size
|
|
The number of images that will be fed through the augmentation functions at once.
|
|
processing_size
|
|
The largest input or output size of the model. This is the size that images are processed
|
|
at.
|
|
"""
|
|
def __init__(self, batch_size: int, processing_size: int) -> None:
|
|
logger.debug(parse_class_init(locals()))
|
|
self._processing_size = processing_size
|
|
self._batch_size = batch_size
|
|
self._constants = ConstantsAugmentation.from_config(processing_size, batch_size)
|
|
logger.debug("[Aug] Initialized %s", self.__class__.__name__)
|
|
|
|
def __repr__(self) -> str:
|
|
"""Pretty print this object"""
|
|
return (f"{self.__class__.__name__}(batch_size={self._batch_size}, "
|
|
f"processing_size={self._processing_size})")
|
|
|
|
# <<< COLOR AUGMENTATION >>> #
|
|
def _random_lab(self, batch: np.ndarray) -> None:
|
|
"""Perform random color/lightness adjustment in L*a*b* color space on a batch of
|
|
images
|
|
|
|
Parameters
|
|
----------
|
|
batch
|
|
The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`,
|
|
`3`) and in `BGR` format of uint8 dtype.
|
|
"""
|
|
randoms = np.random.uniform(-self._constants.color.lab_adjust,
|
|
self._constants.color.lab_adjust,
|
|
size=(self._batch_size, 1, 1, 3)).astype("float32")
|
|
logger.trace("[Aug] Random LAB adjustments: %s", randoms) # type:ignore[attr-defined]
|
|
# Iterating through the images and channels is much faster than numpy.where and slightly
|
|
# faster than numexpr.where.
|
|
for image, rand in zip(batch, randoms):
|
|
for idx in range(rand.shape[-1]):
|
|
adjustment = rand[:, :, idx]
|
|
if adjustment >= 0:
|
|
image[:, :, idx] = ((255 - image[:, :, idx]) * adjustment) + image[:, :, idx]
|
|
else:
|
|
image[:, :, idx] = image[:, :, idx] * (1 + adjustment)
|
|
|
|
def _random_clahe(self, batch: np.ndarray) -> None:
|
|
"""Randomly perform Contrast Limited Adaptive Histogram Equalization on
|
|
a batch of images
|
|
|
|
Parameters
|
|
----------
|
|
batch
|
|
The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`,
|
|
`3`) and in `BGR` format of uint8 dtype.
|
|
"""
|
|
base_contrast = self._constants.color.clahe_base_contrast
|
|
|
|
batch_random = np.random.rand(self._batch_size)
|
|
indices = np.where(batch_random < self._constants.color.clahe_chance)[0]
|
|
if not np.any(indices):
|
|
return
|
|
grid_bases = np.random.randint(self._constants.color.clahe_max_size + 1,
|
|
size=indices.shape[0],
|
|
dtype="uint8")
|
|
grid_sizes = (grid_bases * (base_contrast // 2)) + base_contrast
|
|
logger.trace("[Aug] Adjusting Contrast. Grid Sizes: %s", # type:ignore[attr-defined]
|
|
grid_sizes)
|
|
|
|
clahes = [cv2.createCLAHE(clipLimit=2.0,
|
|
tileGridSize=(grid_size, grid_size))
|
|
for grid_size in grid_sizes] # type:ignore[attr-defined]
|
|
|
|
for idx, clahe in zip(indices, clahes):
|
|
batch[idx, :, :, 0] = clahe.apply(batch[idx, :, :, 0], )
|
|
|
|
def color_adjust(self, batch: np.ndarray) -> np.ndarray:
|
|
"""Perform color augmentation on the passed in batch.
|
|
|
|
The color adjustment parameters are set in :file:`config.train.ini`
|
|
|
|
Parameters
|
|
----------
|
|
batch
|
|
The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`,
|
|
`3`) and in `BGR` format of uint8 dtype.
|
|
|
|
Returns
|
|
----------
|
|
A 4-dimensional array of the same shape as :attr:`batch` with color augmentation applied.
|
|
"""
|
|
logger.trace("[Aug] Augmenting color") # type:ignore[attr-defined]
|
|
batch = batch_convert_color(batch, "BGR2LAB")
|
|
self._random_lab(batch)
|
|
self._random_clahe(batch)
|
|
batch = batch_convert_color(batch, "LAB2BGR")
|
|
return batch
|
|
|
|
# <<< IMAGE AUGMENTATION >>> #
|
|
def transform(self, batch: npt.NDArray[np.uint8], points: npt.NDArray[np.float32] | None
|
|
) -> None:
|
|
"""Perform random transformation on the passed in batch and optional (x, y) points.
|
|
|
|
The transformation parameters are set in :file:`config.train.ini`
|
|
|
|
Parameters
|
|
----------
|
|
batch
|
|
The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`,
|
|
`channels`) and in `BGR` format.
|
|
points
|
|
Any (x, y) points to transform. in shape (batch_size, num_sides, 68, 2). ``None`` if
|
|
there are no points to transform
|
|
"""
|
|
logger.trace("[Aug] Randomly transforming image") # type:ignore[attr-defined]
|
|
rotation = np.random.uniform(-self._constants.transform.rotation,
|
|
self._constants.transform.rotation,
|
|
size=self._batch_size).astype("float32")
|
|
scale = np.random.uniform(1 - self._constants.transform.zoom,
|
|
1 + self._constants.transform.zoom,
|
|
size=self._batch_size).astype("float32")
|
|
|
|
transform = np.random.uniform(-self._constants.transform.shift,
|
|
self._constants.transform.shift,
|
|
size=(self._batch_size, 2)).astype("float32")
|
|
mats = batch_create_matrices(self._processing_size,
|
|
rotation,
|
|
scale=scale,
|
|
translation=transform)
|
|
|
|
for image, mat in zip(batch, mats[:, :2, :]):
|
|
cv2.warpAffine(image,
|
|
mat,
|
|
(self._processing_size, self._processing_size),
|
|
dst=image,
|
|
borderMode=cv2.BORDER_REPLICATE)
|
|
|
|
logger.trace("[Aug] Randomly transformed image") # type:ignore[attr-defined]
|
|
if points is None:
|
|
return
|
|
ones = np.ones((*points.shape[:-1], 1), dtype=points.dtype)
|
|
pts_h = np.concatenate([points, ones], axis=-1)
|
|
points[:] = np.einsum('nij,n...j->n...i', mats, pts_h)[..., :2]
|
|
logger.trace("[Aug] Randomly transformed points") # type:ignore[attr-defined]
|
|
|
|
def random_flip(self, batch: npt.NDArray[np.uint8], points: npt.NDArray[np.float32] | None
|
|
) -> None:
|
|
"""Perform random horizontal flipping on the passed in batch.
|
|
|
|
The probability of flipping an image is set in :file:`config.train.ini`
|
|
|
|
Parameters
|
|
----------
|
|
batch
|
|
The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`,
|
|
`channels`) and in `BGR` format.
|
|
points
|
|
Any (x, y) points to transform. Can be in any shape but the final dimension should be
|
|
shape 2. ``None`` if there are no points to transform
|
|
"""
|
|
logger.trace("[Aug] Randomly flipping image") # type:ignore[attr-defined]
|
|
randoms = np.random.rand(self._batch_size)
|
|
indices = np.where(randoms <= self._constants.transform.flip)[0]
|
|
batch[indices] = batch[indices, :, ::-1]
|
|
logger.trace("[Aug] Randomly flipped %s images of %s", # type:ignore[attr-defined]
|
|
len(indices), self._batch_size)
|
|
if points is None:
|
|
return
|
|
points[indices, ..., 0] = (self._processing_size - 1) - points[indices, ..., 0]
|
|
logger.trace("[Aug] Randomly flipped %s points: %s", # type:ignore[attr-defined]
|
|
len(indices), format_array(points))
|
|
|
|
def _random_warp(self, batch: np.ndarray) -> np.ndarray:
|
|
"""Randomly warp the input batch
|
|
|
|
Parameters
|
|
----------
|
|
batch : :class:`numpy.ndarray`
|
|
The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`,
|
|
`3`) and in `BGR` format.
|
|
|
|
Returns
|
|
----------
|
|
A 4-dimensional array of the same shape as :attr:`batch` with warping applied.
|
|
"""
|
|
logger.trace("[Aug] Randomly warping batch") # type:ignore[attr-defined]
|
|
slices = self._constants.warp.slices
|
|
rands = np.random.normal(size=(self._batch_size, 2, 5, 5),
|
|
scale=self._constants.warp.scale).astype("float32")
|
|
batch_maps = ne.evaluate("m + r", local_dict={"m": self._constants.warp.maps, "r": rands})
|
|
|
|
interpolators = np.array([[cv2.resize(map_, self._constants.warp.pad)[slices, slices]
|
|
for map_ in maps]
|
|
for maps in batch_maps])
|
|
warped_batch = np.array([cv2.remap(image,
|
|
interpolator[0],
|
|
interpolator[1],
|
|
cv2.INTER_LINEAR)
|
|
for image, interpolator in zip(batch, interpolators)])
|
|
|
|
logger.trace("[Aug] Warped image shape: %s", # type:ignore[attr-defined]
|
|
warped_batch.shape)
|
|
return warped_batch
|
|
|
|
def _random_warp_landmarks(self,
|
|
batch: np.ndarray,
|
|
batch_src_points: np.ndarray,
|
|
batch_dst_points: np.ndarray) -> np.ndarray:
|
|
"""From dfaker. Warp the image to a similar set of landmarks from the opposite side
|
|
|
|
batch
|
|
The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`,
|
|
`3`) and in `BGR` format.
|
|
batch_src_points
|
|
A batch of 68 point landmarks for the source faces. This is a 3-dimensional array in
|
|
the shape (`batchsize`, `68`, `2`).
|
|
batch_dst_points
|
|
A batch of randomly chosen closest match destination faces landmarks. This is a
|
|
3-dimensional array in the shape (`batchsize`, `68`, `2`).
|
|
|
|
Returns
|
|
----------
|
|
A 4-dimensional array of the same shape as :attr:`batch` with warping applied.
|
|
"""
|
|
logger.trace("[Aug] Randomly warping landmarks") # type:ignore[attr-defined]
|
|
edge_anchors = self._constants.warp.lm_edge_anchors
|
|
grids = self._constants.warp.lm_grids
|
|
|
|
batch_dst = batch_dst_points + np.random.normal(size=batch_dst_points.shape,
|
|
scale=self._constants.warp.lm_scale)
|
|
|
|
face_cores = [cv2.convexHull(np.concatenate([src[17:], dst[17:]], axis=0))
|
|
for src, dst in zip(batch_src_points.astype("int32"),
|
|
batch_dst.astype("int32"))]
|
|
|
|
batch_src = np.append(batch_src_points, edge_anchors, axis=1)
|
|
batch_dst = np.append(batch_dst, edge_anchors, axis=1)
|
|
|
|
rem_indices = [list(set(idx for fpl in (src, dst)
|
|
for idx, (pty, ptx) in enumerate(fpl)
|
|
if cv2.pointPolygonTest(face_core, (pty, ptx), False) >= 0))
|
|
for src, dst, face_core in zip(batch_src[:, :18, :],
|
|
batch_dst[:, :18, :],
|
|
face_cores)]
|
|
lm_batch_src = [np.delete(src, indices, axis=0)
|
|
for indices, src in zip(rem_indices, batch_src)]
|
|
lm_batch_dst = [np.delete(dst, indices, axis=0)
|
|
for indices, dst in zip(rem_indices, batch_dst)]
|
|
|
|
grid_z = np.array([griddata(dst, src, (grids[0], grids[1]), method="linear")
|
|
for src, dst in zip(lm_batch_src, lm_batch_dst)])
|
|
maps = grid_z.reshape((self._batch_size,
|
|
self._processing_size,
|
|
self._processing_size,
|
|
2)).astype("float32")
|
|
|
|
warped_batch = np.array([cv2.remap(image,
|
|
map_[..., 1],
|
|
map_[..., 0],
|
|
cv2.INTER_LINEAR,
|
|
borderMode=cv2.BORDER_TRANSPARENT)
|
|
for image, map_ in zip(batch, maps)])
|
|
logger.trace("[Aug] Warped batch shape: %s", # type:ignore[attr-defined]
|
|
warped_batch.shape)
|
|
return warped_batch
|
|
|
|
def warp(self,
|
|
batch: np.ndarray,
|
|
to_landmarks: bool = False,
|
|
batch_src_points: np.ndarray | None = None,
|
|
batch_dst_points: np.ndarray | None = None
|
|
) -> np.ndarray:
|
|
|
|
"""Perform random warping on the passed in batch by one of two methods.
|
|
|
|
Parameters
|
|
----------
|
|
batch
|
|
The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`,
|
|
`3`) and in `BGR` format.
|
|
to_landmarks
|
|
If ``False`` perform standard random warping of the input image. If ``True`` perform
|
|
warping to semi-random similar corresponding landmarks from the other side. Default:
|
|
``False``
|
|
batch_src_points
|
|
Only used when :attr:`to_landmarks` is ``True``. A batch of 68 point landmarks for the
|
|
source faces. This is a 3-dimensional array in the shape (`batchsize`, `68`, `2`).
|
|
Default: ``None``
|
|
batch_dst_points
|
|
Only used when :attr:`to_landmarks` is ``True``. A batch of randomly chosen closest
|
|
match destination faces landmarks. This is a 3-dimensional array in the shape
|
|
(`batchsize`, `68`, `2`). Default ``None``
|
|
|
|
Returns
|
|
----------
|
|
A 4-dimensional array of the same shape as :attr:`batch` with warping applied.
|
|
"""
|
|
if to_landmarks:
|
|
assert batch_src_points is not None
|
|
assert batch_dst_points is not None
|
|
return self._random_warp_landmarks(batch, batch_src_points, batch_dst_points)
|
|
return self._random_warp(batch)
|
|
|
|
|
|
__all__ = get_module_objects(__name__)
|