Skip to content

An image augmentation library for tensorflow. It can be used seamlessly with tf.data.Dataset

License

Notifications You must be signed in to change notification settings

looooongChen/tfAugmentor

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

40 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

tfAugmentor

An image augmentation library for tensorflow. The libray is designed to be easily used with tf.data.Dataset. The augmentor accepts tf.data.Dataset object or a nested tuple of numpy array.

Augmentations

Original Flip Rotation Translation
original demo_flip demo_rotation demo_translation
Crop Elactic Deform
demo_crop demo_elastic
Gaussian Blur Contrast Gamma
demo_blur demo_contrast demo_gamma

Demo

Random Rotation Random Translation Random Crop
demo_ratation demo_translation demo_crop
Random Contrast Random Gamma Elastic Deform
demo_contrast demo_gamma demo_elastic

Installation

tfAugmentor is written in Python and can be easily installed via:

pip install tfAugmentor

Required packages:

  • tensorflow (developed under tf 2.4), should work with any 2.x version
  • numpy (numpy=1.20 may cause error of tf.meshgrid, use another version)

Quick Start

tfAugmentor is implemented to work seamlessly with tf.data. The tf.data.Dataset object can be directly processed by tfAugmentor.

To instantiate an Augmentor object, three arguments are required:

class Augmentor(object):
    def __init__(self, signature, image=[], label_map=[]):
		...
  • signature: to present the structure of the dataset
    • a nested tuple of string
    • a list/tuple of dictionary keys, if your dataset is in a dictionary form
  • image: list of string items in signature, which will be treated as normal images
  • label: list of string items in signature, which will be treated as segmentation masks

Note: only the items in 'image' and 'label' will be processed, others will remain untouched

A simple example

import tfAugmentor as tfaug

# new tfAugmentor object
aug = tfaug.Augmentor(signature=('image', ('mask1', 'mask2')), 
                      image=['image'], 
                      label=['mask1', 'mask2'])

# add augumentation operations
aug.flip_left_right(probability=0.5)
aug.rotate90(probability=0.5)
aug.elastic_deform(strength=2, scale=20, probability=1)

# assume we have three numpy arrays
X_image = ... # shape [batch, height, width, channel]
Y_mask1 = ... # shape [batch, height, width, 1]
Y_mask2 = ... # shape [batch, height, width, 1]

# create tf.data.Dataset object
tf_dataset = tf.data.Dataset.from_tensor_slices((X_image, (Y_mask1, Y_mask2))))
# do the actual augmentation
ds1 = aug(tf_dataset)

# or you can directly pass the numpy arrays, a tf.data.Dataset object will be returned 
ds2 = aug((X_image, (Y_mask1, Y_mask2))), keep_size=True)

If the data is passed as a python dictionary, the signature should be the list/tuple of keys. For example:

import tfAugmentor as tfaug

# new tfAugmentor object
aug = tfaug.Augmentor(signature=('image', 'mask1', 'mask2'), 
                      image=['image'], 
                      label=['mask1', 'mask2'])

# add augumentation operations
aug.flip_left_right(probability=0.5)
aug.rotate90(probability=0.5)
aug.elastic_deform(strength=2, scale=20, probability=1)

# assume we have three numpy arrays
X_image = ... # shape [batch, height, width, channel]
Y_mask1 = ... # shape [batch, height, width, 1]
Y_mask2 = ... # shape [batch, height, width, 1]

ds_dict = {'image': X_image,
           'mask1': Y_mask1,
           'mask2': Y_mask2}
# create tf.data.Dataset object
tf_dataset = tf.data.Dataset.from_tensor_slices(ds_dict)
# do the actual augmentation
ds1 = aug(tf_dataset)

# or directly pass the data
ds2 = aug(ds_dict)

Note: All added operations will be executed one by one, but you can create multiply tfAugmentor to realize parallel pipelines

A more complicated example

import tfAugmentor as tfaug

# since 'class' is neither in 'image' nor in 'label', it will not be processed 
aug1 = tfaug.Augmentor((('image_rgb', 'image_depth'), ('semantic_mask', 'class')), 
                       image=['image_rgb', 'image_depth'], 
                       label=['semantic_mask'])
aug2 = tfaug.Augmentor((('image_rgb', 'image_depth'), ('semantic_mask', 'class')), 
                       image=['image_rgb', 'image_depth'], 
                       label=['semantic_mask'])

# add different augumentation operations to aug1 and aug2 
aug1.flip_left_right(probability=0.5)
aug1.random_crop_resize(sacle_range=(0.7, 0.9), probability=0.5)
aug2.elastic_deform(strength=2, scale=20, probability=1)

# assume we have the 1000 data samples
X_rgb = ...  # shape [1000 x 512 x 512 x 3]
X_depth = ... # shape [1000 x 512 x 512 x 1]
Y_semantic_mask = ... # shape [1000 x 512 x 512 x 1]
Y_class = ... # shape [1000 x 1]

# create tf.data.Dataset object
ds_origin = tf.data.Dataset.from_tensor_slices(((X_rgb, X_depth), (Y_semantic_mask, Y_class))))
# do the actual augmentation
ds1 = aug1(ds_origin)
ds2 = aug2(ds_origin)
# combine them
ds = ds_origin.concatenate(ds1)
ds = ds.concatenate(ds1)

Main Features

The argument 'probability' controls the possibility of a certain augmentation taking place.

Mirroring

# flip the image left right  
aug.flip_left_right(probability=1)
# flip the image up down 
aug.flip_up_down(probability=1) 

Rotation

# rotate by 90 degree clockwise
a.rotate90(probability=1) 
# rotate by 180 degree clockwise
a.rotate180(probability=1)
# rotate by 270 degree clockwise 
a.rotate270(probability=1) 
# rotate by a certrain degree, Args: angle - scala, in degree
a.rotate(angle, probability=1) 
# randomly rotate the image
a.random_rotate(probability=1) 

Translation

# tranlate image, Args: offset - [x, y]
a.translate(offset, probability=1):
# randoms translate image 
a.random_translate(translation_range=[-100, 100], probability=1):

Crop and Resize

# randomly crop a sub-image and resize to the original image size
a.random_crop(scale_range=([0.5, 0.8], preserve_aspect_ratio=False, probability=1) 

Elastic Deformation

# performa elastic deformation
a.elastic_deform(scale=10, strength=200, probability=1)

Photometric Adjustment

# adjust image contrast randomly
a.random_contrast(contrast_range=[0.6, 1.4], probability=1)
# perform gamma correction with random gamma values
a.random_gamma(gamma_range=[0.5, 1.5], probability=1)

Noise

# blur the image with gaussian kernel
a.gaussian_blur(sigma=2, probability=1)

Caution

  • If .batch() of tf.data.Dataset is used before augmentation, please set drop_remainder=True. Oherwise, the batch_size will be set to None. The augmention of tfAgmentor requires the batch_size dimension

About

An image augmentation library for tensorflow. It can be used seamlessly with tf.data.Dataset

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages