diff --git a/.gitignore b/.gitignore index 9f0b243..b2762ac 100644 --- a/.gitignore +++ b/.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 diff --git a/rt_bene_model_training/README.md b/rt_bene_model_training/README.md index 507e23e..3b1c201 100644 --- a/rt_bene_model_training/README.md +++ b/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) diff --git a/rt_bene_model_training/evaluate_blink_model.py b/rt_bene_model_training/evaluate_blink_model.py new file mode 100644 index 0000000..8d3bfcf --- /dev/null +++ b/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 diff --git a/rt_bene_model_training/train_and_evaluate.py b/rt_bene_model_training/train_and_evaluate.py new file mode 100644 index 0000000..6464b29 --- /dev/null +++ b/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) diff --git a/rt_bene_model_training/train_blink_model.py b/rt_bene_model_training/train_blink_model.py old mode 100755 new mode 100644 index 02092d2..8b42a42 --- a/rt_bene_model_training/train_blink_model.py +++ b/rt_bene_model_training/train_blink_model.py @@ -55,16 +55,16 @@ 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() @@ -72,58 +72,68 @@ def create_model(backbone, input_shape, lr, metrics): 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()