train.py

"""
Author:
    Ahmed Telili
"""

import argparse
import datetime
import os
import random
import time

import cv2
import h5py
import numpy as np
import pandas as pd
import PIL
import sklearn
import tensorflow as tf
from keras import backend as K
from keras.activations import sigmoid, softmax
from keras.callbacks import (Callback, EarlyStopping, ModelCheckpoint,
                             ReduceLROnPlateau, TensorBoard)
from keras.layers import Layer
from keras.models import load_model
from PIL import Image
from scipy.stats import pearsonr, spearmanr
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle
from tensorflow import keras
from tensorflow.keras import (Input, applications, initializers, layers,
                              models, optimizers)
from tensorflow.keras.layers import (AveragePooling2D, Bidirectional,
                                     Concatenate, Conv2D, Dense, Dropout,
                                     Flatten, GlobalAveragePooling2D,
                                     GlobalMaxPooling2D, LSTM, Lambda,
                                     MaxPooling2D, Reshape, TimeDistributed,
                                     concatenate)
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing import image
from tensorflow_addons.metrics import RSquare
from tensorflow_addons.optimizers import AdamW
from tqdm import tqdm
from tqdm.keras import TqdmCallback

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    tf.config.experimental.set_memory_growth(gpus[0], True)

tf.keras.backend.clear_session()


def logistic_func(X, bayta1, bayta2, bayta3, bayta4):
    # 4-parameter logistic function
    logistic_part = 1 + np.exp(-np.divide(X - bayta3, np.abs(bayta4)))
    yhat = bayta2 + np.divide(bayta1 - bayta2, logistic_part)
    return yhat


def root_mean_squared_error(y_true, y_pred):
    return K.sqrt(K.mean(K.square(y_pred - y_true)))


def build_model(batch_shape):
    model = models.Sequential()
    model.add(TimeDistributed(LSTM(32, return_sequences=True,
                                   kernel_initializer='random_normal',
                                   recurrent_initializer='random_normal',
                                   dropout=0.4, recurrent_dropout=0),
                              input_shape=batch_shape))
    model.add(TimeDistributed(Flatten()))
    model.add(LSTM(16, return_sequences=True,
                   kernel_initializer='random_normal',
                   recurrent_initializer='random_normal',
                   dropout=0.4, recurrent_dropout=0))
    model.add(Flatten())
    model.add(Dense(32, activation='relu',
                    kernel_initializer='random_normal'))
    model.add(layers.Dropout(rate=0.5))
    model.add(layers.Dense(1, activation='relu',
                           kernel_initializer='random_normal'))

    model.compile(optimizer=optimizers.Adam(), loss=root_mean_squared_error,
                  metrics='mae')
    model.summary()
    return model


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="End2End_train")

    parser.add_argument('-nf', '--num_frames', default=30, type=int,
                        help='Number of cropped frames per video.')
    parser.add_argument('-nb', '--num_blocks', default=16, type=int,
                        help='Number of blocks per frame.')
    parser.add_argument('-np', '--num_patches', default=156, type=int,
                        help='Number of cropped patches per frame.')
    parser.add_argument('-b', '--batch_size', default=1, type=int,
                        help='Batch size.')

    if not os.path.exists('./models'):
        os.makedirs('./models')

    args = parser.parse_args()
    
    # For reproducibility 
    idx = [7]

    l = os.listdir('./label')
    l.sort()
    lab_y = np.zeros((196, 3))
    for i in range(len(l)):
        feat = np.load('./label/' + l[i])
        lab_y[i, :] = feat
    np.save('label.npy', lab_y)

    l = os.listdir('./features_X')
    l.sort()
    dens_X = np.zeros((196, 2048))  # 2048: ResNet50 output dim.
    for i in range(len(l)):
        feat = np.load('./features_X/' + l[i])
        dens_X[i, :] = feat
    np.save('features_X.npy', dens_X)

    X = np.load('./features_X.npy')
    y = np.load('./label.npy')
    label_P1 = y[:, 0]

    # Initialization of metrics lists
    score_r2_P1 = []
    score_srocc_P1 = []
    score_plcc_P1 = []

    start_time = time.time()
    for i in idx:
        md = ModelCheckpoint(filepath='./models/vmaf/model.h5', monitor='val_loss', mode='auto',
                             save_best_only=True, verbose=1)
        rd = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=15, min_lr=1e-11, verbose=1, mode='min')
        ear = EarlyStopping(monitor='val_loss', mode='min', patience=50, verbose=1)
        callbacks_k = [md, rd, ear]

        tf.keras.backend.clear_session()

        X_train, X_test, y_train, y_test = train_test_split(X, label_P1, test_size=0.2, random_state=i)
        model = build_model((args.num_blocks, args.num_patches, 2048))

        history = model.fit(X_train, y_train, epochs=200, validation_data=(X_test, y_test),
                            verbose=2, shuffle=True, callbacks=callbacks_k, batch_size=args.batch_size)

        del model

        model = load_model('./models/vmaf/model.h5', compile=False)
        y_pred = model.predict(X_test)

        del model

        y_test = y_test.reshape(y_test.shape[0])
        y_pred = y_pred.reshape(y_pred.shape[0])

        score_r2_P1.append(r2_score(y_test, y_pred))
        score_srocc_P1.append(spearmanr(y_pred, y_test)[0])
        score_plcc_P1.append(pearsonr(y_pred, y_test)[0])

    end_time = time.time()
    print('Time elapsed:', end_time - start_time)
    print('R2 scores:', score_r2_P1)
    print('SROCC scores:', score_srocc_P1)
    print('PLCC scores:', score_plcc_P1)