#!/usr/bin/env python3 """ Processes the augmentation of images for feeding into a Faceswap model. """ from __future__ import annotations import logging from dataclasses import dataclass import cv2 import numexpr as ne import numpy as np from scipy.interpolate import griddata from lib.image import batch_convert_color from lib.logger import parse_class_init from lib.utils import get_module_objects from plugins.train.trainer import trainer_config as cfg logger = logging.getLogger(__name__) @dataclass class ConstantsColor: """ Dataclass for holding constants for enhancing an image (ie contrast/color adjustment) Parameters ---------- clahe_base_contrast : int The base number for Contrast Limited Adaptive Histogram Equalization clahe_chance : float Probability to perform Contrast Limited Adaptive Histogram Equilization clahe_max_size : int Maximum clahe window size lab_adjust : :class:`numpy.ndarray` Adjustment amounts for L*A*B augmentation """ clahe_base_contrast: int """ int : The base number for Contrast Limited Adaptive Histogram Equalization """ clahe_chance: float """ float : Probability to perform Contrast Limited Adaptive Histogram Equilization """ clahe_max_size: int """ int : Maximum clahe window size""" lab_adjust: np.ndarray """ :class:`numpy.ndarray` : Adjustment amounts for L*A*B augmentation """ @dataclass class ConstantsTransform: """ Dataclass for holding constants for transforming an image Parameters ---------- rotation : int Rotation range for transformations zoom : float Zoom range for transformations shift : float Shift range for transformations """ rotation: int """ int : Rotation range for transformations """ zoom: float """ float : Zoom range for transformations """ shift: float """ float : Shift range for transformations """ flip: float """ float : The chance to flip an image """ @dataclass class ConstantsWarp: """ Dataclass for holding constants for warping an image Parameters ---------- maps : :class:`numpy.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 lm_edge_anchors : :class:`numpy.ndarray` The edge anchors for landmark based warping lm_grids : :class:`numpy.ndarray` The grids for landmark based warping """ maps: np.ndarray """ :class:`numpy.ndarray` : The stacked (x, y) mappings for image warping """ pad: tuple[int, int] """ :tuple[int, int] : The padding to apply for image warping """ slices: slice """ slice : The slices for extracting a warped image """ scale: float """ float : The scaling to apply to standard warping """ lm_edge_anchors: np.ndarray """ :class:`numpy.ndarray` : The edge anchors for landmark based warping """ lm_grids: np.ndarray """ :class:`numpy.ndarray` : The grids for landmark based warping """ lm_scale: float """ 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 : :class:`ConstantsColor` The constants for adjusting color/contrast in an image transform : :class:`ConstantsTransform` The constants for image transformation warp : :class:`ConstantsTransform` 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 """ :class:`ConstantsColor` : The constants for adjusting color/contrast in an image """ transform: ConstantsTransform """ :class:`ConstantsTransform` : The constants for image transformation """ warp: ConstantsWarp """ :class:`ConstantsTransform` : 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 : int The size of image to augment the data for Returns ------- clahe_base_contrast : int The base number for Contrast Limited Adaptive Histogram Equalization clahe_chance : float Probability to perform Contrast Limited Adaptive Histogram Equilization clahe_max_size : int Maximum clahe window size """ clahe_base_contrast = max(2, size // 128) clahe_chance = cfg.color_clahe_chance() / 100 clahe_max_size = cfg.color_clahe_max_size() logger.debug("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 ------- :class:`numpy.ndarray` Adjustment amounts for L*A*B augmentation """ amount_l = cfg.color_lightness() / 100. amount_ab = cfg.color_ab() / 100. lab_adjust = np.array([amount_l, amount_ab, amount_ab], dtype="float32") logger.debug("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 : int The size of image to augment the data for Returns ------- :class:`ConstantsColor` 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(retval) return retval @classmethod def _get_transform(cls, size: int) -> ConstantsTransform: """ Load the random transform constants Parameters ---------- size : int The size of image to augment the data for Returns ------- :class:`ConstantsTransform` The constants for image transformation """ retval = ConstantsTransform(rotation=cfg.rotation_range(), zoom=cfg.zoom_amount() / 100., shift=(cfg.shift_range() / 100.) * size, flip=cfg.flip_chance() / 100.) logger.debug(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 : int The size of image to augment the data for batch_size : int The batch size that augmented data is being prepared for Returns ------- edge_anchors : :class:`numpy.ndarray` The edge anchors for landmark based warping grids : :class:`numpy.ndarray` 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("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: int The size of image to augment the data for batch_size : int The batch size that augmented data is being prepared for Returns ------- :class:`ConstantsTransform` 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_mapx = np.broadcast_to(warp_range, (batch_size, 5, 5)).astype("float32") warp_mapy = np.broadcast_to(warp_mapx[0].T, (batch_size, 5, 5)).astype("float32") warp_pad = int(1.25 * size) retval = ConstantsWarp(maps=np.stack((warp_mapx, warp_mapy), 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(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 : int: The size of image to augment the data for batch_size : int The batch size that augmented data is being prepared for """ logger.debug("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 : int The number of images that will be fed through the augmentation functions at once. processing_size: int 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("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 : :class:`numpy.ndarray` 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("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 : :class:`numpy.ndarray` 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("Adjusting Contrast. Grid Sizes: %s", grid_sizes) # type:ignore[attr-defined] clahes = [cv2.createCLAHE(clipLimit=2.0, tileGridSize=(grid_size, grid_size)) for grid_size in grid_sizes] 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 : :class:`numpy.ndarray` The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, `3`) and in `BGR` format of uint8 dtype. Returns ---------- :class:`numpy.ndarray` A 4-dimensional array of the same shape as :attr:`batch` with color augmentation applied. """ logger.trace("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: np.ndarray): """ Perform random transformation on the passed in batch. The transformation parameters are set in :file:`config.train.ini` Parameters ---------- batch : :class:`numpy.ndarray` The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, `channels`) and in `BGR` format. """ logger.trace("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") tform = np.random.uniform(-self._constants.transform.shift, self._constants.transform.shift, size=(self._batch_size, 2)).astype("float32") mats = np.array( [cv2.getRotationMatrix2D((self._processing_size // 2, self._processing_size // 2), rot, scl) for rot, scl in zip(rotation, scale)]).astype("float32") mats[..., 2] += tform for image, mat in zip(batch, mats): cv2.warpAffine(image, mat, (self._processing_size, self._processing_size), dst=image, borderMode=cv2.BORDER_REPLICATE) logger.trace("Randomly transformed image") # type:ignore[attr-defined] def random_flip(self, batch: np.ndarray): """ Perform random horizontal flipping on the passed in batch. The probability of flipping an image is set in :file:`config.train.ini` Parameters ---------- batch : :class:`numpy.ndarray` The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, `channels`) and in `BGR` format. """ logger.trace("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("Randomly flipped %s images of %s", # type:ignore[attr-defined] len(indices), self._batch_size) 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 ---------- :class:`numpy.ndarray` A 4-dimensional array of the same shape as :attr:`batch` with warping applied. """ logger.trace("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}) batch_interp = 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, interp[0], interp[1], cv2.INTER_LINEAR) for image, interp in zip(batch, batch_interp)]) logger.trace("Warped image shape: %s", warped_batch.shape) # type:ignore[attr-defined] 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 : :class:`numpy.ndarray` The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, `3`) and in `BGR` format. batch_src_points : :class:`numpy.ndarray` 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 : :class:`numpy.ndarray` A batch of randomly chosen closest match destination faces landmarks. This is a 3-dimensional array in the shape (`batchsize`, `68`, `2`). Returns ---------- :class:`numpy.ndarray` A 4-dimensional array of the same shape as :attr:`batch` with warping applied. """ logger.trace("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)] lbatch_src = [np.delete(src, idxs, axis=0) for idxs, src in zip(rem_indices, batch_src)] lbatch_dst = [np.delete(dst, idxs, axis=0) for idxs, 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(lbatch_src, lbatch_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("Warped batch shape: %s", warped_batch.shape) # type:ignore[attr-defined] 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 : :class:`numpy.ndarray` The batch should be a 4-dimensional array of shape (`batchsize`, `height`, `width`, `3`) and in `BGR` format. to_landmarks : bool, optional 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 : :class:`numpy.ndarray`, optional 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 : :class:`numpy.ndarray`, optional 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 ---------- :class:`numpy.ndarray` 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__)