extending the augmentations

[DiddyC]

Is it possible to add my own function to use during training? (similarly to: https://github.com/keras-team/keras/issues/3338) I didn't see any answers about this, would appreciate if someone can point me to right direction!

[aleju]

You can use lambda augmenters, e.g.

def func_images(images, random_state, parents, hooks):
    images[:, ::2, :, :] = 0
    return images

def func_keypoints(keypoints_on_images, random_state, parents, hooks):
    return keypoints_on_images

aug = iaa.Lambda(
    func_images=func_images,
    func_keypoints=func_keypoints
)

Alternatively you can also just write your own augmenter, which isn't exactly hard. An easy example is Fliplr for horizontal flips, to be found in imgaug/augmenters/flip.py:

from imgaug.augmenters import Augmenter
from imgaug.parameters import Binomial, StochasticParameter
import numpy as np

class Fliplr(Augmenter):
    def __init__(self, p=0, name=None, deterministic=False, random_state=None):
        super(Fliplr, self).__init__(name=name, deterministic=deterministic, random_state=random_state)

        if ia.is_single_number(p):
            self.p = Binomial(p)
        elif isinstance(p, StochasticParameter):
            self.p = p
        else:
            raise Exception("Expected p to be int or float or StochasticParameter, got %s." % (type(p),))

    def _augment_images(self, images, random_state, parents, hooks):
        nb_images = len(images)
        samples = self.p.draw_samples((nb_images,), random_state=random_state)
        for i in range(nb_images):
            if samples[i] == 1:
                images[i] = np.fliplr(images[i])
        return images

    def _augment_keypoints(self, keypoints_on_images, random_state, parents, hooks):
        nb_images = len(keypoints_on_images)
        samples = self.p.draw_samples((nb_images,), random_state=random_state)
        for i, keypoints_on_image in enumerate(keypoints_on_images):
            if samples[i] == 1:
                width = keypoints_on_image.shape[1]
                for keypoint in keypoints_on_image.keypoints:
                    keypoint.x = (width - 1) - keypoint.x
        return keypoints_on_images

    def get_parameters(self):
        return [self.p]

if you don't need keypoints or your augmenter won't affect them, you can replace _augment_keypoints(...) simply with return keypoints_on_images.

[vic85821]

Hi, I have another question. How could I get the corresponding masks in self-defined lambda function. Because I need to do another process for the masks. For example, in the direction domain mask presents the x, y direction vectors. So if I perform Flipud, it's needed to multiply -1 in y direction vector.

Thanks!

[aleju]

You can call to_deterministic() on the sequence and then augment independently (but with the same parameters) images via augment_images() and masks via augment_segmentation_maps(ia.SegmentationMapOnImage) or augment_heatmaps(ia.HeatmapsOnImage). Though it sounds like your masks are more like maps of vectors, which are currently not supported by the library. Best "solution" is then to remove all augmenters that change the geometry so that you can keep the masks unchanged.