def augment(im, params=None):
    """
    Perform data augmentation on some image using the albumentations package.

    Parameters
    ----------
    im : Numpy array
    params : dict or None
        Contains the data augmentation parameters
        Mandatory keys:
        - h_flip ([0,1] float): probability of performing an horizontal left-right mirroring.
        - v_flip ([0,1] float): probability of performing an vertical up-down mirroring.
        - rot ([0,1] float):  probability of performing a rotation to the image.
        - rot_lim (int):  max degrees of rotation.
        - stretch ([0,1] float):  probability of randomly stretching an image.
        - crop ([0,1] float): randomly take an image crop.
        - zoom ([0,1] float): random zoom applied to crop_size.
            --> Therefore the effective crop size at each iteration will be a
                random number between 1 and crop*(1-zoom). For example:
                  * crop=1, zoom=0: no crop of the image
                  * crop=1, zoom=0.1: random crop of random size between 100% image and 90% of the image
                  * crop=0.9, zoom=0.1: random crop of random size between 90% image and 80% of the image
                  * crop=0.9, zoom=0: random crop of always 90% of the image
                  Image size refers to the size of the shortest side.
        - blur ([0,1] float):  probability of randomly blurring an image.
        - pixel_noise ([0,1] float):  probability of randomly adding pixel noise to an image.
        - pixel_sat ([0,1] float):  probability of randomly using HueSaturationValue in the image.
        - cutout ([0,1] float):  probability of using cutout in the image.

    Returns
    -------
    Numpy array
    """

    ## 1) Crop the image
    effective_zoom = np.random.rand() * params['zoom']
    crop = params['crop'] - effective_zoom

    ly, lx, channels = im.shape
    crop_size = int(crop * min([ly, lx]))
    rand_x = np.random.randint(low=0, high=lx - crop_size + 1)
    rand_y = np.random.randint(low=0, high=ly - crop_size + 1)

    crop = transforms.Crop(x_min=rand_x,
                           y_min=rand_y,
                           x_max=rand_x + crop_size,
                           y_max=rand_y + crop_size)

    im = crop(image=im)['image']

    ## 2) Now add the transformations for augmenting the image pixels
    transform_list = []

    # Add random stretching
    if params['stretch']:
        transform_list.append(
            imgaug_transforms.IAAPerspective(scale=0.1, p=params['stretch']))

    # Add random rotation
    if params['rot']:
        transform_list.append(
            transforms.Rotate(limit=params['rot_lim'], p=params['rot']))

    # Add horizontal flip
    if params['h_flip']:
        transform_list.append(transforms.HorizontalFlip(p=params['h_flip']))

    # Add vertical flip
    if params['v_flip']:
        transform_list.append(transforms.VerticalFlip(p=params['v_flip']))

    # Add some blur to the image
    if params['blur']:
        transform_list.append(
            albumentations.OneOf([
                transforms.MotionBlur(blur_limit=7, p=1.),
                transforms.MedianBlur(blur_limit=7, p=1.),
                transforms.Blur(blur_limit=7, p=1.),
            ],
                                 p=params['blur']))

    # Add pixel noise
    if params['pixel_noise']:
        transform_list.append(
            albumentations.OneOf(
                [
                    transforms.CLAHE(clip_limit=2, p=1.),
                    imgaug_transforms.IAASharpen(p=1.),
                    imgaug_transforms.IAAEmboss(p=1.),
                    transforms.RandomBrightnessContrast(contrast_limit=0,
                                                        p=1.),
                    transforms.RandomBrightnessContrast(brightness_limit=0,
                                                        p=1.),
                    transforms.RGBShift(p=1.),
                    transforms.RandomGamma(p=1.)  #,
                    # transforms.JpegCompression(),
                    # transforms.ChannelShuffle(),
                    # transforms.ToGray()
                ],
                p=params['pixel_noise']))

    # Add pixel saturation
    if params['pixel_sat']:
        transform_list.append(
            transforms.HueSaturationValue(p=params['pixel_sat']))

    # Remove randomly remove some regions from the image
    if params['cutout']:
        ly, lx, channels = im.shape
        scale_low, scale_high = 0.05, 0.25  # min and max size of the squares wrt the full image
        scale = np.random.uniform(scale_low, scale_high)
        transform_list.append(
            transforms.Cutout(num_holes=8,
                              max_h_size=int(scale * ly),
                              max_w_size=int(scale * lx),
                              p=params['cutout']))

    # Compose all image transformations and augment the image
    augmentation_fn = albumentations.Compose(transform_list)
    im = augmentation_fn(image=im)['image']

    return im
Esempio n. 2
0
def augment(im, params=None):
    """
    Perform data augmentation on some image using the albumentations package.

    Parameters
    ----------
    im : Numpy array
    params : dict or None
        Contains the data augmentation parameters
        Mandatory keys:
        - h_flip ([0,1] float): probability of performing an horizontal left-right mirroring.
        - v_flip ([0,1] float): probability of performing an vertical up-down mirroring.
        - rot ([0,1] float):  probability of performing a rotation to the image.
        - rot_lim (int):  max degrees of rotation.
        - stretch ([0,1] float):  probability of randomly stretching an image.
        - expand ([True, False] bool): whether to pad the image to a square shape with background color canvas.
        - crop ([0,1] float): randomly take an image crop.
        - invert_col ([0, 1] float): randomly invert the colors of the image. p=1 -> invert colors (VPR)
        - zoom ([0,1] float): random zoom applied to crop_size.
            --> Therefore the effective crop size at each iteration will be a
                random number between 1 and crop*(1-zoom). For example:
                  * crop=1, zoom=0: no crop of the image
                  * crop=1, zoom=0.1: random crop of random size between 100% image and 90% of the image
                  * crop=0.9, zoom=0.1: random crop of random size between 90% image and 80% of the image
                  * crop=0.9, zoom=0: random crop of always 90% of the image
                  Image size refers to the size of the shortest side.
        - blur ([0,1] float):  probability of randomly blurring an image.
        - pixel_noise ([0,1] float):  probability of randomly adding pixel noise to an image.
        - pixel_sat ([0,1] float):  probability of randomly using HueSaturationValue in the image.
        - cutout ([0,1] float):  probability of using cutout in the image.

    Returns
    -------
    Numpy array
    """
    ## 1) Expand the image by padding it with bg-color canvas
    if params["expand"]:
        desired_size = max(im.shape)
        # check bg
        if np.argmax(im.shape) > 0:
            bgcol = tuple(np.repeat(int(np.mean(im[[0, -1], :, :])), 3))
        else:
            bgcol = tuple(np.repeat(int(np.mean(im[:, [0, -1], :])), 3))

        im = Image.fromarray(im)
        old_size = im.size  # old_size[0] is in (width, height) format

        ratio = float(desired_size) / max(old_size)
        new_size = tuple([int(x * ratio) for x in old_size])
        im = im.resize(new_size, Image.ANTIALIAS)
        # create a new image and paste the resized on it
        new_im = Image.new("RGB", (desired_size, desired_size), color=bgcol)
        new_im.paste(im, ((desired_size - new_size[0]) // 2,
                          (desired_size - new_size[1]) // 2))

        im = np.array(new_im)

    ## 2) Crop the image
    if params["crop"] and params["crop"] != 1:
        effective_zoom = np.random.rand() * params['zoom']
        crop = params['crop'] - effective_zoom

        ly, lx, channels = im.shape
        crop_size = int(crop * min([ly, lx]))
        rand_x = np.random.randint(low=0, high=lx - crop_size + 1)
        rand_y = np.random.randint(low=0, high=ly - crop_size + 1)

        crop = transforms.Crop(x_min=rand_x,
                               y_min=rand_y,
                               x_max=rand_x + crop_size,
                               y_max=rand_y + crop_size)

        im = crop(image=im)['image']

    if params["enhance"]:
        im = Image.fromarray(im)
        enhancer = ImageEnhance.Contrast(im)
        im = np.array(enhancer.enhance(params["enhance"]))

    ## 3) Now add the transformations for augmenting the image pixels
    transform_list = []

    if params['invert_col']:
        transform_list.append(transforms.InvertImg(p=params['invert_col']))

    # Add random stretching
    if params['stretch']:
        transform_list.append(
            imgaug_transforms.IAAPerspective(scale=0.1, p=params['stretch']))

    # Add random rotation
    if params['rot']:
        transform_list.append(
            transforms.Rotate(limit=params['rot_lim'], p=params['rot']))

    # Add horizontal flip
    if params['h_flip']:
        transform_list.append(transforms.HorizontalFlip(p=params['h_flip']))

    # Add vertical flip
    if params['v_flip']:
        transform_list.append(transforms.VerticalFlip(p=params['v_flip']))

    # Add some blur to the image
    if params['blur']:
        transform_list.append(
            albumentations.OneOf([
                transforms.MotionBlur(blur_limit=7, p=1.),
                transforms.MedianBlur(blur_limit=7, p=1.),
                transforms.Blur(blur_limit=7, p=1.),
            ],
                                 p=params['blur']))

    # Add pixel noise
    if params['pixel_noise']:
        transform_list.append(
            albumentations.OneOf(
                [
                    transforms.CLAHE(clip_limit=2, p=1.),
                    imgaug_transforms.IAASharpen(p=1.),
                    imgaug_transforms.IAAEmboss(p=1.),
                    transforms.RandomBrightnessContrast(contrast_limit=0,
                                                        p=1.),
                    transforms.RandomBrightnessContrast(brightness_limit=0,
                                                        p=1.),
                    transforms.RGBShift(p=1.),
                    transforms.RandomGamma(p=1.)  #,
                    # transforms.JpegCompression(),
                    # transforms.ChannelShuffle(),
                    # transforms.ToGray()
                ],
                p=params['pixel_noise']))

    # Add pixel saturation
    if params['pixel_sat']:
        transform_list.append(
            transforms.HueSaturationValue(p=params['pixel_sat']))

    # Remove randomly remove some regions from the image
    if params['cutout']:
        ly, lx, channels = im.shape
        scale_low, scale_high = 0.05, 0.25  # min and max size of the squares wrt the full image
        scale = np.random.uniform(scale_low, scale_high)
        transform_list.append(
            transforms.Cutout(num_holes=8,
                              max_h_size=int(scale * ly),
                              max_w_size=int(scale * lx),
                              p=params['cutout']))

    # Compose all image transformations and augment the image
    augmentation_fn = albumentations.Compose(transform_list)
    im = augmentation_fn(image=im)['image']

    return im