Creating generator object with image augmentation to train Convolutional Neural Networks with Keras

Question

I am currently self-studying on Python generator object and use it to generate training data and do augmentation on-the-fly, then feed it into Convolutional Neural Networks.

Could anyone please help me to review my code? It runs properly, but I need review to make it more efficient and more properly structured. Besides, how can I check that using the generator will consume less memory (compared to just passing regular numpy array to the model)?

Thank you very much!

from keras.utils import to_categorical
from keras.models import Sequential
from keras.layers import Dense, Conv2D, Flatten
import pandas as pd
import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
import tensorflow as tf
from augment import ImageAugment

class Generator():
    def __init__(self, feat, labels, width, height):
        self.feat = feat
        self.labels = labels
        self.width = width
        self.height = height

    def gen(self):
        '''
        Yields generator object for training or evaluation without batching
        Yields:
            im: np.array of (1,width,height,1) of images
            label: np.array of one-hot vector of label (1,num_labels)
        '''
        feat = self.feat
        labels = self.labels
        width = self.width
        height = self.height
        i=0
        while (True):
            im = cv2.imread(feat[i],0)
            im = im.reshape(width,height,1)
            im = np.expand_dims(im,axis=0)
            label = np.expand_dims(labels[i],axis=0)
            yield im,label
            i+=1

            if i>=len(feat):
                i=0


    def gen_test(self):
        '''
        Yields generator object to do prediction
        Yields:
            im: np.array of (1,width,height,1) of images
        '''
        feat = self.feat
        width = self.width
        height = self.height
        i=0
        while (True):
            im = cv2.imread(feat[i],0)
            im = im.reshape(width,height,1)
            im = np.expand_dims(im,axis=0)
            yield im
            i+=1


    def gen_batching(self, batch_size):
        '''
        Yields generator object with batching of batch_size
        Args:
            batch_size (int): batch_size
        Yields:
            feat_batch: np.array of (batch_size,width,height,1) of images
            label_batch: np.array of (batch_size,num_labels)
        '''
        feat = self.feat
        labels = self.labels
        width = self.width
        height = self.height
        num_examples = len(feat)
        num_batch = num_examples/batch_size
        X = []
        for n in range(num_examples):
            im = cv2.imread(feat[n],0)
            try:
                im = im.reshape(width,height,1)
            except:
                print('Error on this image: ', feat[n])
            X.append(im)
        X = np.array(X)

        feat_batch = np.zeros((batch_size,width,height,1))
        label_batch = np.zeros((batch_size,labels.shape[1]))
        while(True):
            for i in range(batch_size):
                index = np.random.randint(X.shape[0],size=1)[0] #shuffle the data
                feat_batch[i] = X[index]
                label_batch[i] = labels[index]
            yield feat_batch,label_batch

    # def on_next(self):
    #     '''
    #     Advance to the next generator object
    #     '''
    #     gen_obj = self.gen_test()
    #     return next(gen_obj)
    #
    # def gen_show(self, pred):
    #     '''
    #     Show the image generator object
    #     '''
    #     i=0
    #     while(True):
    #         image = self.on_next()
    #         image = np.squeeze(image,axis=0)
    #         cv2.imshow('image', image)
    #         cv2.waitKey(0)
    #         i+=1

    def gen_augment(self,batch_size,augment):
        '''
        Yields generator object with batching of batch_size and augmentation.
        The number of examples for 1 batch will be multiplied based on the number of augmentation

        augment represents [speckle, gaussian, poisson]. It means, the augmentation will be done on the augment list element that is 1
        for example, augment = [1,1,0] corresponds to adding speckle noise and gaussian noise
        if batch_size = 100, the number of examples in each batch will become 300

        Args:
            batch_size (int): batch_size
            augment (list): list that defines what kind of augmentation we want to do
        Yields:
            feat_batch: np.array of (batch_size*n_augment,width,height,1) of images
            label_batch: np.array of (batch_size*n_augment,num_labels)
        '''
        feat = self.feat
        labels = self.labels
        width = self.width
        height = self.height

        num_examples = len(feat)
        num_batch = num_examples/batch_size
        X = []
        for n in range(num_examples):
            im = cv2.imread(feat[n],0)
            try:
                im = im.reshape(width,height,1)
            except:
                print('Error on this image: ', feat[n])
            X.append(im)
        X = np.array(X)

        n_augment = augment.count(1)
        print('Number of augmentations: ', n_augment)
        feat_batch = np.zeros(((n_augment+1)*batch_size,width,height,1))
        label_batch = np.zeros(((n_augment+1)*batch_size,labels.shape[1]))

        while(True):
            i=0
            while (i<=batch_size):
                index = np.random.randint(X.shape[0],size=1)[0] #shuffle the data
                aug = ImageAugment(X[index])
                feat_batch[i] = X[index]
                label_batch[i] = labels[index]

                j=0
                if augment[0] == 1:
                    feat_batch[(j*n_augment)+i+batch_size] = aug.add_speckle_noise()
                    label_batch[(j*n_augment)+i+batch_size] = labels[index]
                    j+=1

                if augment[1] == 1:
                    feat_batch[(j*n_augment)+i+batch_size] = aug.add_gaussian_noise()
                    label_batch[(j*n_augment)+i+batch_size] = labels[index]
                    j+=1

                if augment[2] == 1:
                    feat_batch[(j*n_augment)+i+batch_size] = aug.add_poisson_noise()
                    label_batch[(j*n_augment)+i+batch_size] = labels[index]
                    j+=1

                i+=1


            yield feat_batch,label_batch

def CNN_model(width,height):
    # #create model
    model = Sequential()
    model.add(Conv2D(64, kernel_size=3, activation="relu", input_shape=(width,height,1)))
    model.add(Conv2D(32, kernel_size=3, activation="relu"))
    model.add(Flatten())
    model.add(Dense(labels.shape[1], activation="softmax"))

    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    return model


if __name__ == "__main__":
    input_dir = './mnist'
    output_file = 'dataset.csv'

    filename = []
    label = []
    for root,dirs,files in os.walk(input_dir):
        for file in files:
            full_path = os.path.join(root,file)
            filename.append(full_path)
            label.append(os.path.basename(os.path.dirname(full_path)))

    data = pd.DataFrame(data={'filename': filename, 'label':label})
    data.to_csv(output_file,index=False)

    labels = pd.get_dummies(data.iloc[:,1]).values

    X, X_val, y, y_val = train_test_split(
                                            filename, labels,
                                            test_size=0.2,
                                            random_state=1234,
                                            shuffle=True,
                                            stratify=labels
                                            )

    X_train, X_test, y_train, y_test = train_test_split(
                                                        X, y,
                                                        test_size=0.025,
                                                        random_state=1234,
                                                        shuffle=True,
                                                        stratify=y
                                                        )

    width = 28
    height = 28

    test_data = pd.DataFrame(data={'filename': X_test})


    image_gen_train = Generator(X_train,y_train,width,height)
    image_gen_val = Generator(X_val,y_val,width,height)
    image_gen_test = Generator(X_test,None,width,height)


    batch_size = 900
    print('len data: ', len(X_train))
    print('len test data: ', len(X_test))

    #augment represents [speckle, gaussian, poisson]. It means, the augmentation will be done on the augment list element that is 1
    #for example, augment = [1,1,0] corresponds to adding speckle noise and gaussian noise
    augment = [1,1,1]
    model = CNN_model(width,height)

    model.fit_generator(
                        generator=image_gen_train.gen_augment(batch_size=batch_size,augment=augment),
                        steps_per_epoch=np.ceil(len(X_train)/batch_size),
                        epochs=20,
                        verbose=1,
                        validation_data=image_gen_val.gen(),
                        validation_steps=len(X_val)
                        )
    model.save('model_aug_3.h5')
    model = tf.keras.models.load_model('model_aug_3.h5')

    #Try evaluate_generator
    image_gen_test = Generator(X_test,y_test,width,height)
    print(model.evaluate_generator(
                            generator=image_gen_test.gen(),
                            steps=len(X_test)
                            ))

    #Try predict_generator
    image_gen_test = Generator(X_test,None,width,height)
    pred = model.predict_generator(
                            generator=image_gen_test.gen_test(),
                            steps=len(X_test)
                            )
    pred = np.argmax(pred,axis=1)
    # image_gen_test = Generator(X_test,pred,width*3,height*3)
    # image_gen_test.gen_show(pred)
    wrong_pred = []
    for i,ex in enumerate(zip(pred,y_test)):
        if ex[0] != np.argmax(ex[1]):
            wrong_pred.append(i)
    print(wrong_pred)

    # for i in range(len(X_test)):
    #     im = cv2.imread(X_test[i],0)
    #     im = cv2.putText(im, str(pred[i]), (10,15), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
    #     print(i)
    #     cv2.imshow('image',im)
    #     cv2.waitKey(0)
```

Answer 1

Class syntax

The parens are not necessary here:

class Generator():

Type hints

Type hints will add more built-in documentation to your code, and will help the better IDEs give you static analysis hints:

def __init__(self, feat, labels, width, height):

can become, at a guess,

def __init__(self, feat: Sequence[float], labels: Sequence[str], width: int, height: int):

Iteration

First of all, this in gen:

    while (True):

does not require parens. Also, rather than manually maintaining i, you should use for i in itertools.cycle(range(len(feat))).

Similarly, this:

        i=0
        while (i<=batch_size):
            # ...
            i+=1

should just be for i in range(batch_size):.

Bare except

This:

        try:
            im = im.reshape(width,height,1)
        except:
            print('Error on this image: ', feat[n])

should not have an exception catch clause that is so general. The broadest that you should catch, if you do not know what the typical exceptions are, is Exception. Catching everything also prevents KeyboardInterrupt (Ctrl+C break) from working.

Variable names

This:

    X = []

should be lowercase, according to the PEP8 style guide.