Last active
May 3, 2020 11:20
-
-
Save Mistobaan/bccb5823d16537bed2e2823dff1a2155 to your computer and use it in GitHub Desktop.
Tensorflow Image Augmentation Tips and Tricks
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# from https://www.kaggle.com/cdeotte/cutmix-and-mixup-on-gpu-tpu | |
## CutMix Augmentation¶ | |
# The following code does cutmix using the GPU/TPU. | |
# Change the variables SWITCH, CUTMIX_PROB and MIXUP_PROB | |
# in function transform() to control the amount of augmentation during training. | |
# CutMix will occur SWITCH * CUTMIX_PROB often and | |
# MixUp will occur (1-SWITCH) * MIXUP_PROB often during training. | |
def onehot(image,label): | |
CLASSES = 104 | |
return image,tf.one_hot(label,CLASSES) | |
def cutmix(image, label, PROBABILITY = 1.0): | |
# input image - is a batch of images of size [n,dim,dim,3] not a single image of [dim,dim,3] | |
# output - a batch of images with cutmix applied | |
DIM = IMAGE_SIZE[0] | |
CLASSES = 104 | |
imgs = []; labs = [] | |
for j in range(AUG_BATCH): | |
# DO CUTMIX WITH PROBABILITY DEFINED ABOVE | |
P = tf.cast( tf.random.uniform([],0,1)<=PROBABILITY, tf.int32) | |
# CHOOSE RANDOM IMAGE TO CUTMIX WITH | |
k = tf.cast( tf.random.uniform([],0,AUG_BATCH),tf.int32) | |
# CHOOSE RANDOM LOCATION | |
x = tf.cast( tf.random.uniform([],0,DIM),tf.int32) | |
y = tf.cast( tf.random.uniform([],0,DIM),tf.int32) | |
b = tf.random.uniform([],0,1) # this is beta dist with alpha=1.0 | |
WIDTH = tf.cast( DIM * tf.math.sqrt(1-b),tf.int32) * P | |
ya = tf.math.maximum(0,y-WIDTH//2) | |
yb = tf.math.minimum(DIM,y+WIDTH//2) | |
xa = tf.math.maximum(0,x-WIDTH//2) | |
xb = tf.math.minimum(DIM,x+WIDTH//2) | |
# MAKE CUTMIX IMAGE | |
one = image[j,ya:yb,0:xa,:] | |
two = image[k,ya:yb,xa:xb,:] | |
three = image[j,ya:yb,xb:DIM,:] | |
middle = tf.concat([one,two,three],axis=1) | |
img = tf.concat([image[j,0:ya,:,:],middle,image[j,yb:DIM,:,:]],axis=0) | |
imgs.append(img) | |
# MAKE CUTMIX LABEL | |
a = tf.cast(WIDTH*WIDTH/DIM/DIM,tf.float32) | |
if len(label.shape)==1: | |
lab1 = tf.one_hot(label[j],CLASSES) | |
lab2 = tf.one_hot(label[k],CLASSES) | |
else: | |
lab1 = label[j,] | |
lab2 = label[k,] | |
labs.append((1-a)*lab1 + a*lab2) | |
# RESHAPE HACK SO TPU COMPILER KNOWS SHAPE OF OUTPUT TENSOR (maybe use Python typing instead?) | |
image2 = tf.reshape(tf.stack(imgs),(AUG_BATCH,DIM,DIM,3)) | |
label2 = tf.reshape(tf.stack(labs),(AUG_BATCH,CLASSES)) | |
return image2,label2 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment