          return mask_class
  
      def minibatch_ab(self, images, batchsize, side,
                       do_shuffle=True, is_preview=False, is_timelapse=False):
          """ Keep a queue filled to 8x Batch Size """
          logger.debug("Queue batches: (image_count: %s, batchsize: %s, side: '%s', do_shuffle: %s, "
                       "is_preview, %s, is_timelapse: %s)", len(images), batchsize, side, do_shuffle,
                       is_preview, is_timelapse)
          self.batchsize = batchsize
          is_display = is_preview or is_timelapse
 
</s> remove         self.batchsize = batchsize
        is_display = is_preview or is_timelapse
        args = (images, side, is_display, do_shuffle, batchsize)
        batcher = BackgroundGenerator(self.minibatch, thread_count=2, args=args)
 </s> add         self._batchsize = batchsize
        self._processing = ImageAugmentation(batchsize,
                                             is_preview or is_timelapse,
                                             self._model_input_size,
                                             self._model_output_shapes,
                                             self._training_opts.get("coverage_ratio", 0.625),
                                             self._config)
        args = (images, side, do_shuffle, batchsize)
        batcher = BackgroundGenerator(self._minibatch, thread_count=2, args=args) </s> remove     def set_mask_class(self):
        """ Set the mask function to use if using mask """
        mask_type = self.training_opts.get("mask_type", None)
        if mask_type:
            logger.debug("Mask type: '%s'", mask_type)
            mask_class = getattr(masks, mask_type)
        else:
            mask_class = None
        logger.debug("Mask class: %s", mask_class)
        return mask_class
 </s> add         # Batchsize and processing class are set when this class is called by a batcher
        # from lib.training_data
        self._batchsize = 0
        self._processing = None
        logger.debug("Initialized %s", self.__class__.__name__) </s> remove     def minibatch(self, images, side, is_display, do_shuffle, batchsize):
        """ A generator function that yields epoch, batchsize of warped_img
            and batchsize of target_img from the load queue """
        logger.debug("Loading minibatch generator: (image_count: %s, side: '%s', is_display: %s, "
                     "do_shuffle: %s)", len(images), side, is_display, do_shuffle)
        self.validate_samples(images)
 </s> add     def _minibatch(self, images, side, do_shuffle, batchsize):
        """ A generator function that yields the augmented, target and sample images.
        see :func:`minibatch_ab` for more details on the output. """
        logger.debug("Loading minibatch generator: (image_count: %s, side: '%s', do_shuffle: %s)",
                     len(images), side, do_shuffle)
        self._validate_samples(images) </s> remove             batch = list()
            for _ in range(batchsize):
                img_path = next(img_iter)
                data = self.process_face(img_path, side, is_display)
                batch.append(data)
            batch = list(zip(*batch))
            batch = [np.array(x, dtype="float32") for x in batch]
            logger.trace("Yielding batch: (size: %s, item shapes: %s, side:  '%s', "
                         "is_display: %s)",
                         len(batch), [item.shape for item in batch], side, is_display)
            yield batch

        logger.debug("Finished minibatch generator: (side: '%s', is_display: %s)",
                     side, is_display)

    def process_face(self, filename, side, is_display):
        """ Load an image and perform transformation and warping """
        logger.trace("Process face: (filename: '%s', side: '%s', is_display: %s)",
                     filename, side, is_display)
        image = cv2_read_img(filename, raise_error=True)
        if self.mask_class or self.training_opts["warp_to_landmarks"]:
            src_pts = self.get_landmarks(filename, image, side)
        if self.mask_class:
            image = self.mask_class(src_pts, image, channels=4).mask

        image = self.processing.color_adjust(image,
                                             self.training_opts["augment_color"],
                                             is_display)
        if not is_display:
            image = self.processing.random_transform(image)
            if not self.training_opts["no_flip"]:
                image = self.processing.do_random_flip(image)
        sample = image.copy()[:, :, :3]

        if self.training_opts["warp_to_landmarks"]:
            dst_pts = self.get_closest_match(filename, side, src_pts)
            processed = self.processing.random_warp_landmarks(image, src_pts, dst_pts)
 </s> add             img_paths = [next(img_iter) for _ in range(batchsize)]
            yield self._process_batch(img_paths, side)

        logger.debug("Finished minibatch generator: (side: '%s')", side)

    def _process_batch(self, filenames, side):
        """ Performs the augmentation and compiles target images and samples. See
        :func:`minibatch_ab` for more details on the output. """
        logger.trace("Process batch: (filenames: '%s', side: '%s')", filenames, side)
        batch = read_image_batch(filenames)
        processed = dict()
        to_landmarks = self._training_opts["warp_to_landmarks"]

        # Initialize processing training size on first image
        if not self._processing.initialized:
            self._processing.initialize(batch.shape[1])

        # Get Landmarks prior to manipulating the image
        if self._mask_class or to_landmarks:
            batch_src_pts = self._get_landmarks(filenames, batch, side)

        # Color augmentation before mask is added
        if self._training_opts["augment_color"]:
            batch = self._processing.color_adjust(batch)

        # Add mask to batch prior to transforms and warps
        if self._mask_class:
            batch = np.array([self._mask_class(src_pts, image, channels=4).mask
                              for src_pts, image in zip(batch_src_pts, batch)])

        # Random Transform and flip
        batch = self._processing.transform(batch)
        if not self._training_opts["no_flip"]:
            batch = self._processing.random_flip(batch)

        # Add samples to output if this is for display
        if self._processing.is_display:
            processed["samples"] = batch[..., :3].astype("float32") / 255.0

        # Get Targets
        processed.update(self._processing.get_targets(batch))

        # Random Warp
        if to_landmarks:
            warp_kwargs = dict(batch_src_points=batch_src_pts,
                               batch_dst_points=self._get_closest_match(filenames,
                                                                        side,
                                                                        batch_src_pts)) </s> remove         dst_points = landmarks[choice(closest_hashes)]
        logger.trace("Returning: (dst_points: %s)", dst_points)
        return dst_points
 </s> add             batch_closest_hashes.append(closest_hashes)
        logger.trace("Cached closest hashes")
        return batch_closest_hashes


class ImageAugmentation():
    """ Performs augmentation on batches of training images.

    Parameters
    ----------
    batchsize: int
        The number of images that will be fed through the augmentation functions at once.
    is_display: bool
        Whether the images being fed through will be used for Preview or Timelapse. Disables
        the "warp" augmentation for these images.
    input_size: int
        The expected input size for the model. It is assumed that the input to the model is always
        a square image. This is the size, in pixels, of the `width` and the `height` of the input
        to the model.
    output_shapes: list
        A list of tuples defining the output shapes from the model, in the order that the outputs
        are returned. The tuples should be in (`height`, `width`, `channels`) format.
    coverage_ratio: float
        The ratio of the training image to be trained on. Dictates how much of the image will be
        cropped out. Eg: a coverage ratio of 0.625 will result in cropping a 160px box from a 256px
        image (256 * 0.625 = 160).
    config: dict
        The configuration ``dict`` generated from :file:`config.train.ini` containing the trainer \
        plugin configuration options.

    Attributes
    ----------
    initialized: bool
        Flag to indicate whether :class:`ImageAugmentation` has been initialized with the training
        image size in order to cache certain augmentation operations (see :func:`initialize`)
    is_display: bool
        Flag to indicate whether these augmentations are for timelapses/preview images (``True``)
        or standard training data (``False)``
    """
    def __init__(self, batchsize, is_display, input_size, output_shapes, coverage_ratio, config):
        logger.debug("Initializing %s: (batchsize: %s, is_display: %s, input_size: %s, "
                     "output_shapes: %s, coverage_ratio: %s, config: %s)",
                     self.__class__.__name__, batchsize, is_display, input_size, output_shapes,
                     coverage_ratio, config)
        """ A Background iterator to return augmented images, samples and targets.

        The exit point from this class and the sole attribute that should be referenced. Called
        from :mod:`plugins.train.trainer._base`. Returns an iterator that yields images for
        training, preview and timelapses.

        Parameters
        ----------
        images: list
            A list of image paths that will be used to compile the final augmented data from.
        batchsize: int
            The batchsize for this iterator. Images will be returned in ``numpy.ndarray`` s of
            this size from the iterator.
        side: {'a' or 'b'}
            The side of the model that this iterator is for.
        do_shuffle: bool, optional
            Whether data should be shuffled prior to loading from disk. If true, each time the full
            list of filenames are processed, the data will be reshuffled to make sure thay are not
            returned in the same order. Default: ``True``
        is_preview: bool, optional
            Indicates whether this iterator is generating preview images. If ``True`` then certain
            augmentations will not be performed. Default: ``False``
        is_timelapse: bool optional
            Indicates whether this iterator is generating Timelapse images. If ``True``, then
            certain augmentations will not be performed. Default: ``False``

        Yields
        ------
        dict
            The following items are contained in each ``dict`` yielded from this iterator:

            * **feed** (`numpy.ndarray`) - The feed for the model. The array returned is in the \
            format (`batchsize`, `height`, `width`, `channels`). This is the :attr:`x` parameter \
            for :func:`keras.models.model.train_on_batch`.

            * **targets** (`list`) - A list of 4-dimensional ``numpy.ndarray`` s in the order \
            and size of each output of the model as defined in :attr:`model_output_shapes`. the \
            format of these arrays will be (`batchsize`, `height`, `width`, `3`). This is \
            the :attr:`y` parameter for :func:`keras.models.model.train_on_batch` **NB:** \
            masks are not included in the ``targets`` list. If required for feeding into the \
            Keras model, they will need to be added to this list in \
            :mod:`plugins.train.trainer._base` from the ``masks`` key.

            * **masks** (`numpy.ndarray`) - A 4-dimensional array containing the target masks in \
            the format (`batchsize`, `height`, `width`, `1`). **NB:** This item will only exist \
            in the ``dict`` if the :attr:`mask_type` is not ``None``

            * **samples** (`numpy.ndarray`) - A 4-dimensional array containg the samples for \
            feeding to the model's predict function for generating preview and timelapse samples. \
            The array will be in the format (`batchsize`, `height`, `width`, `channels`). **NB:** \
            This item will only exist in the ``dict`` if :attr:`is_preview` or \
            :attr:`is_timelapse` is ``True``
        """
          return mask_class
  
      def minibatch_ab(self, images, batchsize, side,
                       do_shuffle=True, is_preview=False, is_timelapse=False):
        """ A Background iterator to return augmented images, samples and targets.

        The exit point from this class and the sole attribute that should be referenced. Called
        from :mod:`plugins.train.trainer._base`. Returns an iterator that yields images for
        training, preview and timelapses.

        Parameters
        ----------
        images: list
            A list of image paths that will be used to compile the final augmented data from.
        batchsize: int
            The batchsize for this iterator. Images will be returned in ``numpy.ndarray`` s of
            this size from the iterator.
        side: {'a' or 'b'}
            The side of the model that this iterator is for.
        do_shuffle: bool, optional
            Whether data should be shuffled prior to loading from disk. If true, each time the full
            list of filenames are processed, the data will be reshuffled to make sure thay are not
            returned in the same order. Default: ``True``
        is_preview: bool, optional
            Indicates whether this iterator is generating preview images. If ``True`` then certain
            augmentations will not be performed. Default: ``False``
        is_timelapse: bool optional
            Indicates whether this iterator is generating Timelapse images. If ``True``, then
            certain augmentations will not be performed. Default: ``False``

        Yields
        ------
        dict
            The following items are contained in each ``dict`` yielded from this iterator:

            * **feed** (`numpy.ndarray`) - The feed for the model. The array returned is in the \
            format (`batchsize`, `height`, `width`, `channels`). This is the :attr:`x` parameter \
            for :func:`keras.models.model.train_on_batch`.

            * **targets** (`list`) - A list of 4-dimensional ``numpy.ndarray`` s in the order \
            and size of each output of the model as defined in :attr:`model_output_shapes`. the \
            format of these arrays will be (`batchsize`, `height`, `width`, `3`). This is \
            the :attr:`y` parameter for :func:`keras.models.model.train_on_batch` **NB:** \
            masks are not included in the ``targets`` list. If required for feeding into the \
            Keras model, they will need to be added to this list in \
            :mod:`plugins.train.trainer._base` from the ``masks`` key.

            * **masks** (`numpy.ndarray`) - A 4-dimensional array containing the target masks in \
            the format (`batchsize`, `height`, `width`, `1`). **NB:** This item will only exist \
            in the ``dict`` if the :attr:`mask_type` is not ``None``

            * **samples** (`numpy.ndarray`) - A 4-dimensional array containg the samples for \
            feeding to the model's predict function for generating preview and timelapse samples. \
            The array will be in the format (`batchsize`, `height`, `width`, `channels`). **NB:** \
            This item will only exist in the ``dict`` if :attr:`is_preview` or \
            :attr:`is_timelapse` is ``True``
        """
          logger.debug("Queue batches: (image_count: %s, batchsize: %s, side: '%s', do_shuffle: %s, "
                       "is_preview, %s, is_timelapse: %s)", len(images), batchsize, side, do_shuffle,
                       is_preview, is_timelapse)
          self.batchsize = batchsize
          is_display = is_preview or is_timelapse
keep keep keep keep 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