How to build a Custom Data Generator for Keras/tf.Keras where X images are being augmented and corresponding Y labels are also images

user_11725056 0
Viewed 3197

I am working on Image Binarization using UNet and have a dataset of 150 images and their binarized versions too. My idea is to augment the images randomly to make them look like they are differentso I have made a function which inserts any of the 4-5 types of Noises, skewness, shearing and so on to an image. I could have easily used

ImageDataGenerator(preprocess_function=my_aug_function) to augment the images but the problem is that my y target is also an image. Also, I could have used something like:

train_dataset = (
    train_dataset.map(
        encode_single_sample, num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    .batch(batch_size)
    .prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
)

But it has 2 problems:

  1. With larger dataset, it'll blow up the memory as data needs to be already in the memory
  2. This is the crucial part that I need to augment the images on the go to make it look like I have a huge dataset.

Another Solution could be saving augmented images to a directory and making them 30-40K and then loading them. It would be silly thing to do.

Now the idea part is that I can use Sequence as the parent class but How can I keep on augmenting and generating new images on the fly with respective Y binarized images?

I have an idea as the below code. Can somebody help me with the augmentation and generation of y images. I have my X_DIR, Y_DIR where image names for binarised and original are same but stored in different directories.

class DataGenerator(tensorflow.keras.utils.Sequence):
    def __init__(self, files_path, labels_path, batch_size=32, shuffle=True, random_state=42):
        'Initialization'
        self.files = files_path
        self.labels = labels_path
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.random_state = random_state
        self.on_epoch_end()


    def on_epoch_end(self):
        'Updates indexes after each epoch'
        # Shuffle the data here


    def __len__(self):
        return int(np.floor(len(self.files) / self.batch_size))

    def __getitem__(self, index):
        # What do I do here? 


    def __data_generation(self, files):
        # I think this is responsible for Augmentation but no idea how should I implement it and how does it works.
tensorflow2.0 keras tensorflow deep-learning python
3 Answers

Custom Image Data Generator

load Directory data into dataframe for CustomDataGenerator

def data_to_df(data_dir, subset=None, validation_split=None):
    df = pd.DataFrame()
    filenames = []
    labels = []
    
    for dataset in os.listdir(data_dir):
        img_list = os.listdir(os.path.join(data_dir, dataset))
        label = name_to_idx[dataset]
        
        for image in img_list:
            filenames.append(os.path.join(data_dir, dataset, image))
            labels.append(label)
        
    df["filenames"] = filenames
    df["labels"] = labels
    
    if subset == "train":
        split_indexes = int(len(df) * validation_split)
        train_df = df[split_indexes:]
        val_df = df[:split_indexes]
        return train_df, val_df
    
    return df

train_df, val_df = data_to_df(train_dir, subset="train", validation_split=0.2)

Custom Data Generator


import tensorflow as tf
from PIL import Image
import numpy as np

class CustomDataGenerator(tf.keras.utils.Sequence):

    ''' Custom DataGenerator to load img 
    
    Arguments:
        data_frame = pandas data frame in filenames and labels format
        batch_size = divide data in batches
        shuffle = shuffle data before loading
        img_shape = image shape in (h, w, d) format
        augmentation = data augmentation to make model rebust to overfitting
    
    Output:
        Img: numpy array of image
        label : output label for image
    '''
    
    def __init__(self, data_frame, batch_size=10, img_shape=None, augmentation=True, num_classes=None):
        self.data_frame = data_frame
        self.train_len = len(data_frame)
        self.batch_size = batch_size
        self.img_shape = img_shape
        self.num_classes = num_classes
        print(f"Found {self.data_frame.shape[0]} images belonging to {self.num_classes} classes")

    def __len__(self):
        ''' return total number of batches '''
        self.data_frame = shuffle(self.data_frame)
        return math.ceil(self.train_len/self.batch_size)

    def on_epoch_end(self):
        ''' shuffle data after every epoch '''
        # fix on epoch end it's not working, adding shuffle in len for alternative
        pass
    
    def __data_augmentation(self, img):
        ''' function for apply some data augmentation '''
        img = tf.keras.preprocessing.image.random_shift(img, 0.2, 0.3)
        img = tf.image.random_flip_left_right(img)
        img = tf.image.random_flip_up_down(img)
        return img
        
    def __get_image(self, file_id):
        """ open image with file_id path and apply data augmentation """
        img = np.asarray(Image.open(file_id))
        img = np.resize(img, self.img_shape)
        img = self.__data_augmentation(img)
        img = preprocess_input(img)

        return img

    def __get_label(self, label_id):
        """ uncomment the below line to convert label into categorical format """
        #label_id = tf.keras.utils.to_categorical(label_id, num_classes)
        return label_id

    def __getitem__(self, idx):
        batch_x = self.data_frame["filenames"][idx * self.batch_size:(idx + 1) * self.batch_size]
        batch_y = self.data_frame["labels"][idx * self.batch_size:(idx + 1) * self.batch_size]
        # read your data here using the batch lists, batch_x and batch_y
        x = [self.__get_image(file_id) for file_id in batch_x] 
        y = [self.__get_label(label_id) for label_id in batch_y]

        return tf.convert_to_tensor(x), tf.convert_to_tensor(y)
user_122238150

You can use libraries like albumentations and imgaug, both are good but I have heard there are issues with random seed with albumentations. Here's an example of imgaug taken from the documentation here:

seq = iaa.Sequential([
    iaa.Dropout([0.05, 0.2]),      # drop 5% or 20% of all pixels
    iaa.Sharpen((0.0, 1.0)),       # sharpen the image
    iaa.Affine(rotate=(-45, 45)),  # rotate by -45 to 45 degrees (affects segmaps)
    iaa.ElasticTransformation(alpha=50, sigma=5)  # apply water effect (affects segmaps)
], random_order=True)

# Augment images and segmaps.
images_aug = []
segmaps_aug = []
for _ in range(len(input_data)):
    images_aug_i, segmaps_aug_i = seq(image=image, segmentation_maps=segmap)
    images_aug.append(images_aug_i)
    segmaps_aug.append(segmaps_aug_i)

You are going in the right way with the custom generator. In __getitem__, make a batch using batch_x = self.files[index:index+batch_size] and same with batch_y, then augment them using X,y = __data_generation(batch_x, batch_y) which will load images(using any library you like, I prefer opencv), and return the augmented pairs (and any other manipulation).

Your __getitem__ will then return the tuple (X,y)

user_82437970

You can use ImageDataGenerator even if your label is an image. Here is a simple example of how you can do that:

Code:

# Specifying your data augmentation here for both image and label
image_datagen = tf.keras.preprocessing.image.ImageDataGenerator()
mask_datagen = tf.keras.preprocessing.image.ImageDataGenerator()

# Provide the same seed and keyword arguments to the flow methods
seed = 1
image_generator = image_datagen.flow_from_directory(
    data_dir,
    class_mode=None,
    seed=seed)
mask_generator = mask_datagen.flow_from_directory(
    data_dir,
    class_mode=None,
    seed=seed)

# Combine the image and label generator.
train_generator = zip(image_generator, mask_generator)

Now, if you iterate over it you will get:

for image, label in train_generator:
  print(image.shape,label.shape)
  break

Output:

(32, 256, 256, 3) (32, 256, 256, 3)

You can use this train_generator with fit() command.

Code:

model.fit_generator(
    train_generator,
    steps_per_epoch=2000,
    epochs=50)

With flow_from_directory your memory won't be cluttered and Imagedatagenerator will take care of the augmentation part.

user_95579700
Related