Conditional Augmenter

aug.py

import imgaug.augmenters as iaa
import numpy as np


class Conditional(iaa.Augmenter):
    """Apply child augmenter(s) for images wich a given conditional comes true.

    Let ``C`` be one or more child augmenters given to
    :class:`~imgaug.augmenters.meta.Conditional`.
    Let ``p`` be the fraction of images (or other data) to augment.
    Let ``I`` be the input images (or other data).
    Let ``N`` be the number of input images (or other entities).
    Then the images of ``I`` which pass a conditional will be augmented using ``C``.

    **Supported dtypes**:

        * ``uint8``: yes; fully tested
        * ``uint16``: yes; tested
        * ``uint32``: yes; tested
        * ``uint64``: yes; tested
        * ``int8``: yes; tested
        * ``int16``: yes; tested
        * ``int32``: yes; tested
        * ``int64``: yes; tested
        * ``float16``: yes; tested
        * ``float32``: yes; tested
        * ``float64``: yes; tested
        * ``float128``: yes; tested
        * ``bool``: yes; tested

    Parameters
    ----------
    func_images : callable
        The function to call for each batch of images.
        It must follow the form::

            function(images)
        and must return a boolean array with ``True``, if that image should be augmented
         or ``False`` otherwise.

    then_list : None or imgaug.augmenters.meta.Augmenter or list of imgaug.augmenters.meta.Augmenter, optional
        Augmenter(s) to apply to the matched images.
        If this is a list of augmenters, it will be converted to a
        :class:`~imgaug.augmenters.meta.Sequential`.

    else_list : None or imgaug.augmenters.meta.Augmenter or list of imgaug.augmenters.meta.Augmenter, optional
        Augmenter(s) to apply to the non matched images.
        These augmenters will be applied only when the ones in `then_list`
        are *not* applied (either-or-relationship).
        If this is a list of augmenters, it will be converted to a
        :class:`~imgaug.augmenters.meta.Sequential`.

    seed : None or int or imgaug.random.RNG or numpy.random.Generator or numpy.random.BitGenerator or numpy.random.SeedSequence or numpy.random.RandomState, optional
        See :func:`~imgaug.augmenters.meta.Augmenter.__init__`.

    name : None or str, optional
        See :func:`~imgaug.augmenters.meta.Augmenter.__init__`.

    random_state : None or int or imgaug.random.RNG or numpy.random.Generator or numpy.random.BitGenerator or numpy.random.SeedSequence or numpy.random.RandomState, optional
        Old name for parameter `seed`.
        Its usage will not yet cause a deprecation warning,
        but it is still recommended to use `seed` now.
        Outdated since 0.4.0.

    deterministic : bool, optional
        Deprecated since 0.4.0.
        See method ``to_deterministic()`` for an alternative and for
        details about what the "deterministic mode" actually does.

    Examples
    --------
    >>> import imgaug.augmenters as iaa
    >>> aug = iaa.Conditional(lambda images: np.mean(images, axis=tuple(range(1, images.ndim))) < 127, iaa.Multiply(2.0))

    """

    def __init__(self, func_images, then_list=None, else_list=None, seed=None, name: str = None,
                 random_state="deprecated", deterministic="deprecated"):
        super(Conditional, self).__init__(
            seed=seed, name=name,
            random_state=random_state, deterministic=deterministic)

        self.func_images = func_images
        self.then_list = iaa.handle_children_list(then_list, self.name, "then", default=None)
        self.else_list = iaa.handle_children_list(else_list, self.name, "else", default=None)

    # Added in 0.4.0.
    def _augment_batch_(self, batch, random_state, parents, hooks):
        with batch.propagation_hooks_ctx(self, hooks, parents):
            matches = self.func_images(batch.images)

            # create lists/arrays of images for if and else lists (one for each)
            # note that np.where returns tuple(array([0, 5, 9, ...])) or
            # tuple(array([]))
            indices_then_list = np.where(matches)[0]
            indices_else_list = np.where(~matches)[0]

            indice_lists = [indices_then_list, indices_else_list]
            augmenter_lists = [self.then_list, self.else_list]

            # For then_list: collect augmentables to be processed by then_list
            # augmenters, apply them to the list, then map back to the output
            # list. Analogous for else_list.
            # TODO maybe this would be easier if augment_*() accepted a list
            #      that can contain Nones
            for indices, augmenters in zip(indice_lists, augmenter_lists):
                if augmenters is not None and len(augmenters) > 0:
                    batch_sub = batch.subselect_rows_by_indices(indices)
                    batch_sub = augmenters.augment_batch_(
                        batch_sub,
                        parents=parents + [self],
                        hooks=hooks
                    )
                    batch = batch.invert_subselect_rows_by_indices_(indices,
                                                                    batch_sub)

            return batch

    def _to_deterministic(self):
        aug = self.copy()
        aug.then_list = (aug.then_list.to_deterministic()
                         if aug.then_list is not None
                         else aug.then_list)
        aug.else_list = (aug.else_list.to_deterministic()
                         if aug.else_list is not None
                         else aug.else_list)
        aug.deterministic = True
        aug.random_state = self.random_state.derive_rng_()
        return aug

    def get_parameters(self):
        """See :func:`~imgaug.augmenters.meta.Augmenter.get_parameters`."""
        return []

    def get_children_lists(self):
        """See :func:`~imgaug.augmenters.meta.Augmenter.get_children_lists`."""
        result = []
        if self.then_list is not None:
            result.append(self.then_list)
        if self.else_list is not None:
            result.append(self.else_list)
        return result

    def __str__(self):
        pattern = (
            "%s("
            "func_images=<lambda>, name=%s, then_list=%s, else_list=%s, deterministic=%s"
            ")")
        return pattern % (
            self.__class__.__name__, self.name, self.then_list,
            self.else_list, self.deterministic)