Commit f197c7ca authored by sjjsmuel's avatar sjjsmuel

shuffle fix and add augmentation

parent 7162de84
......@@ -10,21 +10,41 @@ import tensorflow as tf
Augmentation functions
'''
def flip(x, label, size):
x = tf.image.random_flip_left_right(x)
x = tf.image.random_flip_up_down(x)
return x, label
def color(x: tf.Tensor, label, size):
x = tf.image.random_hue(x, 0.1)
x = tf.image.random_saturation(x, 0.5, 1.7)
x = tf.image.random_brightness(x, 0.2)
x = tf.image.random_contrast(x, 0.5, 1.5)
return x, label
def rotate(x: tf.Tensor, label, size):
return tf.image.rot90(x, tf.random.uniform(shape=[], minval=0, maxval=4, dtype=tf.int32)), label
def zoom(x: tf.Tensor, label, size):
img = x['img']
mask = x['mouth']
mask = tf.expand_dims(mask, -1)
random_flr = tf.random.uniform(shape=[], minval=0, maxval=2, dtype=tf.int32)
img = tf.cond(random_flr == 0, lambda: img, lambda: tf.image.flip_left_right(img))
mask = tf.cond(random_flr == 0, lambda: mask, lambda: tf.image.flip_left_right(mask))
random_fud = tf.random.uniform(shape=[], minval=0, maxval=2, dtype=tf.int32)
img = tf.cond(random_fud == 0, lambda: img, lambda: tf.image.flip_up_down(img))
mask = tf.cond(random_fud == 0, lambda: mask, lambda: tf.image.flip_up_down(mask))
mask = tf.squeeze(mask)
return {'img': img, 'mouth': mask}, label
def color(x, label, size):
img = x['img']
mask = x['mouth']
img = tf.image.random_hue(img, 0.1)
img = tf.image.random_saturation(img, 0.5, 1.7)
img = tf.image.random_brightness(img, 0.2)
img = tf.image.random_contrast(img, 0.5, 1.5)
return {'img': img, 'mouth': mask}, label
def rotate(x, label, size):
img = x['img']
mask = x['mouth']
mask = tf.expand_dims(mask, -1)
random_value = tf.random.uniform(shape=[], minval=0, maxval=4, dtype=tf.int32)
img = tf.image.rot90(img, random_value)
mask = tf.image.rot90(mask, random_value)
mask = tf.squeeze(mask)
return {'img': img, 'mouth': mask}, label
def zoom(x, label, size):
# Generate 20 crop settings, ranging from a 1% to 30% crop.
scales = list(np.arange(0.6, 1.0, 0.03))
boxes = np.zeros((len(scales), 4))
......@@ -34,18 +54,39 @@ def zoom(x: tf.Tensor, label, size):
x2 = y2 = 0.5 + (0.5 * scale)
boxes[i] = [x1, y1, x2, y2]
def random_crop(img):
def random_crop(dataset_dict):
img = x['img']
mask = x['mouth']
mask = tf.expand_dims(mask, -1)
# Create different crops for an image
crops = tf.image.crop_and_resize([img], boxes=boxes, box_indices=np.zeros(len(scales)), crop_size=size)
# Return a random crop
return crops[tf.random.uniform(shape=[], minval=0, maxval=len(scales), dtype=tf.int32)]
crops_image = tf.image.crop_and_resize([img], boxes=boxes, box_indices=np.zeros(len(scales)), crop_size=size)
crops_mask = tf.image.crop_and_resize([mask], boxes=boxes, box_indices=np.zeros(len(scales)), crop_size=size)
# Select a random crop
random_index = tf.random.uniform(shape=[], minval=0, maxval=len(scales), dtype=tf.int32)
img = crops_image[random_index]
mask = crops_mask[random_index]
mask = tf.squeeze(mask)
return {'img':img, 'mouth': mask}
choice = tf.random.uniform(shape=[], minval=0., maxval=1., dtype=tf.float32)
# Only apply cropping 80% of the time
return tf.cond(choice < 0.8, lambda: x, lambda: random_crop(x)), label
'''
Clipping helper
'''
def clip(x, y):
img = x['img']
mask = x['mouth']
img = tf.clip_by_value(img, 0, 1)
return {'img': img, 'mouth': mask}, y
'''
Data Loader Class
provides main functions for providing the data in the correct for for the training and testing
'''
class DataLoader(object):
def __init__(self, data_path, batch_size, img_width, img_height, channels, should_size_dataset1 = 0, should_size_dataset2=0, split_size=0.0, augment=False, annotation=None):
......@@ -177,9 +218,9 @@ class DataLoader(object):
filenames_list, mouths_list, labels_list = self._load_data_as_filename_lists(self.data_path)
# Shuffle the lists
dataset_list = list(zip(filenames_list, labels_list))
dataset_list = list(zip(filenames_list, mouths_list, labels_list))
random.shuffle(dataset_list)
filenames_list, labels_list = zip(*dataset_list)
filenames_list, mouths_list, labels_list = zip(*dataset_list)
# Will skipp the splitting if it would result in a full copy of dataset in dataset_1 or dataset_2
......@@ -209,14 +250,14 @@ class DataLoader(object):
for f in self.AUGMENTATIONS:
self.dataset_1 = self.dataset_1.map(lambda x,y: tf.cond(tf.random.uniform([], 0, 1) > 0.1, lambda: f(x,y, (self.IMG_WIDTH, self.IMG_HEIGHT)), lambda: (x,y)),
num_parallel_calls=self.NR_THREADS)
self.dataset_1 = self.dataset_1.map(lambda x,y: (tf.clip_by_value(x, 0, 1),y))
self.dataset_1 = self.dataset_1.map(lambda x,y: clip(x,y))
if self.dataset_2:
self.dataset_2 = self.dataset_2.map(self.decode_img, num_parallel_calls=self.NR_THREADS)
self.dataset_2 = self.dataset_2.cache()
for f in self.AUGMENTATIONS:
self.dataset_2 = self.dataset_2.map(lambda x, y: tf.cond(tf.random.uniform([], 0, 1) > 0.1, lambda: f(x, y, (self.IMG_WIDTH, self.IMG_HEIGHT)),
lambda: (x, y)), num_parallel_calls=self.NR_THREADS)
self.dataset_2 = self.dataset_2.map(lambda x, y: (tf.clip_by_value(x, 0, 1), y))
self.dataset_2 = self.dataset_2.map(lambda x, y: clip(x, y))
self.dataset_1 = self.dataset_1.repeat(self.dataset_1_repeat_factor).shuffle(self.dataset_1_size*self.dataset_1_repeat_factor).batch(self.batch_size, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment