Class AugmentorList

tensorpack/dataflow/imgaug/base.py:219–263

Augment an image by a list of augmentors

class AugmentorList(ImageAugmentor):
    """
    Augment an image by a list of augmentors
    """

    def __init__(self, augmentors):
        """
        Args:
            augmentors (list): list of :class:`ImageAugmentor` instance to be applied.
        """
        assert isinstance(augmentors, (list, tuple)), augmentors
        self.augmentors = augmentors
        super(AugmentorList, self).__init__()

    def reset_state(self):
        """ Will reset state of each augmentor """
        super(AugmentorList, self).reset_state()
        for a in self.augmentors:
            a.reset_state()

    def get_transform(self, img):
        check_dtype(img)
        assert img.ndim in [2, 3], img.ndim

        from .transform import LazyTransform, TransformList
        # The next augmentor requires the previous one to finish.
        # So we have to use LazyTransform
        tfms = []
        for idx, a in enumerate(self.augmentors):
            if idx == 0:
                t = a.get_transform(img)
            else:
                t = LazyTransform(a.get_transform)

            if isinstance(t, TransformList):
                tfms.extend(t.tfms)
            else:
                tfms.append(t)
        return TransformList(tfms)

    def __str__(self):
        repr_each_aug = ",\n".join(["  " + repr(x) for x in self.augmentors])
        return "imgaug.AugmentorList([\n{}])".format(repr_each_aug)

    __repr__ = __str__


Augmentor = ImageAugmentor