added pytorch training

moved tensorflow to subdirectory

Author: ahmed-alhindawi

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
-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.
+Evaluation code for a 3-fold evaluation is provided in the [evaluate_blink_model.py](tensorflow/evaluate_blink_model.py) file. An example to train and evaluate an ensemble of models can be found in [train_and_evaluate.py](tensorflow/train_and_evaluate.py). Please run `python train_and_evaluate.py --help` to see the required arguments.
 
 ![Results](../assets/rt_bene_precision_recall.png)
 

rt_bene_model_training/__init__.py

@@ -0,0 +1,72 @@
+import os
+
+import numpy as np
+from PIL import Image
+from torch.utils import data
+from torchvision import transforms
+from tqdm import tqdm
+
+
+class RTBENEH5Dataset(data.Dataset):
+
+    def __init__(self, h5_file, subject_list=None, transform=None, loader_desc="train"):
+        self._h5_file = h5_file
+        self._transform = transform
+        self._subject_labels = []
+        self._positive_labels = 0
+        self._total_labels = 0
+
+        assert subject_list is not None, "Must pass a list of subjects to load the data for"
+
+        if self._transform is None:
+            self._transform = transforms.Compose([transforms.Resize((224, 224), Image.BICUBIC),
+                                                  transforms.ToTensor(),
+                                                  transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                                                       std=[0.229, 0.224, 0.225])])
+
+        _wanted_subjects = ["s{:03d}".format(_i) for _i in subject_list]
+
+        for grp_s_n in tqdm(_wanted_subjects, desc="Loading ({}) subject metadata...".format(loader_desc)):  # subjects
+            for grp_i_n, grp_i in h5_file[grp_s_n].items():  # images
+                if "image" in grp_i.keys() and "label" in grp_i.keys():
+                    image_dataset = grp_i["image"]
+                    label = grp_i["label"]
+                    if label == 1.0:
+                        self._positive_labels = self._positive_labels + 1
+                    self._total_labels = self._total_labels + 1
+
+                    for _i in range(len(image_dataset)):
+                        self._subject_labels.append(["/" + grp_s_n + "/" + grp_i_n, _i])
+
+    @staticmethod
+    def get_class_weights(h5_file, subject_list):
+        positive = 0
+        total = 0
+        _wanted_subjects = ["s{:03d}".format(_i) for _i in subject_list]
+
+        for grp_s_n in tqdm(_wanted_subjects, desc="Loading class weights..."):
+            for grp_i_n, grp_i in h5_file[grp_s_n].items():  # images
+                if "image" in grp_i.keys() and "label" in grp_i.keys():
+                    label = grp_i["label"][()][0]
+                    if label == 1.0:
+                        positive = positive + 1
+                    total = total + 1
+
+        negative = total - positive
+        weight_for_0 = (negative + positive) / negative
+        weight_for_1 = (negative + positive) / positive
+        return {0: weight_for_0, 1: weight_for_1}
+
+    def __len__(self):
+        return len(self._subject_labels)
+
+    def __getitem__(self, index):
+        _sample = self._subject_labels[index]
+        assert type(_sample[0]) == str, "Sample not found at index {}".format(index)
+        _img = self._h5_file[_sample[0] + "/image"][_sample[1]][()]
+        label_data = self._h5_file[_sample[0] + "/label"][()].astype(np.float32)
+
+        # Load data and get label
+        _transformed_img = self._transform(Image.fromarray(_img, 'RGB'))
+
+        return _transformed_img, label_data

rt_bene_model_training/pytorch/__init__.py

@@ -0,0 +1,75 @@
+from __future__ import print_function, division, absolute_import
+
+import argparse
+import os
+
+import h5py
+import numpy as np
+from PIL import Image, ImageFilter, ImageOps
+from torchvision import transforms
+from tqdm import tqdm
+
+script_path = os.path.dirname(os.path.realpath(__file__))
+
+# Augmentations following `prepare_dataset.m`: randomly crop and resize the image 10 times,
+# along side two blurring stages, grayscaling and histogram normalisation
+_required_size = (224, 224)
+_transforms_list = [transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),  # equivalent to random 5px from each edge
+                    transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),
+                    transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),
+                    transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),
+                    transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),
+                    transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),
+                    transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),
+                    transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),
+                    transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),
+                    transforms.RandomResizedCrop(size=_required_size, scale=(0.85, 1.0)),
+                    transforms.Grayscale(num_output_channels=3),
+                    lambda x: x.filter(ImageFilter.GaussianBlur(radius=1)),
+                    lambda x: x.filter(ImageFilter.GaussianBlur(radius=3)),
+                    lambda x: ImageOps.equalize(x)]  # histogram equalisation
+
+
+def load_and_augment(file_path, augment=False):
+    image = Image.open(file_path).resize(_required_size)
+    augmented_images = [np.array(trans(image)) for trans in _transforms_list if augment is True]
+    augmented_images.append(np.array(image))
+
+    return np.array(augmented_images, dtype=np.uint8)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Estimate gaze from images')
+    parser.add_argument('--rt_bene_root', type=str, required=True, nargs='?', help='Path to the base directory of RT_GENE')
+    parser.add_argument('--augment_dataset', type=bool, required=False, default=False, help="Whether to augment the dataset with predefined transforms")
+    parser.add_argument('--compress', action='store_true', dest="compress")
+    parser.add_argument('--no-compress', action='store_false', dest="compress")
+    parser.set_defaults(compress=False)
+    args = parser.parse_args()
+
+    _compression = "lzf" if args.compress is True else None
+
+    subject_path = [os.path.join(args.rt_bene_root, "s{:03d}_noglasses/".format(_i)) for _i in range(0, 17)]
+
+    hdf_file = h5py.File(os.path.abspath(os.path.join(args.rt_bene_root, "rtbene_dataset.hdf5")), mode='w')
+    for subject_id, subject_data in enumerate(subject_path):
+        subject_id = str("s{:03d}".format(subject_id))
+        subject_grp = hdf_file.create_group(subject_id)
+        with open(os.path.join(args.rt_bene_root, "{}_blink_labels.csv".format(subject_id)), "r") as f:
+            _lines = f.readlines()
+
+            for line in tqdm(_lines, desc="Subject {}".format(subject_id)):
+
+                split = line.split(",")
+                image_name = split[0]
+                image_grp = subject_grp.create_group(image_name)
+                image_path = os.path.join(subject_data, "natural/left/", "{}".format(split[0]))
+                if os.path.exists(image_path):
+                    label = float(split[1].strip("\n"))
+                    if label != 0.5:  # paper removed 0.5s
+                        image_data = load_and_augment(image_path, augment=args.augment_dataset)
+                        image_grp.create_dataset("image", data=image_data, compression=_compression)
+                        image_grp.create_dataset("label", data=[label])
+
+    hdf_file.flush()
+    hdf_file.close()

rt_bene_model_training/pytorch/rtbene_dataset.py

@@ -1,96 +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
+#!/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