          if np.random.random() < random_flip:
              logger.trace("Flip within threshold. Flipping")
              retval = image[:, ::-1]
          else:
              logger.trace("Flip outside threshold. Not Flipping")
              retval = image
          logger.trace("Randomly flipped image")
          return retval
  
      def random_warp(self, image):
          """ get pair of random warped images from aligned face image """
          logger.trace("Randomly warping image")
 
</s> remove     def random_transform(self, image):
        """ Randomly transform an image """
        logger.trace("Randomly transforming image")
        height, width = image.shape[0:2]

        rotation_range = self.config.get("rotation_range", 10)
        rotation = np.random.uniform(-rotation_range, rotation_range)

        zoom_range = self.config.get("zoom_range", 5) / 100
        scale = np.random.uniform(1 - zoom_range, 1 + zoom_range)

        shift_range = self.config.get("shift_range", 5) / 100
        tnx = np.random.uniform(-shift_range, shift_range) * width
        tny = np.random.uniform(-shift_range, shift_range) * height

        mat = cv2.getRotationMatrix2D(  # pylint:disable=no-member
            (width // 2, height // 2), rotation, scale)
        mat[:, 2] += (tnx, tny)
        result = cv2.warpAffine(  # pylint:disable=no-member
            image, mat, (width, height),
            borderMode=cv2.BORDER_REPLICATE)  # pylint:disable=no-member

        logger.trace("Randomly transformed image")
        return result

    def do_random_flip(self, image):
        """ Perform flip on image if random number is within threshold """
        logger.trace("Randomly flipping image")
        random_flip = self.config.get("random_flip", 50) / 100
        if np.random.random() < random_flip:
            logger.trace("Flip within threshold. Flipping")
            retval = image[:, ::-1]
 </s> add         Returns the targets as a list of 4-dimensional ``numpy.ndarray`` s of shape (`batchsize`,
        `height`, `width`, 3). If the :attr:`batch` is 4 channels, then the masks will be split
        from the batch, with the largest output masks being returned in their own item.
        """
        batch = [tgt.astype("float32") / 255.0 for tgt in batch]
        if all(tgt.shape[-1] == 4 for tgt in batch):
            logger.trace("Batch contains mask")
            sizes = [item.shape[1] for item in batch]
            mask_batch = np.expand_dims(batch[sizes.index(max(sizes))][..., -1], axis=-1)
            batch = [item[..., :3] for item in batch]
            logger.trace("batch shapes: %s, mask_batch shape: %s",
                         [tgt.shape for tgt in batch], mask_batch.shape)
            retval = dict(targets=batch, masks=mask_batch) </s> remove     def random_warp(self, image):
        """ get pair of random warped images from aligned face image """
        logger.trace("Randomly warping image")
        height, width = image.shape[0:2]
        coverage = self.get_coverage(image) // 2
        try:
            assert height == width and height % 2 == 0
        except AssertionError as err:
            msg = ("Training images should be square with an even number of pixels across each "
                   "side. An image was found with width: {}, height: {}."
                   "\nMost likely this is a frame rather than a face within your training set. "
                   "\nMake sure that the only images within your training set are faces generated "
                   "from the Extract process.".format(width, height))
            raise FaceswapError(msg) from err

        range_ = np.linspace(height // 2 - coverage, height // 2 + coverage, 5, dtype='float32')
        mapx = np.broadcast_to(range_, (5, 5)).copy()
        mapy = mapx.T
        # mapx, mapy = np.float32(np.meshgrid(range_,range_)) # instead of broadcast

        pad = int(1.25 * self.input_size)
        slices = slice(pad // 10, -pad // 10)
        dst_slices = [slice(0, (size + 1), (size // 4)) for size in self.output_sizes]
        interp = np.empty((2, self.input_size, self.input_size), dtype='float32')
 </s> add     # <<< COLOR AUGMENTATION >>> #
    def color_adjust(self, batch):
        """ Perform color augmentation on the passed in batch. </s> remove         logger.trace("Target image shapes: %s", [tgt.shape for tgt in target_images])
        return self.compile_images(warped_image, target_images)
 </s> add         logger.trace("Randomly transformed image")
        return batch </s> remove     def separate_mask(image):
        """ Return the image and the mask from a 4 channel image """
        mask = None
        if image.shape[2] == 4:
            logger.trace("Image contains mask")
            mask = np.expand_dims(image[:, :, -1], axis=2)
            image = image[:, :, :3]
        else:
            logger.trace("Image has no mask")
        return image, mask

    def get_coverage(self, image):
        """ Return coverage value for given image """
        coverage = int(image.shape[0] * self.coverage_ratio)
        logger.trace("Coverage: %s", coverage)
        return coverage
 </s> add     def _separate_target_mask(batch):
        """ Return the batch and the batch of final masks </s> remove     def get_transpose_axes(num):
        if num % 2 == 0:
            logger.debug("Even number of images to stack")
            y_axes = list(range(1, num - 1, 2))
            x_axes = list(range(0, num - 1, 2))
        else:
            logger.debug("Odd number of images to stack")
            y_axes = list(range(0, num - 1, 2))
            x_axes = list(range(1, num - 1, 2))
        return y_axes, x_axes, [num - 1]

    images_shape = np.array(images.shape)
    new_axes = get_transpose_axes(len(images_shape))
    new_shape = [np.prod(images_shape[x]) for x in new_axes]
    logger.debug("Stacked images")
    return np.transpose(
        images,
        axes=np.concatenate(new_axes)
        ).reshape(new_shape)
 </s> add         Returns
        ----------
        numpy.ndarray
            A 4-dimensional array of the same shape as :attr:`batch` with transformation applied.
        """
        if not self.is_display:
            logger.trace("Randomly flipping image")
            randoms = np.random.rand(self._batchsize)
            indices = np.where(randoms > self._config.get("random_flip", 50) / 100)[0]
            batch[indices] = batch[indices, :, ::-1]
            logger.trace("Randomly flipped %s images of %s", len(indices), self._batchsize)
        return batch

    def warp(self, batch, to_landmarks=False, **kwargs):
        """ Perform random warping on the passed in batch by one of two methods.

        Parameters
        ----------
        batch: 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``
        kwargs: dict
            If :attr:`to_landmarks` is ``True`` the following additional kwargs must be passed in:

            * **batch_src_points** (`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** (`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
        ----------
        numpy.ndarray
            A 4-dimensional array of the same shape as :attr:`batch` with warping applied.
        """
        if to_landmarks:
            return self._random_warp_landmarks(batch, **kwargs).astype("float32") / 255.0
        return self._random_warp(batch).astype("float32") / 255.0

    def _random_warp(self, batch):
        """ Randomly warp the input batch """
        logger.trace("Randomly warping batch")
        mapx = self._constants["warp_mapx"]
        mapy = self._constants["warp_mapy"]
        pad = self._constants["warp_pad"]
        slices = self._constants["warp_slices"]

        rands = np.random.normal(size=(self._batchsize, 2, 5, 5),
                                 scale=self._scale).astype("float32")
        batch_maps = np.stack((mapx, mapy), axis=1) + rands
        batch_interp = np.array([[cv2.resize(map_, (pad, 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)
        return warped_batch

    def _random_warp_landmarks(self, batch, batch_src_points, batch_dst_points):
        """ From dfaker. Warp the image to a similar set of landmarks from the opposite side """
        logger.trace("Randomly warping landmarks")
        edge_anchors = self._constants["warp_lm_edge_anchors"]
        grids = self._constants["warp_lm_grids"]
        slices = self._constants["tgt_slices"]

        batch_dst = (batch_dst_points + np.random.normal(size=batch_dst_points.shape,
                                                         scale=2.0)).astype("int32")

        face_cores = [cv2.convexHull(np.concatenate([src[17:], dst[17:]], axis=0))
                      for src, dst in zip(batch_src_points, batch_dst)]

        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)]
        batch_src = [np.delete(src, idxs, axis=0) for idxs, src in zip(rem_indices, batch_src)]
        batch_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(batch_src, batch_dst)])
        maps = grid_z.reshape(self._batchsize,
                              self._training_size,
                              self._training_size,
                              2).astype("float32")
        warped_batch = np.array([cv2.remap(image,
                                           map_[..., 1],
                                           map_[..., 0],
                                           cv2.INTER_LINEAR,
                                           cv2.BORDER_TRANSPARENT)
                                 for image, map_ in zip(batch, maps)])
        warped_batch = np.array([cv2.resize(image[slices, slices, :],
                                            (self._input_size, self._input_size),
                                            cv2.INTER_AREA)
                                 for image in warped_batch])
        logger.trace("Warped batch shape: %s", warped_batch.shape)
        return warped_batch
            logger.trace("Batch has no mask")
            retval = dict(targets=batch)
          if np.random.random() < random_flip:
              logger.trace("Flip within threshold. Flipping")
              retval = image[:, ::-1]
          else:
            logger.trace("Batch has no mask")
            retval = dict(targets=batch)
            logger.trace("Batch has no mask")
            retval = dict(targets=batch)
            logger.trace("Batch has no mask")
            retval = dict(targets=batch)
          return retval
  
      def random_warp(self, image):
          """ get pair of random warped images from aligned face image """
          logger.trace("Randomly warping image")
keep keep keep keep replace replace replace keep keep keep keep keep
 Optimize Data Augmentation (#881)

* Move image utils to lib.image
* Add .pylintrc file
* Remove some cv2 pylint ignores
* TrainingData: Load images from disk in batches
* TrainingData: get_landmarks to batch
* TrainingData: transform and flip to batches
* TrainingData: Optimize color augmentation
* TrainingData: Optimize target and random_warp
* TrainingData - Convert _get_closest_match for batching
* TrainingData: Warp To Landmarks optimized
* Save models to threadpoolexecutor
* Move stack_images, Rename ImageManipulation. ImageAugmentation Docstrings
* Masks: Set dtype and threshold for lib.masks based on input face
* Docstrings and Documentation 
lib/training_data.py