RT-BENE model evaluation (#73)

* Add evaluation for RT-BENE

* Ignore some more models

* Add testing code description to README

Co-authored-by: Kevin Cortacero <cortacero.kevin@gmail.com>

.gitignore

@@ -20,6 +20,9 @@ rt_gene/model_nets/all_subjects_mpii_prl_utmv_2_02.h5
 rt_gene/model_nets/all_subjects_mpii_prl_utmv_3_02.h5
 rt_gene/model_nets/blink_model_1.h5
 rt_gene/model_nets/blink_model_2.h5
+rt_gene/model_nets/rt-bene_mobilenetv2_fold1_best.h5
+rt_gene/model_nets/rt-bene_mobilenetv2_fold2_best.h5
+rt_gene/model_nets/rt-bene_mobilenetv2_fold3_best.h5
 rt_gene/model_nets/Model_allsubjects1_pytorch.model
 rt_gene/model_nets/Model_allsubjects2_pytorch.model
 rt_gene/model_nets/Model_allsubjects3_pytorch.model

rt_bene_model_training/README.md

@@ -31,7 +31,7 @@ For pip users: `pip install tensorflow-gpu numpy tqdm opencv-python scikit-learn
 This code was used to train the blink estimator for RT-BENE. The labels for the RT-BENE blink dataset are contained in the [rt_bene_dataset](../rt_bene_dataset) directory. The images corresponding to the labels can be downloaded from the RT-GENE dataset (labels are only available for the "noglasses" part): [download](https://zenodo.org/record/2529036) [(alternative link)](https://goo.gl/tfUaDm). Please run `python train_blink_model.py --help` to see the required arguments to train the model.
 
 ## Model testing code
-The evaluation code will be provided soon; please check back.
+Evaluation code for a 3-fold evaluation is provided in the [evaluate_blink_model.py](./evaluate_blink_model.py) file. An example to train and evaluate an ensemble of models can be found in [train_and_evaluate.py](./train_and_evaluate.py). Please run `python train_and_evaluate.py --help` to see the required arguments.
 
 ![Results](../rt_bene_precision_recall.png)
 

rt_bene_model_training/evaluate_blink_model.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python
+
+import gc
+
+import tensorflow as tf
+from tensorflow.keras.models import load_model
+
+from sklearn.metrics import confusion_matrix, roc_curve, auc, average_precision_score
+
+import numpy as np
+
+tf.compat.v1.disable_eager_execution()
+
+config = tf.compat.v1.ConfigProto()
+config.gpu_options.allow_growth = True
+tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config))
+
+
+fold_infos = {
+    'fold1': [2],
+    'fold2': [1],
+    'fold3': [0],
+    'all': [2, 1, 0]
+}
+
+model_metrics = [tf.keras.metrics.BinaryAccuracy()]
+
+
+def estimate_metrics(testing_fold, model_instance):
+    threshold = 0.5
+    p = model_instance.predict(x=testing_fold['x'], verbose=0)
+    p = p >= threshold
+    matrix = confusion_matrix(testing_fold['y'], p)
+    ap = average_precision_score(testing_fold['y'], p)
+    fpr, tpr, thresholds = roc_curve(testing_fold['y'], p)
+    roc = auc(fpr, tpr)
+    return matrix, ap, roc
+
+
+def get_metrics_from_matrix(matrix):
+    tp, tn, fp, fn = matrix[1, 1], matrix[0, 0], matrix[0, 1], matrix[1, 0]
+    precision = tp / (tp + fp)
+    recall = tp / (tp + fn)
+    f1score = 2. * (precision * recall) / (precision + recall)
+    return precision, recall, f1score
+
+
+def threefold_evaluation(dataset, model_paths_fold1, model_paths_fold2, model_paths_fold3, input_size):
+    folds = ['fold1', 'fold2', 'fold3']
+    aps = []
+    rocs = []
+    recalls = []
+    precisions = []
+    f1scores = []
+    models = []
+
+    for fold_to_eval_on, model_paths in zip(folds, [model_paths_fold1, model_paths_fold2, model_paths_fold3]):
+        if len(model_paths_fold1) > 1:
+            models = [load_model(model_path, compile=False) for model_path in model_paths]
+            img_input_l = tf.keras.Input(shape=input_size, name='img_input_L')
+            img_input_r = tf.keras.Input(shape=input_size, name='img_input_R')
+            tensors = [model([img_input_r, img_input_l]) for model in models]
+            output_layer = tf.keras.layers.average(tensors)
+            model_instance = tf.keras.Model(inputs=[img_input_r, img_input_l], outputs=output_layer)
+        else:
+            model_instance = load_model(model_paths[0])
+        model_instance.compile()
+
+        testing_fold = dataset.get_training_data(fold_infos[fold_to_eval_on])  # get the testing fold subjects
+
+        matrix, ap, roc = estimate_metrics(testing_fold, model_instance)
+        aps.append(ap)
+        rocs.append(roc)
+        precision, recall, f1score = get_metrics_from_matrix(matrix)
+        recalls.append(recall)
+        precisions.append(precision)
+        f1scores.append(f1score)
+
+        del model_instance, testing_fold
+        # noinspection PyUnusedLocal
+        for model in models:
+            del model
+        gc.collect()
+
+    evaluation = {'AP': {}, 'ROC': {}, 'precision': {}, 'recall': {}, 'f1score': {}}
+    evaluation['AP']['avg'] = np.mean(np.array(aps))
+    evaluation['AP']['std'] = np.std(np.array(aps))
+    evaluation['ROC']['avg'] = np.mean(np.array(rocs))
+    evaluation['ROC']['std'] = np.std(np.array(rocs))
+    evaluation['precision']['avg'] = np.mean(np.array(precisions))
+    evaluation['precision']['std'] = np.std(np.array(precisions))
+    evaluation['recall']['avg'] = np.mean(np.array(recalls))
+    evaluation['recall']['std'] = np.std(np.array(recalls))
+    evaluation['f1score']['avg'] = np.mean(np.array(f1scores))
+    evaluation['f1score']['std'] = np.std(np.array(f1scores))
+    return evaluation

rt_bene_model_training/train_and_evaluate.py

@@ -0,0 +1,52 @@
+from evaluate_blink_model import threefold_evaluation
+from train_blink_model import ThreefoldTraining
+from dataset_manager import RTBeneDataset
+from pathlib import Path
+import argparse
+import pprint
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("model_save_root", help="target folder to save the models (auto-saved)")
+    parser.add_argument("csv_subject_list", help="path to the dataset csv file")
+    parser.add_argument("--ensemble_size", type=int, default=1, help="number of models to train for the ensemble")
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--epochs", type=int, default=15)
+    parser.add_argument("--input_size", type=tuple, help="input size of images", default=(96, 96))
+
+    args = parser.parse_args()
+
+    fold_list = ['fold1', 'fold2', 'fold3']
+    ensemble_size = args.ensemble_size  # 1 is considered as single model
+    epochs = args.epochs
+    batch_size = args.batch_size
+    input_size = args.input_size
+    csv_subject_list = args.csv_subject_list
+    model_save_root = args.model_save_root
+
+    dataset = RTBeneDataset(csv_subject_list, input_size)
+
+    threefold_training = ThreefoldTraining(dataset, epochs, batch_size, input_size)
+
+    all_evaluations = {}
+
+    for backbone in ['densenet121', 'resnet50', 'mobilenetv2']:
+        models_fold1 = []
+        models_fold2 = []
+        models_fold3 = []
+
+        for i in range(1, ensemble_size + 1):
+            model_save_path = Path(model_save_root + backbone + '/' + str(i))
+            model_save_path.mkdir(parents=True, exist_ok=True)
+            threefold_training.train(backbone, str(model_save_path) + '/')
+
+            models_fold1.append(str(model_save_path) + '/rt-bene_' + backbone + '_fold1_best.h5')
+            models_fold2.append(str(model_save_path) + '/rt-bene_' + backbone + '_fold2_best.h5')
+            models_fold3.append(str(model_save_path) + '/rt-bene_' + backbone + '_fold3_best.h5')
+
+        evaluation = threefold_evaluation(dataset, models_fold1, models_fold2, models_fold3, input_size)
+        all_evaluations[backbone] = evaluation
+
+    threefold_training.free()
+
+    pprint.pprint(all_evaluations)

rt_bene_model_training/train_blink_model.py

@@ -55,75 +55,85 @@ def create_model(backbone, input_shape, lr, metrics):
     base = create_model_base(backbone, input_shape)
 
     # define the 2 inputs (left and right eyes)
-    left_input = Input(shape=input_shape)
+    left_input  = Input(shape=input_shape)
     right_input = Input(shape=input_shape)
 
     # get the 2 outputs using shared layers
-    out_left = base(left_input)
+    out_left  = base(left_input)
     out_right = base(right_input)
 
     # average the predictions
     merged = Average()([out_left, out_right])
-    model = Model(inputs=[right_input, left_input], outputs=merged)
+    model  = Model(inputs=[right_input, left_input], outputs=merged)
 
     model.compile(loss='binary_crossentropy', optimizer=Adam(lr=lr), metrics=metrics)
     model.summary()
 
     return model
 
 
+class ThreefoldTraining(object):
+    def __init__(self, dataset, epochs, batch_size, input_size):
+        self.fold_map = {'fold1': [0, 1], 'fold2': [0, 2], 'fold3': [1, 2]}
+        self.model_metrics = [tf.keras.metrics.BinaryAccuracy(), tf.keras.metrics.Recall(), tf.keras.metrics.Precision()]
+        self.dataset = dataset
+        self.validation_set = dataset.get_validation_data()
+        self.epochs = epochs
+        self.batch_size = batch_size
+        self.input_size = [input_size[0], input_size[1], 3]
+        self.learning_rate = 1e-4
+
+    def train(self, backbone, model_save_path):
+        for fold_name, training_subjects_fold in self.fold_map.items():
+            training_set = self.dataset.get_training_data(training_subjects_fold)
+            positive = training_set['positive']
+            negative = training_set['negative']
+
+            print('Number of positive samples in training data: {} ({:.2f}% of total)'.
+                  format(positive, 100 * float(positive) / len(training_set['y'])))
+
+            model_instance = create_model(backbone, self.input_size, self.learning_rate, self.model_metrics)
+            name = 'rt-bene_' + backbone + '_' + fold_name
+
+            weight_for_0 = 1. / negative * (negative + positive)
+            weight_for_1 = 1. / positive * (negative + positive)
+            class_weight = {0: weight_for_0, 1: weight_for_1}
+
+            save_best = ModelCheckpoint(model_save_path + name + '_best.h5', monitor='val_loss', verbose=1,
+                                        save_best_only=True, save_weights_only=False, mode='min', period=1)
+            auto_save = ModelCheckpoint(model_save_path + name + '_auto_{epoch:02d}.h5', verbose=1,
+                                        save_best_only=False, save_weights_only=False, period=1)
+
+            # train the model
+            model_instance.fit(x=training_set['x'], y=training_set['y'],
+                               batch_size=self.batch_size, epochs=self.epochs,
+                               verbose=1,
+                               validation_data=(self.validation_set['x'], self.validation_set['y']),
+                               callbacks=[save_best, auto_save],
+                               class_weight=class_weight)
+            # noinspection PyUnusedLocal
+            model_instance, training_set = None, None
+            del model_instance, training_set
+            gc.collect()
+
+    def free(self):
+        self.validation_set = None
+        del self.validation_set
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("model_base", choices=['densenet121', 'resnet50', 'mobilenetv2'])
-    parser.add_argument("model_path", help="target folder to save the models (auto-saved)")
+    parser.add_argument("backbone", choices=['densenet121', 'resnet50', 'mobilenetv2'])
+    parser.add_argument("model_save_path", help="target folder to save the models (auto-saved)")
     parser.add_argument("csv_subject_list", help="path to the dataset csv file")
     parser.add_argument("--batch_size", type=int, default=64)
-    parser.add_argument("--epochs", type=int, default=8)
+    parser.add_argument("--epochs", type=int, default=15)
     parser.add_argument("--input_size", type=tuple, help="input size of images", default=(96, 96))
 
     args = parser.parse_args()
-    model_base = args.model_base
-    epochs = args.epochs
-    batch_size = args.batch_size
-    input_size = args.input_size
-
-    dataset = RTBeneDataset(args.csv_subject_list, input_size)
-    fold_infos = [
-        ([0, 1], 'fold1'),
-        ([0, 2], 'fold2'),
-        ([1, 2], 'fold3')
-    ]
-
-    validation_set = dataset.get_validation_data()
-
-    # 3 folds training
-    for subjects_train, fold_name in fold_infos:
-        training_fold = dataset.get_training_data(subjects_train)
-        positive = training_fold['positive']
-        negative = training_fold['negative']
-
-        print('Number of positive samples in training data: {} ({:.2f}% of total)'.
-              format(positive, 100 * float(positive) / len(training_fold['y'])))
-
-        model_metrics = [tf.keras.metrics.BinaryAccuracy(), tf.keras.metrics.Recall(), tf.keras.metrics.Precision()]
-        model_instance = create_model(args.model_base, [input_size[0], input_size[1], 3], 1e-4, model_metrics)
-        name = 'rt-bene_' + args.model_base + '_' + fold_name
-
-        weight_for_0 = 1. / negative * (negative + positive)
-        weight_for_1 = 1. / positive * (negative + positive)
-        class_weight = {0: weight_for_0, 1: weight_for_1}
-
-        save_best = ModelCheckpoint(args.model_path + name + '_best.h5', monitor='val_loss', verbose=1, 
-                                    save_best_only=True, save_weights_only=False, mode='min', period=1)
-        auto_save = ModelCheckpoint(args.model_path + name + '_auto_{epoch:02d}.h5', verbose=1,
-                                    save_best_only=False, save_weights_only=False, period=1)
-
-        # train the model
-        model_instance.fit(x=training_fold['x'], y=training_fold['y'], batch_size=batch_size, epochs=epochs, verbose=1,
-                           validation_data=(validation_set['x'], validation_set['y']), callbacks=[save_best, auto_save],
-                           class_weight=class_weight)
-        model_instance, training_fold = None, None
-        del model_instance, training_fold
-        gc.collect()
-    validation_set = None
-    del validation_set
+
+    rtbene_dataset = RTBeneDataset(args.csv_subject_list, args.input_size)
+
+    threefold_training = ThreefoldTraining(rtbene_dataset, args.epochs, args.batch_size, args.input_size)
+    threefold_training.train(args.backbone, args.model_save_path + '/')
+    threefold_training.free()