Keras implementation of layer which performs augmentations of images using GPU.
Keras implementation of layer which performs augmentations of images using GPU. It can be a good and fast replacement for ImageDataGenerator or independent augmentation libraries like imgaug, albumentations etc.
Important Note: It's prototype version which I believe can be improved a lot in terms of speed and usability. I'd really like to see it in official Keras and TF repository. Feel free to add more transforms or parameters using pull requests.
Python 3.*, numpy, tensorflow>=2.5.0, tensorflow_addons>=0.14.0
pip install keras-augm-layer
You need to define initial model, then choose set of transforms and add AugmLayer as first layer of your model. See example below:
from keras.applications.mobilenet import MobileNet, preprocess_input
from keras.layers import Dense, Input
from keras.models import Model
from keras_augm_layer import *
def get_classification_model():
base_model = MobileNet((128, 128, 3), depth_multiplier=1, alpha=0.25,
include_top=False, pooling='avg', weights='imagenet')
x = base_model.layers[-1].output
x = Dense(1, use_bias=False)(x)
model = Model(inputs=base_model.inputs, outputs=x)
return model
transforms = [
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
RandomRotate(angle=45, p=0.5),
RandomRotate90(p=0.5),
RGBShift(p=0.1, r_shift_limit=20, g_shift_limit=20, b_shift_limit=20),
RandomBrightness(p=0.999, max_delta=0.1),
RandomContrast(0.5, 1.5, p=0.9),
RandomHue(0.5, p=0.9),
RandomSaturation(0.5, 1.5, p=0.9),
RandomGaussNoise((10, 50), p=0.99),
ToGray(p=0.5),
JpegCompression(p=0.9, quality_lower=5, quality_upper=99),
]
inp = Input((None, None, 3))
x = AugmLayer(transforms, output_dim=(128, 128, 3), preproc_input=preprocess_input)(inp)
base_model = get_classification_model()
x = base_model(x)
augm_model = Model(inputs=inp, outputs=x)
preprocess_input
function, because it must be applied to image before sent on model input.Before starting training process you most likely will want to check your set of transforms. See example how to debug below:
# Define set of transforms
transforms = [
HorizontalFlip(p=0.5),
VerticalFlip(p=0.5),
]
# Create Keras Model only with single AugmLayer (don't use preprocess_input and specify training=True)
inp = Input((None, None, 3))
x = AugmLayer(transforms, output_dim=(128, 128, 3))(inp, training=True)
augm_model = Model(inputs=inp, outputs=x)
# Read and predict images
files = glob.glob("../input/*.jpg")
images = []
for f in files:
images.append(cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB))
images = np.array(images)
# augm_images will contain augmented images
augm_images = augm_model.predict(images)
# Now show all augmented images
def show_image(im, name='image'):
cv2.imshow(name, cv2.cvtColor(im.astype(np.uint8), cv2.COLOR_RGB2BGR))
cv2.waitKey(0)
cv2.destroyAllWindows()
for i in range(augm_images.shape[0]):
show_image(augm_images[i])
You can use this script for debug: check_augmentations.py
AugmLayer has the same behaviour as Dropout layer. It's applied transformations only during training phase. So you can save model as is. But exists the better way. If you create augm_model as it was in Usage section you can just pull out initial classification model with construction below:
model = augm_model.layers[2]
model.save(...)
Set of augm | Albumentations (Threads 1) | Albumentations (Threads 4) | Albumentations (Threads 6) | AugmLayer |
---|---|---|---|---|
Flips + RandomRotate45 | 15s | 14s | 14s | 19s |
JpegCompression | 38s | 15s | 15s | 21s |
Colors + Noise | 147s | 68s | 62s | 32s |
Many Augm | 176s | 76s | 69s | 41s |
Some observations: GPU augmentations became much faster than CPU in following cases: