augmentation 0.2

A library, based on PyTorch, that performs data augmentation on the GPU

GPU accelerated data augmentation

This code is an extension on the imgaug package that provides flexible use of data augmentation. In particular, we provide standard augmentation strategies on the GPU, as some of these can be intensive on the CPU. Our GPU translation is based on PyTorch. The current version supports both 2D and 3D data augmentation, however the 3D augmentation is still in test phase.

Installation

Install the dependencies and you are ready to go!

pip install -r requirements.txt

Usage

Import the required modules:

import imageio
import torch
from augmentation_2d import FlipX, FlipY, Rotate90, AddNoise, RandomDeformation

Read an image into a tensor and transfer it to the GPU:

x = imageio.imread('img/elaine.png').astype('float')
shape = x.shape
x = torch.Tensor(x).unsqueeze(0).unsqueeze(0).cuda()

Flip the image along the x and y-axis:

flipx = FlipX(shape)
y = flipx(x)
flipy = FlipY(shape)
y = flipy(x)

Rotate the image 90 degrees:

rotate = Rotate90(shape)
y = rotate(x)

Rotate the image randomly:

rotate = RotateRandom(shape)
y = rotate(x)

Add noise to the image:

noise = AddNoise(sigma_min=20, sigma_max=20)
y = noise(x)

Apply random deformations to the image:

deform = RandomDeformation(shape, sigma=0.01)
y = deform(x)