Ceretai/task-4-active_learning/2_Active_learning/simple_race_class_model.py

'''
Simple initial model for race classification. 

Based on https://towardsdatascience.com/building-a-multi-output-convolutional-neural-network-with-keras-ed24c7bc1178
'''

import numpy as np 
import pandas as pd
import os
import glob
import pandas as pd
from PIL import Image

import tensorflow as tf
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.models import Model
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import MaxPooling2D
from tensorflow.keras.layers import Activation
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import Lambda
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Input
#import tensorflow as tf

TRAIN_TEST_SPLIT = 0.7
IM_WIDTH = IM_HEIGHT = 224
dataset_dict = {
    'race_id': {
        0: 'East Asian', 
        1: 'White', 
        2: 'Latino_Hispanic', 
        3: 'Southeast Asian', 
        4: 'Black',
        5: 'Indian',
        6: 'Middle Eastern'
    }
}

dataset_dict['race_alias'] = dict((r, i) for i, r in dataset_dict['race_id'].items())

class FFDataGenerator():
    '''
    Data generator for the FairFace dataset. This class should be used
    when training our Keras multi-output model.
    '''

    def __init__(self, df):
        self.df = df

    def generate_split_indexes(self):
        p = np.random.permutation(len(self.df))
        train_up_to = int(len(self.df) * TRAIN_TEST_SPLIT)
        train_idx = p[:train_up_to]
        test_idx = p[train_up_to:]

        train_up_to = int(train_up_to * TRAIN_TEST_SPLIT)
        train_idx, valid_idx = train_idx[:train_up_to], train_idx[train_up_to:]

        #Reads race
        self.df['race_id'] = self.df['race'].map(lambda race: dataset_dict['race_alias'][race])

        return train_idx, valid_idx, test_idx

    def preprocess_image(self, img_path):
        '''
        Used to perform some minor preprocessing on the image before 
        inputting into the network. 
        '''

        im = Image.open('../1_Public_datasets/1_FairFace/' + img_path)
        im = im.resize((IM_WIDTH, IM_HEIGHT))
        im = np.array(im) / 255.0

        return im

    def generate_images(self, image_idx, is_training, batch_size = 16):
        '''
        Used to generate a batch with images when training/testing
        /validating our Keras model.
        '''

        # arrays to store our batched data
        images, races = [], []
        while True:
            for idx in image_idx:
                person = self.df.iloc[idx]
                race = person['race_id']
                file = person['file']

                im = self.preprocess_image(file)

                races.append(to_categorical(race, len(dataset_dict['race_id'])))
                images.append(im)

                #yielding condition
                if len(images) >= batch_size:
                    yield np.array(images), [np.array(races)]
                    images, races = [], []

                if not is_training:
                    break

class FFOutputModel():
    '''
    Used to generate our multi-output model. Convolutional Layers is
    defined on the make_default_hidden_layers method.
    '''

    def make_default_hidden_layers(self, inputs):
        '''
        Used to generate a default set of hidden layers. The
        structure used in this network is defined as:

        Conv2d -> BatchNormalization -> Pooling -> Dropout
        '''

        x = Conv2D(16, (3, 3), padding = 'same')(inputs)
        x = Activation('relu')(x)
        x = BatchNormalization(axis = -1)(x)
        x = MaxPooling2D(pool_size = (3, 3))(x)
        x = Dropout(0.25)(x)

        x = Conv2D(32, (3, 3), padding = 'same')(inputs)
        x = Activation('relu')(x)
        x = BatchNormalization(axis = -1)(x)
        x = MaxPooling2D(pool_size = (2, 2))(x)
        x = Dropout(0.25)(x)

        x = Conv2D(32, (3, 3), padding = 'same')(inputs)
        x = Activation('relu')(x)
        x = BatchNormalization(axis = -1)(x)
        x = MaxPooling2D(pool_size = (2, 2))(x)
        x = Dropout(0.25)(x)

        return x

    def build_race_branch(self, inputs, num_races):
        '''
        Used to build the race branch of our face recognition
        network.
        This branch is composed of thre Conv -> BN -> Pool ->
        Dropout blocks,
        followed by the Dense output layer.
        '''

        x = Lambda(lambda c: tf.image.rgb_to_grayscale(c))(inputs)

        x = self.make_default_hidden_layers(inputs)

        x = Flatten()(x)
        x = Dense(128)(x)
        x = Activation('relu')(x)
        x = BatchNormalization()(x)
        x = Dropout(0.5)(x)
        x = Dense(num_races)(x)
        x = Activation('softmax', name = 'race_output')(x)

        return x

    def assemble_full_model(self, width, height, num_races):
        '''
        Used to assemble our model CNN.
        '''

        input_shape = (height, width, 3)

        inputs = Input(shape = input_shape)

        race_branch = self.build_race_branch(inputs, num_races)

        model = Model(inputs = inputs, outputs = [race_branch], name = 'face_net')

        return model