DataLoader.py

import glob
import math
import random
import numpy as np
import tensorflow as tf



'''
    Augmentation functions    
'''
def flip(x, 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)
    # rotate either 0 , 90, 180 or 270 degrees (0 or 1 or 2 or 3 times 90 degrees)
    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))

    for i, scale in enumerate(scales):
        x1 = y1 = 0.5 - (0.5 * scale)
        x2 = y2 = 0.5 + (0.5 * scale)
        boxes[i] = [x1, y1, x2, y2]

    def random_crop(dataset_dict):
        img = x['img']
        mask = x['mouth']
        mask = tf.expand_dims(mask, -1)
        # Create different crops for an image
        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, shuffle=True):
        self.data_path = data_path
        self.batch_size = batch_size
        self.split_size = split_size
        # -1 means detecting automatically how many threads are possible
        self.NR_THREADS = -1
        self.IMG_WIDTH = img_width
        self.IMG_HEIGHT = img_height
        self.CHANNELS = channels
        self.AUGMENT = augment
        self.AUGMENTATIONS = [flip, color, zoom, rotate]
        self.annotation = annotation
        self.shuffle = shuffle

        self.classes = [item.name for item in data_path.glob('*') if item.name != '.DS_Store']
        self.n_classes = len(self.classes)

        self.dataset_1 = None
        self.dataset_2 = None
        self.dataset_1_size = -1
        self.dataset_2_size = -1
        self.dataset_1_class_elem_count = {}
        self.dataset_2_class_elem_count = {}
        self.dataset_1_size_min_count = should_size_dataset1
        self.dataset_2_size_min_count = should_size_dataset2
        self.dataset_1_repeat_factor = 1
        self.dataset_2_repeat_factor = 1

    def _create_dataset(self, filenames_list, mouths_list, labels_list):
        # Creating constants for dataset
        filenames = tf.constant(filenames_list)
        mouths = tf.constant(mouths_list, dtype='float32')
        labels = tf.constant(labels_list)
        return tf.data.Dataset.from_tensor_slices(({'img': filenames, 'mouth': mouths}, labels))


    def decode_img(self, input, one_hot_encodings):
        #input = (self.decode_img(img), input[1])
        filepath = input['img']
        # load the raw data from the file as a string
        img = tf.io.read_file(filepath)
        # convert the compressed string to a 3D uint8 tensor
        img = tf.image.decode_jpeg(img, channels=self.CHANNELS)
        # Use `convert_image_dtype` to convert to floats in the [0,1] range.
        img = tf.image.convert_image_dtype(img, tf.float32)
        # resize the image to the desired size.
        input['img'] = tf.image.resize(img, [self.IMG_WIDTH, self.IMG_HEIGHT])
        return input, one_hot_encodings

    '''
        Count the actual number of element for each class in each dataset.
        This is the base information to calculate the right weights to fix the imbalance of the dataset.
    '''
    def count_class_elements(self, one_hot_encodings):
        class_counts = {k: 0 for k in range(self.n_classes)}
        for one_hot_encoding in one_hot_encodings:
            class_counts[one_hot_encoding.index(1)] += 1
        return class_counts

    def _get_one_hot_for_label(self, classname):
        index = self.classes.index(classname)
        one_hot = [0] * self.n_classes
        one_hot[index] = 1
        return one_hot

    def _load_data_as_filename_lists(self, data_path):
        filenames = []
        mouth_masks = []
        one_hot_encodings = []

        # Length of list including all files from previous directories (directory == one class of training data)
        last_length = 0
        for class_index, classname in enumerate(self.classes):
            filenames_of_class = [item for item in glob.glob(str(data_path) + '/' + classname + '/*')]
            # Add the Filenames to the list of all filenames
            filenames.extend(filenames_of_class)
            # Add the Filenames to the list of all filenames
            one_hot_encodings.extend([self._get_one_hot_for_label(classname)] * len(filenames_of_class))

        if self.annotation:
            for file_full_path in filenames:
                filename = file_full_path.split('/')[-1]
                annotations = self.annotation.get_annotations(filename, ['mouth'])
                mask = np.ones(shape=(self.IMG_WIDTH, self.IMG_HEIGHT))
                if len(annotations) > 0:
                    annotations = annotations[0][1] # if there is an annotation of the mouth present, extract the corners
                    width_orig, height_orig = annotations[0]
                    width_scalar = self.IMG_WIDTH / width_orig
                    height_scalar = self.IMG_HEIGHT / height_orig
                    x1, y1 = annotations[1]
                    x2, y2 = annotations[2]
                    x1 = round(x1 * width_scalar)
                    x2 = round(x2 * width_scalar)
                    y1 = round(y1 * height_scalar)
                    y2 = round(y2 * height_scalar)

                    mask[y1:y2+1,x1:x2+1] = 0
                mouth_masks.append(mask)
        else:
            mouth_masks = [np.ones(shape=(self.IMG_WIDTH, self.IMG_HEIGHT))] * len(filenames) # create empty lists if no additional annotation for the mouth is provided
        return filenames, mouth_masks, one_hot_encodings

    def _filename_and_labels_split(self, filenames_list, mouths_list, labels_list):
        splitting_element = math.ceil(len(filenames_list) * self.split_size)
        filenames2_list = filenames_list[splitting_element:]
        filenames_list = filenames_list[:splitting_element]
        mouths2_list = mouths_list[splitting_element:]
        mouths_list = mouths_list[:splitting_element]
        labels2_list = labels_list[splitting_element:]
        labels_list = labels_list[:splitting_element]
        return filenames_list, mouths_list, labels_list, filenames2_list, mouths2_list, labels2_list

    def _get_weights(self, dataset_class_elem_count, dataset_size):
        weights = {k: 0 for k in range(self.n_classes)}
        for key, value in  dataset_class_elem_count.items():
            weights[key]  = dataset_size / (self.n_classes * value)
        return weights

    def get_weights_dataset1(self):
        return self._get_weights(self.dataset_1_class_elem_count, self.dataset_1_size)

    def get_weights_dataset2(self):
        return self._get_weights(self.dataset_2_class_elem_count, self.dataset_2_size)

    def load_dataset(self):
        print('Loading dataset information')
        # Load filenames and labels as two corresponding lists of strings
        filenames_list, mouths_list, labels_list = self._load_data_as_filename_lists(self.data_path)

        # Shuffle the lists
        if self.shuffle:
            dataset_list = list(zip(filenames_list, mouths_list, labels_list))
            random.shuffle(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
        if self.split_size not in [0, 1]:
            filenames_list, mouths_list, labels_list, filenames2_list, mouths2_list, labels2_list = self._filename_and_labels_split(filenames_list, labels_list)
            self.dataset_2 = self._create_dataset(filenames2_list, mouths2_list, labels2_list)
            if self.shuffle:
                self.dataset_2 = self.dataset_2.shuffle(buffer_size=(len(filenames2_list)))
            self.dataset_2_size = len(filenames2_list)
            if self.dataset_2_size_min_count != 0:
                self.dataset_2_repeat_factor = math.ceil(self.dataset_2_size_min_count / self.dataset_2_size)
            self.dataset_2_class_elem_count = self.count_class_elements(labels2_list)

        # Creating actual TF Dataset 1
        self.dataset_1 = self._create_dataset(filenames_list, mouths_list, labels_list)
        if self.shuffle:
            self.dataset_1 = self.dataset_1.shuffle(buffer_size=(len(filenames_list)))
        self.dataset_1_size = len(filenames_list)
        if self.dataset_1_size_min_count != 0:
            self.dataset_1_repeat_factor = math.ceil(self.dataset_1_size_min_count / self.dataset_1_size)

        self.dataset_1_class_elem_count = self.count_class_elements(labels_list)

        # Load images
        print('Loading images of dataset into memory.')
        self.dataset_1 = self.dataset_1.map(self.decode_img, num_parallel_calls=self.NR_THREADS)
        self.dataset_1 = self.dataset_1.cache().repeat(self.dataset_1_repeat_factor)

        if self.AUGMENT:
            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: 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().repeat(self.dataset_2_repeat_factor)
                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: clip(x, y))

        if self.shuffle:
            self.dataset_1 = self.dataset_1.shuffle(self.dataset_1_size)
        self.dataset_1 = self.dataset_1.batch(self.batch_size, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)

        if self.dataset_2:
            if self.shuffle:
                self.dataset_2 = self.dataset_2.shuffle(self.dataset_2_size)
            self.dataset_2 = self.dataset_2.batch(self.batch_size, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)
            return self.dataset_1, self.dataset_2
        else:
            return self.dataset_1