Initial commit

DeepDeformationMapRegistration/data_generator.py

@@ -0,0 +1,497 @@
+import sys, os
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+
+import numpy as np
+from tensorflow import keras
+import os
+import h5py
+import random
+from PIL import Image
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.operators import min_max_norm
+
+
+class DataGeneratorManager(keras.utils.Sequence):
+    def __init__(self, dataset_path, batch_size=32, shuffle=True, num_samples=None, validation_split=None, validation_samples=None,
+                 clip_range=[0., 1.], voxelmorph=False, segmentations=False,
+                 seg_labels: dict = {'bg': 0, 'vessels': 1, 'tumour': 2, 'parenchyma': 3}):
+        # Get the list of files
+        self.__list_files = self.__get_dataset_files(dataset_path)
+        self.__list_files.sort()
+        self.__dataset_path = dataset_path
+        self.__shuffle = shuffle
+        self.__total_samples = len(self.__list_files)
+        self.__validation_split = validation_split
+        self.__clip_range = clip_range
+        self.__batch_size = batch_size
+
+        self.__validation_samples = validation_samples
+
+        self.__voxelmorph = voxelmorph
+        self.__segmentations = segmentations
+        self.__seg_labels = seg_labels
+
+        if num_samples is not None:
+            self.__num_samples = self.__total_samples if num_samples > self.__total_samples else num_samples
+        else:
+            self.__num_samples = self.__total_samples
+
+        self.__internal_idxs = np.arange(self.__num_samples)
+
+        # Split it accordingly
+        if validation_split is None:
+            self.__validation_num_samples = None
+            self.__validation_idxs = list()
+            if self.__shuffle:
+                random.shuffle(self.__internal_idxs)
+            self.__training_idxs = self.__internal_idxs
+
+            self.__validation_generator = None
+        else:
+            self.__validation_num_samples = int(np.ceil(self.__num_samples * validation_split))
+            if self.__shuffle:
+                self.__validation_idxs = np.random.choice(self.__internal_idxs, self.__validation_num_samples)
+            else:
+                self.__validation_idxs = self.__internal_idxs[0: self.__validation_num_samples]
+            self.__training_idxs = np.asarray([idx for idx in self.__internal_idxs if idx not in self.__validation_idxs])
+            # Build them DataGenerators
+            self.__validation_generator = DataGenerator(self, 'validation')
+
+        self.__train_generator = DataGenerator(self, 'train')
+        self.reshuffle_indices()
+
+    @property
+    def dataset_path(self):
+        return self.__dataset_path
+
+    @property
+    def dataset_list_files(self):
+        return self.__list_files
+
+    @property
+    def train_idxs(self):
+        return self.__training_idxs
+
+    @property
+    def validation_idxs(self):
+        return self.__validation_idxs
+
+    @property
+    def batch_size(self):
+        return self.__batch_size
+
+    @property
+    def clip_rage(self):
+        return self.__clip_range
+
+    @property
+    def shuffle(self):
+        return self.__shuffle
+
+    def get_generator_idxs(self, generator_type):
+        if generator_type == 'train':
+            return self.train_idxs
+        elif generator_type == 'validation':
+            return self.validation_idxs
+        else:
+            raise ValueError('Invalid generator type: ', generator_type)
+
+    @staticmethod
+    def __get_dataset_files(search_path):
+        """
+        Get the path to the dataset files
+        :param  search_path: dir path to search for the hd5 files
+        :return:
+        """
+        file_list = list()
+        for root, dirs, files in os.walk(search_path):
+            file_list.sort()
+            for data_file in files:
+                file_name, extension = os.path.splitext(data_file)
+                if extension.lower() == '.hd5':
+                    file_list.append(os.path.join(root, data_file))
+
+        if not file_list:
+            raise ValueError('No files found to train in ', search_path)
+
+        print('Found {} files in {}'.format(len(file_list), search_path))
+        return file_list
+
+    def reshuffle_indices(self):
+        if self.__validation_num_samples is None:
+            if self.__shuffle:
+                random.shuffle(self.__internal_idxs)
+            self.__training_idxs = self.__internal_idxs
+        else:
+            if self.__shuffle:
+                self.__validation_idxs = np.random.choice(self.__internal_idxs, self.__validation_num_samples)
+            else:
+                self.__validation_idxs = self.__internal_idxs[0: self.__validation_num_samples]
+            self.__training_idxs = np.asarray([idx for idx in self.__internal_idxs if idx not in self.__validation_idxs])
+
+            # Update the indices
+            self.__validation_generator.update_samples(self.__validation_idxs)
+
+        self.__train_generator.update_samples(self.__training_idxs)
+
+    def get_generator(self, type='train'):
+        if type.lower() == 'train':
+            return self.__train_generator
+        elif type.lower() == 'validation':
+            if self.__validation_generator is not None:
+                return self.__validation_generator
+            else:
+                raise Warning('No validation generator available. Set a non-zero validation_split to build one.')
+        else:
+            raise ValueError('Unknown dataset type "{}". Expected "train" or "validation"'.format(type))
+
+    @property
+    def is_voxelmorph(self):
+        return self.__voxelmorph
+
+    @property
+    def give_segmentations(self):
+        return self.__segmentations
+
+    @property
+    def seg_labels(self):
+        return self.__seg_labels
+
+
+class DataGenerator(DataGeneratorManager):
+    def __init__(self, GeneratorManager: DataGeneratorManager, dataset_type='train'):
+        self.__complete_list_files = GeneratorManager.dataset_list_files
+        self.__list_files = [self.__complete_list_files[idx] for idx in GeneratorManager.get_generator_idxs(dataset_type)]
+        self.__batch_size = GeneratorManager.batch_size
+        self.__total_samples = len(self.__list_files)
+        self.__clip_range = GeneratorManager.clip_rage
+        self.__manager = GeneratorManager
+        self.__shuffle = GeneratorManager.shuffle
+
+        self.__seg_labels = GeneratorManager.seg_labels
+
+        self.__num_samples = len(self.__list_files)
+        self.__internal_idxs = np.arange(self.__num_samples)
+        # These indices are internal to the generator, they are not the same as the dataset_idxs!!
+
+        self.__dataset_type = dataset_type
+
+        self.__last_batch = 0
+        self.__batches_per_epoch = int(np.floor(len(self.__internal_idxs) / self.__batch_size))
+
+        self.__voxelmorph = GeneratorManager.is_voxelmorph
+        self.__segmentations = GeneratorManager.is_voxelmorph and GeneratorManager.give_segmentations
+
+    @staticmethod
+    def __get_dataset_files(search_path):
+        """
+        Get the path to the dataset files
+        :param  search_path: dir path to search for the hd5 files
+        :return:
+        """
+        file_list = list()
+        for root, dirs, files in os.walk(search_path):
+            for data_file in files:
+                file_name, extension = os.path.splitext(data_file)
+                if extension.lower() == '.hd5':
+                    file_list.append(os.path.join(root, data_file))
+
+        if not file_list:
+            raise ValueError('No files found to train in ', search_path)
+
+        print('Found {} files in {}'.format(len(file_list), search_path))
+        return file_list
+
+    def update_samples(self, new_sample_idxs):
+        self.__list_files = [self.__complete_list_files[idx] for idx in new_sample_idxs]
+        self.__num_samples = len(self.__list_files)
+        self.__internal_idxs = np.arange(self.__num_samples)
+
+    def on_epoch_end(self):
+        """
+        To be executed at the end of each epoch. Reshuffle the assigned samples
+        :return:
+        """
+        if self.__shuffle:
+            random.shuffle(self.__internal_idxs)
+        self.__last_batch = 0
+
+    def __len__(self):
+        """
+        Number of batches per epoch
+        :return:
+        """
+        return self.__batches_per_epoch
+
+    def __getitem__(self, index):
+        """
+        Generate one batch of data
+        :param index: epoch index
+        :return:
+        """
+        idxs = self.__internal_idxs[index * self.__batch_size:(index + 1) * self.__batch_size]
+
+        fix_img, mov_img, fix_vessels, mov_vessels, fix_tumour, mov_tumour, disp_map = self.__load_data(idxs)
+
+        try:
+            fix_img = min_max_norm(fix_img).astype(np.float32)
+            mov_img = min_max_norm(mov_img).astype(np.float32)
+        except ValueError:
+            print(idxs, fix_img.shape, mov_img.shape)
+            er_str = 'ERROR:\t[file]:\t{}\t[idx]:\t{}\t[fix_img.shape]:\t{}\t[mov_img.shape]:\t{}\t'.format(self.__list_files[idxs], idxs, fix_img.shape, mov_img.shape)
+            raise ValueError(er_str)
+
+        fix_vessels[fix_vessels > 0.] = self.__seg_labels['vessels']
+        mov_vessels[mov_vessels > 0.] = self.__seg_labels['vessels']
+
+        # fix_tumour[fix_tumour > 0.] = self.__seg_labels['tumour']
+        # mov_tumour[mov_tumour > 0.] = self.__seg_labels['tumour']
+
+        # https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit
+        # A generator or keras.utils.Sequence returning (inputs, targets) or (inputs, targets, sample_weights)
+        # The second element must match the outputs of the model, in this case (image, displacement map)
+        if self.__voxelmorph:
+            zero_grad = np.zeros([fix_img.shape[0], *C.DISP_MAP_SHAPE])
+            if self.__segmentations:
+                inputs = [mov_vessels, fix_vessels, mov_img, fix_img, zero_grad]
+                outputs = []  #[fix_img, zero_grad]
+            else:
+                inputs = [mov_img, fix_img]
+                outputs = [fix_img, zero_grad]
+            return (inputs, outputs)
+        else:
+            return (fix_img, mov_img, fix_vessels, mov_vessels), # (None, fix_seg, fix_seg, fix_img)
+
+    def next_batch(self):
+        if self.__last_batch > self.__batches_per_epoch:
+            raise ValueError('No more batches for this epoch')
+        batch = self.__getitem__(self.__last_batch)
+        self.__last_batch += 1
+        return batch
+
+    def __load_data(self, idx_list):
+        """
+        Build the batch with the samples in idx_list
+        :param idx_list:
+        :return:
+        """
+        if isinstance(idx_list, (list, np.ndarray)):
+            fix_img = np.empty((0, ) + C.IMG_SHAPE)
+            mov_img = np.empty((0, ) + C.IMG_SHAPE)
+            disp_map = np.empty((0, ) + C.DISP_MAP_SHAPE)
+
+            # fix_segm = np.empty((0, ) + const.IMG_SHAPE)
+            # mov_segm = np.empty((0, ) + const.IMG_SHAPE)
+
+            fix_vessels = np.empty((0, ) + C.IMG_SHAPE)
+            mov_vessels = np.empty((0, ) + C.IMG_SHAPE)
+            fix_tumors = np.empty((0, ) + C.IMG_SHAPE)
+            mov_tumors = np.empty((0, ) + C.IMG_SHAPE)
+            for idx in idx_list:
+                data_file = h5py.File(self.__list_files[idx], 'r')
+
+                fix_img = np.append(fix_img, [data_file[C.H5_FIX_IMG][:]], axis=0)
+                mov_img = np.append(mov_img, [data_file[C.H5_MOV_IMG][:]], axis=0)
+
+                # fix_segm = np.append(fix_segm, [data_file[const.H5_FIX_PARENCHYMA_MASK][:]], axis=0)
+                # mov_segm = np.append(mov_segm, [data_file[const.H5_MOV_PARENCHYMA_MASK][:]], axis=0)
+
+                disp_map = np.append(disp_map, [data_file[C.H5_GT_DISP][:]], axis=0)
+
+                fix_vessels = np.append(fix_vessels, [data_file[C.H5_FIX_VESSELS_MASK][:]], axis=0)
+                mov_vessels = np.append(mov_vessels, [data_file[C.H5_MOV_VESSELS_MASK][:]], axis=0)
+                fix_tumors = np.append(fix_tumors, [data_file[C.H5_FIX_TUMORS_MASK][:]], axis=0)
+                mov_tumors = np.append(mov_tumors, [data_file[C.H5_MOV_TUMORS_MASK][:]], axis=0)
+
+                data_file.close()
+        else:
+            data_file = h5py.File(self.__list_files[idx_list], 'r')
+
+            fix_img = np.expand_dims(data_file[C.H5_FIX_IMG][:], 0)
+            mov_img = np.expand_dims(data_file[C.H5_MOV_IMG][:], 0)
+
+            # fix_segm = np.expand_dims(data_file[const.H5_FIX_PARENCHYMA_MASK][:], 0)
+            # mov_segm = np.expand_dims(data_file[const.H5_MOV_PARENCHYMA_MASK][:], 0)
+
+            fix_vessels = np.expand_dims(data_file[C.H5_FIX_VESSELS_MASK][:], axis=0)
+            mov_vessels = np.expand_dims(data_file[C.H5_MOV_VESSELS_MASK][:], axis=0)
+            fix_tumors = np.expand_dims(data_file[C.H5_FIX_TUMORS_MASK][:], axis=0)
+            mov_tumors = np.expand_dims(data_file[C.H5_MOV_TUMORS_MASK][:], axis=0)
+
+            disp_map = np.expand_dims(data_file[C.H5_GT_DISP][:], 0)
+
+            data_file.close()
+
+        return fix_img, mov_img, fix_vessels, mov_vessels, fix_tumors, mov_tumors, disp_map
+
+    def get_single_sample(self):
+        fix_img, mov_img, fix_segm, mov_segm, _ = self.__load_data(0)
+        # return X, y
+        return np.expand_dims(np.concatenate([mov_img, fix_img, mov_segm, mov_segm], axis=-1), axis=0)
+
+    def get_random_sample(self, num_samples):
+        idxs = np.random.randint(0, self.__num_samples, num_samples)
+        fix_img, mov_img, fix_segm, mov_segm, disp_map = self.__load_data(idxs)
+
+        return (fix_img, mov_img, fix_segm, mov_segm, disp_map), [self.__list_files[f] for f in idxs]
+
+    def get_input_shape(self):
+        input_batch, _ = self.__getitem__(0)
+        if self.__voxelmorph:
+            ret_val = list(input_batch[0].shape)
+            ret_val[-1] = 2
+            ret_val = (None, ) + tuple(ret_val[1:])
+        else:
+            ret_val = input_batch.shape
+            ret_val = (None, ) + ret_val[1:]
+        return ret_val # const.BATCH_SHAPE_SEGM
+
+    def who_are_you(self):
+        return self.__dataset_type
+
+    def print_datafiles(self):
+        return self.__list_files
+
+
+class DataGeneratorManager2D:
+    FIX_IMG_H5 = 'input/1'
+    MOV_IMG_H5 = 'input/0'
+    def __init__(self, h5_file_list, batch_size=32, data_split=0.7, img_size=None,
+                 fix_img_tag=FIX_IMG_H5, mov_img_tag=MOV_IMG_H5, multi_loss=False):
+        self.__file_list = h5_file_list #h5py.File(h5_file, 'r')
+        self.__batch_size = batch_size
+        self.__data_split = data_split
+
+        self.__initialize()
+
+        self.__train_generator = DataGenerator2D(self.__train_file_list,
+                                                 batch_size=self.__batch_size,
+                                                 img_size=img_size,
+                                                 fix_img_tag=fix_img_tag,
+                                                 mov_img_tag=mov_img_tag,
+                                                 multi_loss=multi_loss)
+        self.__val_generator = DataGenerator2D(self.__val_file_list,
+                                               batch_size=self.__batch_size,
+                                               img_size=img_size,
+                                               fix_img_tag=fix_img_tag,
+                                               mov_img_tag=mov_img_tag,
+                                               multi_loss=multi_loss)
+
+    def __initialize(self):
+        num_samples = len(self.__file_list)
+        random.shuffle(self.__file_list)
+
+        data_split = int(np.floor(num_samples * self.__data_split))
+        self.__val_file_list = self.__file_list[0:data_split]
+        self.__train_file_list = self.__file_list[data_split:]
+
+    @property
+    def train_generator(self):
+        return self.__train_generator
+
+    @property
+    def validation_generator(self):
+        return self.__val_generator
+
+
+class DataGenerator2D(keras.utils.Sequence):
+    FIX_IMG_H5 = 'input/1'
+    MOV_IMG_H5 = 'input/0'
+
+    def __init__(self, file_list: list, batch_size=32, img_size=None, fix_img_tag=FIX_IMG_H5, mov_img_tag=MOV_IMG_H5, multi_loss=False):
+        self.__file_list = file_list  # h5py.File(h5_file, 'r')
+        self.__file_list.sort()
+        self.__batch_size = batch_size
+        self.__idx_list = np.arange(0, len(self.__file_list))
+        self.__multi_loss = multi_loss
+
+        self.__tags = {'fix_img': fix_img_tag,
+                       'mov_img': mov_img_tag}
+
+        self.__batches_seen = 0
+        self.__batches_per_epoch = 0
+
+        self.__img_size = img_size
+
+        self.__initialize()
+
+    def __len__(self):
+        return self.__batches_per_epoch
+
+    def __initialize(self):
+        random.shuffle(self.__idx_list)
+
+        if self.__img_size is None:
+            f = h5py.File(self.__file_list[0], 'r')
+            self.input_shape = f[self.__tags['fix_img']].shape  # Already defined in super()
+            f.close()
+        else:
+            self.input_shape = self.__img_size
+
+        if self.__multi_loss:
+            self.input_shape = (self.input_shape, (*self.input_shape[:-1], 2))
+
+        self.__batches_per_epoch = int(np.ceil(len(self.__file_list) / self.__batch_size))
+
+    def __load_and_preprocess(self, fh, tag):
+        img = fh[tag][:]
+
+        if (self.__img_size is not None) and (img[..., 0].shape != self.__img_size):
+            im = Image.fromarray(img[..., 0])  # Can't handle the 1 channel
+            img = np.array(im.resize(self.__img_size[:-1], Image.LANCZOS)).astype(np.float32)
+            img = img[..., np.newaxis]
+
+        if img.max() > 1. or img.min() < 0.:
+            try:
+                img = min_max_norm(img).astype(np.float32)
+            except ValueError:
+                print(fh, tag, img.shape)
+                er_str = 'ERROR:\t[file]:\t{}\t[tag]:\t{}\t[img.shape]:\t{}\t'.format(fh, tag, img.shape)
+                raise ValueError(er_str)
+        return img.astype(np.float32)
+
+    def __getitem__(self, idx):
+        idxs = self.__idx_list[idx * self.__batch_size:(idx + 1) * self.__batch_size]
+
+        fix_imgs, mov_imgs = self.__load_samples(idxs)
+
+        zero_grad = np.zeros((*fix_imgs.shape[:-1], 2))
+
+        inputs = [mov_imgs, fix_imgs]
+        outputs = [fix_imgs, zero_grad]
+
+        if self.__multi_loss:
+            return [mov_imgs, fix_imgs, zero_grad],
+        else:
+            return (inputs, outputs)
+
+    def __load_samples(self, idx_list):
+        if self.__multi_loss:
+            img_shape = (0, *self.input_shape[0])
+        else:
+            img_shape = (0, *self.input_shape)
+
+        fix_imgs = np.empty(img_shape)
+        mov_imgs = np.empty(img_shape)
+        for i in idx_list:
+            f = h5py.File(self.__file_list[i], 'r')
+            fix_imgs = np.append(fix_imgs, [self.__load_and_preprocess(f, self.__tags['fix_img'])], axis=0)
+            mov_imgs = np.append(mov_imgs, [self.__load_and_preprocess(f, self.__tags['mov_img'])], axis=0)
+            f.close()
+
+        return fix_imgs, mov_imgs
+
+    def on_epoch_end(self):
+        np.random.shuffle(self.__idx_list)
+
+    def get_single_sample(self):
+        idx = random.randint(0, len(self.__idx_list))
+        fix, mov = self.__load_samples([idx])
+        return mov, fix
+
+

DeepDeformationMapRegistration/layers.py

@@ -0,0 +1,189 @@
+import os, sys
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+
+import tensorflow.keras.layers as kl
+import tensorflow.keras.backend as K
+import tensorflow as tf
+import numpy as np
+
+from DeepDeformationMapRegistration.utils.operators import soft_threshold
+
+
+class UncertaintyWeighting(kl.Layer):
+    def __init__(self, num_loss_fns=1, num_reg_fns=0, loss_fns: list = [tf.keras.losses.mean_squared_error],
+                 reg_fns: list = list(), prior_loss_w=[1.], manual_loss_w=[1.], prior_reg_w=[1.], manual_reg_w=[1.],
+                 **kwargs):
+        assert isinstance(loss_fns, list) and (num_loss_fns == len(loss_fns) or len(loss_fns) == 1)
+        assert isinstance(reg_fns, list) and (num_reg_fns == len(reg_fns))
+        self.num_loss = num_loss_fns
+        if len(loss_fns) == 1 and self.num_loss > 1:
+            self.loss_fns = loss_fns * self.num_loss
+        else:
+            self.loss_fns = loss_fns
+
+        if len(prior_loss_w) == 1:
+            self.prior_loss_w = prior_loss_w * num_loss_fns
+        else:
+            self.prior_loss_w = prior_loss_w
+        self.prior_loss_w = np.log(self.prior_loss_w)
+
+        if len(manual_loss_w) == 1:
+            self.manual_loss_w = manual_loss_w * num_loss_fns
+        else:
+            self.manual_loss_w = manual_loss_w
+
+        self.num_reg = num_reg_fns
+        if self.num_reg != 0:
+            if len(reg_fns) == 1 and self.num_reg > 1:
+                self.reg_fns = reg_fns * self.num_reg
+            else:
+                self.reg_fns = reg_fns
+
+            self.is_placeholder = True
+            if self.num_reg != 0:
+                if len(prior_reg_w) == 1:
+                    self.prior_reg_w = prior_reg_w * num_reg_fns
+                else:
+                    self.prior_reg_w = prior_reg_w
+                self.prior_reg_w = np.log(self.prior_reg_w)
+
+                if len(manual_reg_w) == 1:
+                    self.manual_reg_w = manual_reg_w * num_reg_fns
+                else:
+                    self.manual_reg_w = manual_reg_w
+
+        else:
+            self.prior_reg_w = list()
+            self.manual_reg_w = list()
+
+        super(UncertaintyWeighting, self).__init__(**kwargs)
+
+    def build(self, input_shape=None):
+        self.log_loss_vars = self.add_weight(name='loss_log_vars', shape=(self.num_loss,),
+                                             initializer=tf.keras.initializers.Constant(self.prior_loss_w),
+                                             trainable=True)
+        self.loss_weights = tf.math.softmax(self.log_loss_vars, name='SM_loss_weights')
+
+        if self.num_reg != 0:
+            self.log_reg_vars = self.add_weight(name='loss_reg_vars', shape=(self.num_reg,),
+                                                initializer=tf.keras.initializers.Constant(self.prior_reg_w),
+                                                trainable=True)
+            if self.num_reg == 1:
+                self.reg_weights = tf.math.exp(self.log_reg_vars, name='EXP_reg_weights')
+            else:
+                self.reg_weights = tf.math.softmax(self.log_reg_vars, name='SM_reg_weights')
+
+        super(UncertaintyWeighting, self).build(input_shape)
+
+    def multi_loss(self, ys_true, ys_pred, regs_true, regs_pred):
+        loss_values = list()
+        loss_names_loss = list()
+        loss_names_reg = list()
+
+        for y_true, y_pred, loss_fn, man_w in zip(ys_true, ys_pred, self.loss_fns, self.manual_loss_w):
+            loss_values.append(tf.keras.backend.mean(man_w * loss_fn(y_true, y_pred)))
+            loss_names_loss.append(loss_fn.__name__)
+
+        loss_values = tf.convert_to_tensor(loss_values, dtype=tf.float32, name="step_loss_values")
+        loss = tf.math.multiply(self.loss_weights, loss_values, name='step_weighted_loss')
+
+        if self.num_reg != 0:
+            loss_reg = list()
+            for reg_true, reg_pred, reg_fn, man_w in zip(regs_true, regs_pred, self.reg_fns, self.manual_reg_w):
+                loss_reg.append(K.mean(man_w * reg_fn(reg_true, reg_pred)))
+                loss_names_reg.append(reg_fn.__name__)
+
+            reg_values = tf.convert_to_tensor(loss_reg, dtype=tf.float32, name="step_reg_values")
+            loss = loss + tf.math.multiply(self.reg_weights, reg_values, name='step_weighted_reg')
+
+        for i, loss_name in enumerate(loss_names_loss):
+            self.add_metric(tf.slice(self.loss_weights, [i], [1]), name='LOSS_WEIGHT_{}_{}'.format(i, loss_name),
+                            aggregation='mean')
+            self.add_metric(tf.slice(loss_values, [i], [1]), name='LOSS_VALUE_{}_{}'.format(i, loss_name),
+                            aggregation='mean')
+        if self.num_reg != 0:
+            for i, loss_name in enumerate(loss_names_reg):
+                self.add_metric(tf.slice(self.reg_weights, [i], [1]), name='REG_WEIGHT_{}_{}'.format(i, loss_name),
+                                aggregation='mean')
+                self.add_metric(tf.slice(reg_values, [i], [1]), name='REG_VALUE_{}_{}'.format(i, loss_name),
+                                aggregation='mean')
+
+        return K.sum(loss)
+
+    def call(self, inputs):
+        ys_true = inputs[:self.num_loss]
+        ys_pred = inputs[self.num_loss:self.num_loss*2]
+        reg_true = inputs[-self.num_reg*2:-self.num_reg]
+        reg_pred = inputs[-self.num_reg:]             # The last terms are the regularization ones which have no GT
+        loss = self.multi_loss(ys_true, ys_pred, reg_true, reg_pred)
+        self.add_loss(loss, inputs=inputs)
+        # We won't actually use the output, but we need something for the TF graph
+        return K.concatenate(inputs, -1)
+
+    def get_config(self):
+        base_config = super(UncertaintyWeighting, self).get_config()
+        base_config['num_loss_fns'] = self.num_loss
+        base_config['num_reg_fns'] = self.num_reg
+
+        return base_config
+
+
+def distance_map(coord1, coord2, dist, img_shape_w_channel=(64, 64, 1)):
+    max_dist = np.max(img_shape_w_channel)
+    dm_p = np.ones(img_shape_w_channel, np.float32)*max_dist
+    dm_n = np.ones(img_shape_w_channel, np.float32)*max_dist
+
+    for c1, c2, d in zip(coord1, coord2, dist):
+        dm_p[c1, c2, 0] = d if dm_p[c1, c2, 0] > d else dm_p[c1, c2]
+        d_n = 64. - max_dist
+        dm_n[c1, c2, 0] = d_n if dm_n[c1, c2, 0] > d_n else dm_n[c1, c2]
+
+    return dm_p/max_dist, dm_n/max_dist
+
+
+def volume_to_ov_and_dm(in_volume: tf.Tensor):
+    # This one is run as a preprocessing step
+    def get_ov_projections_and_dm(volume):
+        # tf.sign returns -1, 0, 1 depending on the sign of the elements of the input (negative, zero, positive)
+        i, j, k, c = tf.where(volume > 0.0)
+        top = tf.sign(tf.reduce_sum(volume, axis=0), name='ov_top')
+        right = tf.sign(tf.reduce_sum(volume, axis=1), name='ov_right')
+        front = tf.sign(tf.reduce_sum(volume, axis=2), name='ov_front')
+
+        top_p, top_n = tf.py_func(distance_map, [j, k, i], tf.float32)
+        right_p, right_n = tf.py_func(distance_map, [i, k, j], tf.float32)
+        front_p, front_n = tf.py_func(distance_map, [i, j, k], tf.float32)
+
+        return [front, right, top], [front_p, front_n, top_p, top_n, right_p, right_n]
+
+    if len(in_volume.shape.as_list()) > 4:
+        return tf.map_fn(get_ov_projections_and_dm, in_volume, [tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32])
+    else:
+        return get_ov_projections_and_dm(in_volume)
+
+
+def ov_and_dm_to_volume(ov_projections):
+    front, right, top = ov_projections
+
+    def get_volume(front: tf.Tensor, right: tf.Tensor, top: tf.Tensor):
+        front_shape = front.shape.as_list()  # Assume (H, W, C)
+        top_shape = top.shape.as_list()
+
+        front_vol = tf.tile(tf.expand_dims(front, 2), [1, 1, top_shape[0], 1])
+        right_vol = tf.tile(tf.expand_dims(right, 1), [1, front_shape[1], 1, 1])
+        top_vol = tf.tile(tf.expand_dims(top, 0), [front_shape[0], 1, 1, 1])
+        sum = tf.add(tf.add(front_vol, right_vol), top_vol)
+        return soft_threshold(sum, 2., 'get_volume')
+
+    if len(front.shape.as_list()) > 3:
+        return tf.map_fn(lambda x: get_volume(x[0], x[1], x[2]), ov_projections, tf.float32)
+    else:
+        return get_volume(front, right, top)
+
+# TODO: Recovering the coordinates from the distance maps to prevent artifacts
+#   will the gradients be backpropagated??!?!!?!?!

DeepDeformationMapRegistration/losses.py

@@ -0,0 +1,47 @@
+import os, sys
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+
+import tensorflow as tf
+from scipy.ndimage import generate_binary_structure
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.operators import soft_threshold
+
+
+class HausdorffDistance:
+    def __init__(self, ndim=3, nerosion=10):
+        self.ndims = ndim
+        self.conv = getattr(tf.nn, 'conv%dd' % self.ndims)
+        self.nerosions = nerosion
+
+    def _erode(self, in_tensor, kernel):
+        out = 1. - tf.squeeze(self.conv(tf.expand_dims(1. - in_tensor, 0), kernel, [1] * (self.ndims + 2), 'SAME'), axis=0)
+        return soft_threshold(out, 0.5, name='soft_thresholding')
+
+    def _erosion_distance_single(self, y_true, y_pred):
+        diff = tf.math.pow(y_pred - y_true, 2)
+        alpha = 2.
+
+        norm = 1 / self.ndims * 2 + 1
+        kernel = generate_binary_structure(self.ndims, 1).astype(int) * norm
+        kernel = tf.constant(kernel, tf.float32)
+        kernel = tf.expand_dims(tf.expand_dims(kernel, -1), -1)
+
+        ret = 0.
+        for i in range(self.nerosions):
+            for j in range(i + 1):
+                er = self._erode(diff, kernel)
+            ret += tf.reduce_sum(tf.multiply(er, tf.pow(i + 1., alpha)))
+
+        return tf.multiply(C.IMG_SIZE ** -self.ndims, ret)  # Divide by the image size
+
+    def loss(self, y_true, y_pred):
+        batched_dist = tf.map_fn(lambda x: self._erosion_distance_single(x[0], x[1]), (y_true, y_pred),
+                                 dtype=tf.float32)
+
+        return batched_dist  # tf.reduce_mean(batched_dist)
+

DeepDeformationMapRegistration/networks.py

@@ -0,0 +1,63 @@
+import os, sys
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+
+import tensorflow as tf
+import voxelmorph as vxm
+from voxelmorph.tf.modelio import LoadableModel, store_config_args
+
+
+class VxmWeaklySupervised(LoadableModel):
+
+    @store_config_args
+    def __init__(self, inshape, all_labels: [list, tuple], nb_unet_features=None, int_steps=5, bidir=False, **kwargs):
+        """
+        Parameters:
+            inshape: Input shape. e.g. (192, 192, 192)
+            all_labels: List of all labels included in training segmentations.
+            hot_labels: List of labels to output as one-hot maps.
+            nb_unet_features: Unet convolutional features. See VxmDense documentation for more information.
+            int_steps: Number of flow integration steps. The warp is non-diffeomorphic when this value is 0.
+            kwargs: Forwarded to the internal VxmDense model.
+        """
+
+        fix_segm = tf.keras.Input((*inshape, len(all_labels)), name='fix_segmentations_input')
+        mov_segm = tf.keras.Input((*inshape, len(all_labels)), name='mov_segmentations_input')
+
+        mov_img = tf.keras.Input((*inshape, 1), name='mov_image_input')
+
+        unet_input_model = tf.keras.Model(inputs=[mov_segm, fix_segm], outputs=[mov_segm, fix_segm])
+
+        vxm_model = vxm.networks.VxmDense(inshape=inshape,
+                                          nb_unet_features=nb_unet_features,
+                                          input_model=unet_input_model,
+                                          int_steps=int_steps,
+                                          bidir=bidir, **kwargs)
+
+        pred_img = vxm.layers.SpatialTransformer(interp_method='linear', indexing='ij', name='pred_fix_img')(
+            [mov_img, vxm_model.references.pos_flow])
+
+        inputs = [mov_segm, fix_segm, mov_img] # mov_img, mov_segm, fix_segm
+        outputs = [pred_img] + vxm_model.outputs
+
+        self.references = LoadableModel.ReferenceContainer()
+        self.references.pred_segm = vxm_model.outputs[0]
+        self.references.pred_img = pred_img
+        self.references.pos_flow = vxm_model.references.pos_flow
+
+        super().__init__(inputs=inputs, outputs=outputs)
+
+    def get_registration_model(self):
+        return tf.keras.Model(self.inputs, self.references.pos_flow)
+
+    def register(self, mov_img, mov_segm, fix_segm):
+        return self.get_registration_model().predict([mov_segm, fix_segm, mov_img])
+
+    def apply_transform(self, mov_img, mov_segm, fix_segm, interp_method='linear'):
+        warp_model = self.get_registration_model()
+        img_input = tf.keras.Input(shape=mov_img.shape[1:], name='input_img')
+        pred_img = vxm.layers.SpatialTransformer(interp_method=interp_method)([img_input, warp_model.output])
+        return tf.keras.Model(warp_model.inputs, pred_img).predict([mov_segm, fix_segm, mov_img])

DeepDeformationMapRegistration/utils/acummulated_optimizer.py

@@ -0,0 +1,57 @@
+from tensorflow.keras.optimizers import Optimizer
+from tensorflow.keras import backend as K
+
+
+class AccumOptimizer(Optimizer):
+    """Optimizer
+    Inheriting Optimizer class, wrapping the original optimizer
+    to achieve a new corresponding optimizer of gradient accumulation.
+    # Arguments
+        optimizer: an instance of keras optimizer (supporting
+                    all keras optimizers currently available);
+        steps_per_update: the steps of gradient accumulation
+    # Returns
+        a new keras optimizer.
+    """
+    def __init__(self, optimizer, steps_per_update=1, **kwargs):
+        super(AccumOptimizer, self).__init__(**kwargs)
+        self.optimizer = optimizer
+        with K.name_scope(self.__class__.__name__):
+            self.steps_per_update = steps_per_update
+            self.iterations = K.variable(0, dtype='int64', name='iterations')
+            self.cond = K.equal(self.iterations % self.steps_per_update, 0)
+            self.lr = self.optimizer.lr
+            self.optimizer.lr = K.switch(self.cond, self.optimizer.lr, 0.)
+            for attr in ['momentum', 'rho', 'beta_1', 'beta_2']:
+                if hasattr(self.optimizer, attr):
+                    value = getattr(self.optimizer, attr)
+                    setattr(self, attr, value)
+                    setattr(self.optimizer, attr, K.switch(self.cond, value, 1 - 1e-7))
+            for attr in self.optimizer.get_config():
+                if not hasattr(self, attr):
+                    value = getattr(self.optimizer, attr)
+                    setattr(self, attr, value)
+            # Cover the original get_gradients method with accumulative gradients.
+            def get_gradients(loss, params):
+                return [ag / self.steps_per_update for ag in self.accum_grads]
+            self.optimizer.get_gradients = get_gradients
+    def get_updates(self, loss, params):
+        self.updates = [
+            K.update_add(self.iterations, 1),
+            K.update_add(self.optimizer.iterations, K.cast(self.cond, 'int64')),
+        ]
+        # gradient accumulation
+        self.accum_grads = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
+        grads = self.get_gradients(loss, params)
+        for g, ag in zip(grads, self.accum_grads):
+            self.updates.append(K.update(ag, K.switch(self.cond, g, ag + g)))
+        # inheriting updates of original optimizer
+        self.updates.extend(self.optimizer.get_updates(loss, params)[1:])
+        self.weights.extend(self.optimizer.weights)
+        return self.updates
+    def get_config(self):
+        iterations = K.eval(self.iterations)
+        K.set_value(self.iterations, 0)
+        config = self.optimizer.get_config()
+        K.set_value(self.iterations, iterations)
+        return config

DeepDeformationMapRegistration/utils/cmd_args_parser.py

@@ -0,0 +1,111 @@
+import sys, getopt
+import DeepDeformationMapRegistration.utils.constants as C
+import os
+
+
+def parse_arguments(argv):
+
+    try:
+        opts, args = getopt.getopt(argv, "hg:b:l:r:d:t:i:f:x:p:q:", ["gpu-num=",
+                                                             "batch-size=",
+                                                             "loss=",
+                                                             "remote=",
+                                                             "debug=",
+                                                             "debug-training=",
+                                                             "debug-input-data=",
+                                                             "destination-folder=",
+                                                             "destination-folder-fix=",
+                                                             "training-dataset=",
+                                                             "test-dataset=",
+                                                             "help"])
+    except getopt.GetoptError:
+        print('\n\t\t--gpu-num:\t\tGPU number to use'
+              '\n\t\t--batch-size:\t\tsize of the training batch'
+              '\n\t\t--loss:\t\tLoss function: ncc, mse, dssim'
+              '\n\t\t--remote:\t\tExecuting the script in The Beast: "True"/"False". Def: False'
+              '\n\t\t--debug:\t\tEnable debugging logs: "True"/"False". Def: False'
+              '\n\t\t--debug-training:\t\tEnable debugging training logs: "True"/"False". Def: False'
+              '\n\t\t--debug-input-data:\t\tEnable debugging input data logs: "True"/"False". Def: False'
+              '\n\t\t--destination-folder:\t\tName of the folder where to save the generated training files'
+              '\n\t\t--destination-folder-fixed:\t\tSame as --destination-folder but do not add the timestamp'
+              '\n\t\t--training-dataset:\t\tPath to the training dataset file'
+              '\n\t\t--test-dataset:\t\tPath to the test dataset file'
+              '\n')
+        sys.exit(2)
+
+    for opt, arg in opts:
+        if opt in ('--help', '-h'):
+            print('\n\t\t--gpu-num:\t\tGPU number to use\n\t\t--batch-size:\t\tsize of the training batch'
+                  '\n\t\t--loss:\t\tLoss function: ncc, mse, dssim\n')
+            continue
+        elif opt in ('--gpu_num', '-g'):
+            old = C.GPU_NUM
+            C.GPU_NUM = arg
+            os.environ['CUDA_VISIBLE_DEVICES'] = C.GPU_NUM
+            print('\t\tGPU_NUM: {} -> {}'.format(old, C.GPU_NUM))
+
+        elif opt in ('--batch-size', '-b'):
+            old = C.BATCH_SIZE
+            C.BATCH_SIZE = int(arg)
+            print('\t\tBATCH_SIZE: {} -> {}'.format(old, C.BATCH_SIZE))
+
+        elif opt in ('--destination-folder', '-f'):
+            old = C.DESTINATION_FOLDER
+            C.DESTINATION_FOLDER = arg + '_' + C.CUR_DATETIME
+            print('\t\tDESTINATION_FOLDER: {} -> {}'.format(old, C.DESTINATION_FOLDER))
+
+        elif opt in ('--destination-folder-fixed', '-x'):
+            old = C.DESTINATION_FOLDER
+            C.DESTINATION_FOLDER = arg
+            print('\t\tDESTINATION_FOLDER: {} -> {}'.format(old, C.DESTINATION_FOLDER))
+
+        elif opt in ('--training-dataset', '-p'):
+            old = C.TRAINING_DATASET
+            C.TRAINING_DATASET = arg
+            print('\t\tTRAINING_DATASET: {} -> {}'.format(old, C.TRAINING_DATASET))
+
+        elif opt in ('--test-dataset', '-q'):
+            old = C.TEST_DATASET
+            C.TEST_DATASET = arg
+            print('\t\tTEST_DATASET: {} -> {}'.format(old, C.TEST_DATASET))
+
+        elif opt in ('--remote', '-r'):
+            old = C.REMOTE
+            if arg.lower() in ('1', 'true', 't'):
+                C.REMOTE = True
+            else:
+                C.REMOTE = False
+            print('\t\tREMOTE: {} -> {}'.format(old, C.REMOTE))
+
+        elif opt in ('--debug', '-d'):
+            old = C.DEBUG
+            if arg.lower() in ('1', 'true', 't'):
+                C.DEBUG = True
+            else:
+                C.DEBUG = False
+            print('\t\tDEBUG: {} -> {}'.format(old, C.DEBUG))
+
+        elif opt in ('--debug-training', '-t'):
+            old = C.DEBUG_TRAINING
+            if arg.lower() in ('1', 'true', 't'):
+                C.DEBUG_TRAINING = True
+            else:
+                C.DEBUG_TRAINING = False
+            print('\t\tDEBUG_TRAINING: {} -> {}'.format(old, C.DEBUG_TRAINING))
+
+        elif opt in ('--debug-input-data', '-i'):
+            old = C.DEBUG_INPUT_DATA
+            if arg.lower() in ('1', 'true', 't'):
+                C.DEBUG_INPUT_DATA = True
+            else:
+                C.DEBUG_INPUT_DATA = False
+            print('\t\tDEBUG_INPUT_DATA: {} -> {}'.format(old, C.DEBUG_INPUT_DATA))
+
+        elif opt in ('--loss', '-l'):
+            old = C.LOSS_FNC
+            if arg in ('ncc', 'mse', 'dssim', 'dice'):
+                C.LOSS_FNC = arg
+            else:
+                print('Invalid option for --loss. Expected: "mse", "ncc" or "dssim", got {}'.format(arg))
+                sys.exit(2)
+            print('\t\tLOSS_FNC: {} -> {}'.format(old, C.LOSS_FNC))

DeepDeformationMapRegistration/utils/conf_file_utils.py

@@ -0,0 +1,52 @@
+import DeepDeformationMapRegistration.utils.constants as C
+import re
+import os
+
+
+class ConfigurationFile:
+    def __init__(self,
+                 file_path: str):
+        self.__file = file_path
+        self.__load_configuration()
+
+    def __load_configuration(self):
+        fd = open(self.__file, 'r')
+        file_lines = fd.readlines()
+
+        for line in file_lines:
+            if '#' not in line and line != '\n':
+                match = re.match('(.*)=(.*)', line)
+                if match[1] in C.__dict__.keys():
+                    # Careful with eval!!
+                    try:
+                        new_val = eval(match[2])
+                    except NameError:
+                        new_val = match[2]
+                    old = C.__dict__[match[1]]
+                    C.__dict__[match[1]] = new_val
+
+                    # Special case
+                    if match[1] == 'GPU_NUM':
+                        C.__dict__[match[1]] = str(new_val)
+                        os.environ['CUDA_VISIBLE_DEVICES'] = C.GPU_NUM
+
+                    if match[1] == 'EPOCHS':
+                        C.__dict__[match[1]] = new_val
+                        C.__dict__['SAVE_EPOCH'] = new_val // 10
+                        C.__dict__['VERBOSE_EPOCH'] = new_val // 10
+
+                    if match[1] == 'SAVE_EPOCH' or match[1] == 'VERBOSE_EPOCH':
+                        if new_val is not None:
+                            C.__dict__[match[1]] = C.__dict__['EPOCHS'] // new_val
+                        else:
+                            C.__dict__[match[1]] = None
+
+                    if match[1] == 'VALIDATION_ERR_LIMIT_COUNTER':
+                        C.__dict__[match[1]] = new_val
+                        C.__dict__['VALIDATION_ERR_LIMIT_COUNTER_BACKUP'] = new_val
+
+
+                    print('INFO: Updating constant {}: {} -> {}'.format(match[1], old, C.__dict__[match[1]]))
+                else:
+                    print('ERROR: Unknown constant {}'.format(match[1]))
+

DeepDeformationMapRegistration/utils/constants.py

@@ -0,0 +1,487 @@
+"""
+Constants
+"""
+import tensorflow as tf
+import os
+import datetime
+import numpy as np
+
+# RUN CONFIG
+REMOTE = False  # os.popen('hostname').read().encode('utf-8') == 'medtech-beast' #os.environ.get('REMOTE') == 'True'
+
+# Remote execution
+DEV_ORDER = 'PCI_BUS_ID'
+GPU_NUM = '0'
+
+# Dataset generation constants
+# See batchGenerator __next__ method: return [in_mov, in_fix], [disp_map, out_img]
+MOVING_IMG = 0
+FIXED_IMG = 1
+MOVING_PARENCHYMA_MASK = 2
+FIXED_PARENCHYMA_MASK = 3
+MOVING_VESSELS_MASK = 4
+FIXED_VESSELS_MASK = 5
+MOVING_TUMORS_MASK = 6
+FIXED_TUMORS_MASK = 7
+MOVING_SEGMENTATIONS = 8  # Compination of vessels and tumors
+FIXED_SEGMENTATIONS = 9  # Compination of vessels and tumors
+DISP_MAP_GT = 0
+PRED_IMG_GT = 1
+DISP_VECT_GT = 2
+DISP_VECT_LOC_GT = 3
+
+IMG_SIZE = 64  # Assumed a square image
+IMG_SHAPE = (IMG_SIZE, IMG_SIZE, IMG_SIZE, 1)  # (IMG_SIZE, IMG_SIZE, 1)
+DISP_MAP_SHAPE = (IMG_SIZE, IMG_SIZE, IMG_SIZE, 3)
+BATCH_SHAPE = (None, IMG_SIZE, IMG_SIZE, IMG_SIZE, 2)  # Expected batch shape by the network
+BATCH_SHAPE_SEGM = (None, IMG_SIZE, IMG_SIZE, IMG_SIZE, 3)  # Expected batch shape by the network
+IMG_BATCH_SHAPE = (None, IMG_SIZE, IMG_SIZE, IMG_SIZE, 1)  # Batch shape for single images
+
+RAW_DATA_BASE_DIR = './data'
+DEFORMED_DATA_NAME = 'deformed'
+GROUND_TRUTH_DATA_NAME = 'groundTruth'
+GROUND_TRUTH_COORDS_FILE = 'centerlineCoords_GT.txt'
+DEFORMED_COORDS_FILE = 'centerlineCoords_DF.txt'
+H5_MOV_IMG = 'input/{}'.format(MOVING_IMG)
+H5_FIX_IMG = 'input/{}'.format(FIXED_IMG)
+H5_MOV_PARENCHYMA_MASK = 'input/{}'.format(MOVING_PARENCHYMA_MASK)
+H5_FIX_PARENCHYMA_MASK = 'input/{}'.format(FIXED_PARENCHYMA_MASK)
+H5_MOV_VESSELS_MASK = 'input/{}'.format(MOVING_VESSELS_MASK)
+H5_FIX_VESSELS_MASK = 'input/{}'.format(FIXED_VESSELS_MASK)
+H5_MOV_TUMORS_MASK = 'input/{}'.format(MOVING_TUMORS_MASK)
+H5_FIX_TUMORS_MASK = 'input/{}'.format(FIXED_TUMORS_MASK)
+H5_FIX_SEGMENTATIONS = 'input/{}'.format(FIXED_SEGMENTATIONS)
+H5_MOV_SEGMENTATIONS = 'input/{}'.format(MOVING_SEGMENTATIONS)
+
+H5_GT_DISP = 'output/{}'.format(DISP_MAP_GT)
+H5_GT_IMG = 'output/{}'.format(PRED_IMG_GT)
+H5_GT_DISP_VECT = 'output/{}'.format(DISP_VECT_GT)
+H5_GT_DISP_VECT_LOC = 'output/{}'.format(DISP_VECT_LOC_GT)
+H5_PARAMS_INTENSITY_RANGE = 'parameters/intensity'
+TRAINING_PERC = 0.8
+VALIDATION_PERC = 1 - TRAINING_PERC
+MAX_ANGLE = 45.0  # degrees
+MAX_FLIPS = 2  # Axes to flip over
+NUM_ROTATIONS = 5
+MAX_WORKERS = 10
+
+# Training constants
+MODEL = 'unet'
+BATCH_NORM = False
+TENSORBOARD = False
+LIMIT_NUM_SAMPLES = None  # If you don't want to use all the samples in the training set. None to use all
+TRAINING_DATASET = 'data/training.hd5'
+TEST_DATASET = 'data/test.hd5'
+VALIDATION_DATASET = 'data/validation.hd5'
+LOSS_FNC = 'mse'
+LOSS_SCHEME = 'unidirectional'
+NUM_EPOCHS = 10
+DATA_FORMAT = 'channels_last'  # or 'channels_fist'
+DATA_DIR = './data'
+MODEL_CHECKPOINT = './model_checkpoint'
+BATCH_SIZE = 8
+EPOCHS = 100
+SAVE_EPOCH = EPOCHS // 10  # Epoch when to save the model
+VERBOSE_EPOCH = EPOCHS // 10
+VALIDATION_ERR_LIMIT = 0.2  # Stop training if reached this limit
+VALIDATION_ERR_LIMIT_COUNTER = 10  # Number of successive times the validation error was smaller than the threshold
+VALIDATION_ERR_LIMIT_COUNTER_BACKUP = 10
+THRESHOLD = 0.5  # Threshold to select the centerline in the interpolated images
+RESTORE_TRAINING = True  # look for previously saved models to resume training
+EARLY_STOP_PATIENCE = 10
+LOG_FIELD_NAMES = ['time', 'epoch', 'step',
+                   'training_loss_mean', 'training_loss_std',
+                   'training_loss1_mean', 'training_loss1_std',
+                   'training_loss2_mean', 'training_loss2_std',
+                   'training_loss3_mean', 'training_loss3_std',
+                   'training_ncc1_mean', 'training_ncc1_std',
+                   'training_ncc2_mean', 'training_ncc2_std',
+                   'training_ncc3_mean', 'training_ncc3_std',
+                   'validation_loss_mean', 'validation_loss_std',
+                   'validation_loss1_mean', 'validation_loss1_std',
+                   'validation_loss2_mean', 'validation_loss2_std',
+                   'validation_loss3_mean', 'validation_loss3_std',
+                   'validation_ncc1_mean', 'validation_ncc1_std',
+                   'validation_ncc2_mean', 'validation_ncc2_std',
+                   'validation_ncc3_mean', 'validation_ncc3_std']
+LOG_FIELD_NAMES_SHORT = ['time', 'epoch', 'step',
+                         'training_loss_mean', 'training_loss_std',
+                         'training_loss1_mean', 'training_loss1_std',
+                         'training_loss2_mean', 'training_loss2_std',
+                         'training_ncc1_mean', 'training_ncc1_std',
+                         'training_ncc2_mean', 'training_ncc2_std',
+                         'validation_loss_mean', 'validation_loss_std',
+                         'validation_loss1_mean', 'validation_loss1_std',
+                         'validation_loss2_mean', 'validation_loss2_std',
+                         'validation_ncc1_mean', 'validation_ncc1_std',
+                         'validation_ncc2_mean', 'validation_ncc2_std']
+LOG_FIELD_NAMES_UNET = ['time', 'epoch', 'step', 'reg_smooth_coeff', 'reg_jacob_coeff',
+                        'training_loss_mean', 'training_loss_std',
+                        'training_loss_dissim_mean', 'training_loss_dissim_std',
+                        'training_reg_smooth_mean', 'training_reg_smooth_std',
+                        'training_reg_jacob_mean', 'training_reg_jacob_std',
+                        'training_ncc_mean', 'training_ncc_std',
+                        'training_dice_mean', 'training_dice_std',
+                        'training_owo_mean', 'training_owo_std',
+                        'validation_loss_mean', 'validation_loss_std',
+                        'validation_loss_dissim_mean', 'validation_loss_dissim_std',
+                        'validation_reg_smooth_mean', 'validation_reg_smooth_std',
+                        'validation_reg_jacob_mean', 'validation_reg_jacob_std',
+                        'validation_ncc_mean', 'validation_ncc_std',
+                        'validation_dice_mean', 'validation_dice_std',
+                        'validation_owo_mean', 'validation_owo_std']
+CUR_DATETIME = datetime.datetime.now().strftime("%H%M_%d%m%Y")
+DESTINATION_FOLDER = 'training_log_' + CUR_DATETIME
+CSV_DELIMITER = ";"
+CSV_QUOTE_CHAR = '"'
+REG_SMOOTH = 0.0
+REG_MAG = 1.0
+REG_TYPE = 'l2'
+MAX_DISP_DM = 10.
+MAX_DISP_DM_TF = tf.constant((MAX_DISP_DM,), tf.float32, name='MAX_DISP_DM')
+MAX_DISP_DM_PERC = 0.25
+
+W_SIM = 0.7
+W_REG = 0.3
+W_INV = 0.1
+
+# Loss function parameters
+REG_SMOOTH1 = 1 / 100000
+REG_SMOOTH2 = 1 / 5000
+REG_SMOOTH3 = 1 / 5000
+LOSS1 = 1.0
+LOSS2 = 0.6
+LOSS3 = 0.3
+REG_JACOBIAN = 0.1
+
+LOSS_COEFFICIENT = 1.0
+REG_COEFFICIENT = 1.0
+
+DICE_SMOOTH = 1.
+
+CC_WINDOW = [9,9,9]
+
+# Adam optimizer
+LEARNING_RATE = 1e-3
+B1 = 0.9
+B2 = 0.999
+LEARNING_RATE_DECAY = 0.01
+LEARNING_RATE_DECAY_STEP = 10000  # Update the learning rate every LEARNING_RATE_DECAY_STEP steps
+OPTIMIZER = 'adam'
+
+# Network architecture constants
+LAYER_MAXPOOL = 0
+LAYER_UPSAMP = 1
+LAYER_CONV = 2
+AFFINE_TRANSF = False
+OUTPUT_LAYER = 3
+DROPOUT = True
+DROPOUT_RATE = 0.2
+MAX_DATA_SIZE = (1000, 1000, 1)
+PLATEAU_THR = 0.01  # A slope between +-PLATEAU_THR will be considered a plateau for the LR updating function
+ENCODER_FILTERS = [4, 8, 16, 32, 64]
+
+# SSIM
+SSIM_FILTER_SIZE = 11  # Size of Gaussian filter
+SSIM_FILTER_SIGMA = 1.5  # Width of Gaussian filter
+SSIM_K1 = 0.01  # Def. 0.01
+SSIM_K2 = 0.03  # Recommended values 0 < K2 < 0.4
+MAX_VALUE = 1.0  # Maximum intensity values
+
+# Mathematic constants
+EPS = 1e-8
+EPS_tf = tf.constant(EPS, dtype=tf.float32)
+LOG2 = tf.math.log(tf.constant(2, dtype=tf.float32))
+
+# Debug constants
+VERBOSE = False
+DEBUG = False
+DEBUG_TRAINING = False
+DEBUG_INPUT_DATA = False
+
+# Plotting
+FONT_SIZE = 10
+DPI = 200  # Dots Per Inch
+
+# Coordinates
+B = 0  # Batch dimension
+H = 1  # Height dimension
+W = 2  # Width dimension
+D = 3  # Depth
+C = -1  # Channel dimension
+
+D_DISP = 2
+W_DISP = 1
+H_DISP = 0
+
+DIMENSIONALITY = 3
+
+# Interpolation type
+BIL_INTERP = 0
+TPS_INTERP = 1
+CUADRATIC_C = 0.5
+
+# Data augmentation
+MAX_DISP = 5  # Test = 15
+NUM_ROT = 5
+NUM_FLIPS = 2
+MAX_ANGLE = 10
+
+# Thin Plate Splines implementation constants
+TPS_NUM_CTRL_PTS_PER_AXIS = 4
+TPS_NUM_CTRL_PTS = np.power(TPS_NUM_CTRL_PTS_PER_AXIS, DIMENSIONALITY)
+TPS_REG = 0.01
+DISP_SCALE = 2  # Scaling of the output of the CNN to increase the range of tanh
+
+
+class CoordinatesGrid:
+    def __init__(self):
+        self.__grid = 0
+        self.__grid_fl = 0
+        self.__norm = False
+        self.__num_pts = 0
+        self.__batches = False
+        self.__shape = None
+        self.__shape_flat = None
+
+    def set_coords_grid(self, img_shape: tf.TensorShape, num_ppa: int = None, batches: bool = False,
+                        img_type: tf.DType = tf.float32, norm: bool = False):
+        self.__batches = batches
+        not_batches = not batches  # Just to not make a too complex code when indexing the values
+        if num_ppa is None:
+            num_ppa = img_shape
+        if norm:
+            x_coords = tf.linspace(-1., 1.,
+                                   num_ppa[W - int(not_batches)])  # np.linspace works fine, tf had some issues...
+            y_coords = tf.linspace(-1., 1., num_ppa[H - int(not_batches)])  # num_ppa: number of points per axis
+            z_coords = tf.linspace(-1., 1., num_ppa[D - int(not_batches)])
+        else:
+            x_coords = tf.linspace(0., img_shape[W - int(not_batches)] - 1.,
+                                   num_ppa[W - int(not_batches)])  # np.linspace works fine, tf had some issues...
+            y_coords = tf.linspace(0., img_shape[H - int(not_batches)] - 1.,
+                                   num_ppa[H - int(not_batches)])  # num_ppa: number of points per axis
+            z_coords = tf.linspace(0., img_shape[D - int(not_batches)] - 1., num_ppa[D - int(not_batches)])
+
+        coords = tf.meshgrid(x_coords, y_coords, z_coords, indexing='ij')
+        self.__num_pts = num_ppa[W - int(not_batches)] * num_ppa[H - int(not_batches)] * num_ppa[D - int(not_batches)]
+
+        grid = tf.stack([coords[0], coords[1], coords[2]], axis=-1)
+        grid = tf.cast(grid, img_type)
+
+        grid_fl = tf.stack([tf.reshape(coords[0], [-1]),
+                            tf.reshape(coords[1], [-1]),
+                            tf.reshape(coords[2], [-1])], axis=-1)
+        grid_fl = tf.cast(grid_fl, img_type)
+
+        grid_homogeneous = tf.stack([tf.reshape(coords[0], [-1]),
+                                     tf.reshape(coords[1], [-1]),
+                                     tf.reshape(coords[2], [-1]),
+                                     tf.ones_like(tf.reshape(coords[0], [-1]))], axis=-1)
+
+        self.__shape = np.asarray([num_ppa[W - int(not_batches)], num_ppa[H - int(not_batches)], num_ppa[D - int(not_batches)], 3])
+        total_num_pts = np.prod(self.__shape[:-1])
+        self.__shape_flat = np.asarray([total_num_pts, 3])
+        if batches:
+            grid = tf.expand_dims(grid, axis=0)
+            grid = tf.tile(grid, [img_shape[B], 1, 1, 1, 1])
+            grid_fl = tf.expand_dims(grid_fl, axis=0)
+            grid_fl = tf.tile(grid_fl, [img_shape[B], 1, 1])
+            grid_homogeneous = tf.expand_dims(grid_homogeneous, axis=0)
+            grid_homogeneous = tf.tile(grid_homogeneous, [img_shape[B], 1, 1])
+            self.__shape = np.concatenate([np.asarray((img_shape[B],)), self.__shape])
+            self.__shape_flat = np.concatenate([np.asarray((img_shape[B],)), self.__shape_flat])
+
+        self.__norm = norm
+        self.__grid_fl = grid_fl
+        self.__grid = grid
+        self.__grid_homogeneous = grid_homogeneous
+
+    @property
+    def grid(self):
+        return self.__grid
+
+    @property
+    def size(self):
+        return self.__len__()
+
+    def grid_flat(self, transpose=False):
+        if transpose:
+            if self.__batches:
+                ret = tf.transpose(self.__grid_fl, (0, 2, 1))
+            else:
+                ret = tf.transpose(self.__grid_fl)
+        else:
+            ret = self.__grid_fl
+        return ret
+
+    def grid_homogeneous(self, transpose=False):
+        if transpose:
+            if self.__batches:
+                ret = tf.transpose(self.__grid_homogeneous, (0, 2, 1))
+            else:
+                ret = tf.transpose(self.__grid_homogeneous)
+        else:
+            ret = self.__grid_homogeneous
+        return ret
+
+    @property
+    def is_normalized(self):
+        return self.__norm
+
+    def __len__(self):
+        return tf.size(self.__grid)
+
+    @property
+    def number_pts(self):
+        return self.__num_pts
+
+    @property
+    def shape_grid_flat(self):
+        return self.__shape_flat
+
+    @property
+    def shape(self):
+        return self.__shape
+
+
+
+COORDS_GRID = CoordinatesGrid()
+
+
+class VisualizationParameters:
+    def __init__(self):
+        self.__scale = None  # See https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.axes.Axes.quiver.html
+        self.__spacing = 5
+
+    def set_spacing(self, img_shape: tf.TensorShape):
+        self.__spacing = int(5 * np.log(img_shape[W]))
+
+    @property
+    def spacing(self):
+        return self.__spacing
+
+    def set_arrow_scale(self, scale: int):
+        self.__scale = scale
+
+    @property
+    def arrow_scale(self):
+        return self.__scale
+
+
+QUIVER_PARAMS = VisualizationParameters()
+
+# Configuration file
+CONF_FILE_NAME = 'configuration.txt'
+
+
+def summary():
+    return '##### CONFIGURATION: REMOTE {}  DEBUG {} DEBUG TRAINING {}' \
+           '\n\t\tLEARNING RATE: {}' \
+           '\n\t\tBATCH SIZE: {}' \
+           '\n\t\tLIMIT NUM SAMPLES: {}' \
+           '\n\t\tLOSS_FNC: {}' \
+           '\n\t\tTRAINING_DATASET: {} ({:.1f}%/{:.1f}%)' \
+           '\n\t\tTEST_DATASET: {}'.format(REMOTE, DEBUG, DEBUG_TRAINING, LEARNING_RATE, BATCH_SIZE, LIMIT_NUM_SAMPLES,
+                                           LOSS_FNC, TRAINING_DATASET, TRAINING_PERC * 100, (1 - TRAINING_PERC) * 100,
+                                           TEST_DATASET)
+
+
+# LOG Severity levers
+# https://docs.python.org/2/library/logging.html#logging-levels
+INF = 20  # Information
+WAR = 30  # Warning
+ERR = 40  # Error
+DEB = 10  # Debug
+CRI = 50  # Critical
+
+SEVERITY_STR = {INF: 'INFO',
+                WAR: 'WARNING',
+                ERR: 'ERROR',
+                DEB: 'DEBUG',
+                CRI: 'CRITICAL'}
+
+HL_LOG_FIELD_NAMES = ['Time', 'Epoch', 'Step',
+                      'train_loss', 'train_loss_std',
+                      'train_loss1', 'train_loss1_std',
+                      'train_loss2', 'train_loss2_std',
+                      'train_loss3', 'train_loss3_std',
+                      'train_NCC', 'train_NCC_std',
+                      'val_loss', 'val_loss_std',
+                      'val_loss1', 'val_loss1_std',
+                      'val_loss2', 'val_loss2_std',
+                      'val_loss3', 'val_loss3_std',
+                      'val_NCC', 'val_NCC_std']
+
+# Sobel filters
+SOBEL_W_2D = tf.constant([[-1., 0., 1.],
+                          [-2., 0., 2.],
+                          [-1., 0., 1.]], dtype=tf.float32, name='sobel_w_2d')
+SOBEL_W_3D = tf.tile(tf.expand_dims(SOBEL_W_2D, axis=-1), [1, 1, 3])
+SOBEL_H_3D = tf.transpose(SOBEL_W_3D, [1, 0, 2])
+SOBEL_D_3D = tf.transpose(SOBEL_W_3D, [2, 1, 0])
+
+aux = tf.expand_dims(tf.expand_dims(SOBEL_W_3D, axis=-1), axis=-1)
+SOBEL_FILTER_W_3D_IMAGE = aux
+SOBEL_FILTER_W_3D = tf.tile(aux, [1, 1, 1, 3, 3])
+# tf.nn.conv3d expects the filter in [D, H, W, C_in, C_out] order
+SOBEL_FILTER_W_3D = tf.transpose(SOBEL_FILTER_W_3D, [2, 0, 1, 3, 4], name='sobel_filter_i_3d')
+
+aux = tf.expand_dims(tf.expand_dims(SOBEL_H_3D, axis=-1), axis=-1)
+SOBEL_FILTER_H_3D_IMAGE = aux
+SOBEL_FILTER_H_3D = tf.tile(aux, [1, 1, 1, 3, 3])
+SOBEL_FILTER_H_3D = tf.transpose(SOBEL_FILTER_H_3D, [2, 0, 1, 3, 4], name='sobel_filter_j_3d')
+
+aux = tf.expand_dims(tf.expand_dims(SOBEL_D_3D, axis=-1), axis=-1)
+SOBEL_FILTER_D_3D_IMAGE = aux
+SOBEL_FILTER_D_3D = tf.tile(aux, [1, 1, 1, 3, 3])
+SOBEL_FILTER_D_3D = tf.transpose(SOBEL_FILTER_D_3D, [2, 1, 0, 3, 4], name='sobel_filter_k_3d')
+
+# Filters for spatial integration of the displacement map
+INTEG_WIND_SIZE = IMG_SIZE
+INTEG_STEPS = 4  # VoxelMorph default value for the integration of the stationary velocity field. >4 memory alloc issue
+INTEG_FILTER_D = tf.ones([INTEG_WIND_SIZE, 1, 1, 1, 1], name='integrate_h_filter')
+INTEG_FILTER_H = tf.ones([1, INTEG_WIND_SIZE, 1, 1, 1], name='integrate_w_filter')
+INTEG_FILTER_W = tf.ones([1, 1, INTEG_WIND_SIZE, 1, 1], name='integrate_d_filter')
+
+# Laplacian filter
+LAPLACIAN_27_P = tf.constant(np.asarray([np.ones((3, 3)),
+                                         [[1, 1, 1],
+                                          [1, -26, 1],
+                                          [1, 1, 1]],
+                                         np.ones((3, 3))]), tf.float32)
+LAPLACIAN_27_P = tf.expand_dims(tf.expand_dims(LAPLACIAN_27_P, axis=-1), axis=-1)
+LAPLACIAN_27_P = tf.tile(LAPLACIAN_27_P, [1, 1, 1, 3, 3], name='laplacian_27_p')
+
+
+LAPLACIAN_7_P = tf.constant(np.asarray([[[0, 0, 0],
+                                         [0, 1, 0],
+                                         [0, 0, 0]],
+                                        [[0, 1, 0],
+                                         [1, -6, 1],
+                                         [0, 1, 0]],
+                                        [[0, 0, 0],
+                                         [0, 1, 0],
+                                         [0, 0, 0]]]), tf.float32)
+LAPLACIAN_7_P = tf.expand_dims(tf.expand_dims(LAPLACIAN_7_P, axis=-1), axis=-1)
+LAPLACIAN_7_P = tf.tile(LAPLACIAN_7_P, [1, 1, 1, 3, 3], name='laplacian_7_p')
+
+# Constants for bias loss
+ZERO_WARP = tf.zeros((1,) + DISP_MAP_SHAPE, name='zero_warp')
+BIAS_WARP_WEIGHT = 1e-02
+BIAS_AFFINE_WEIGHT = 1e-02
+
+# Overlapping score
+OS_SCALE = 10
+EPS_1 = 1.0
+EPS_1_tf = tf.constant(EPS_1)
+
+# LDDMM
+GAUSSIAN_KERNEL_SHAPE = (8, 8, 8)
+
+# Constants for MultiLoss layer
+PRIOR_W = [1., 1 / 60, 1.]
+MANUAL_W = [1.] * len(PRIOR_W)
+
+REG_PRIOR_W = [1e-3]
+REG_MANUAL_W = [1.] * len(REG_PRIOR_W)
+

DeepDeformationMapRegistration/utils/misc.py

@@ -0,0 +1,28 @@
+import os
+import errno
+import nibabel as nb
+import numpy as np
+import re
+
+def try_mkdir(dir):
+    try:
+        os.makedirs(dir)
+    except OSError as err:
+        if err.errno == errno.EEXIST:
+            print("Directory " + dir + " already exists")
+        else:
+            raise ValueError("Can't create dir " + dir)
+    else:
+        print("Created directory " + dir)
+
+
+def function_decorator(new_name):
+    """"
+    Change the __name__ property of a function using new_name.
+    :param new_name:
+    :return:
+    """
+    def decorator(func):
+        func.__name__ = new_name
+        return func
+    return decorator

DeepDeformationMapRegistration/utils/nifty_utils.py

@@ -0,0 +1,42 @@
+import os
+import errno
+import nibabel as nb
+import numpy as np
+import re
+import zipfile
+import tensorflow as tf
+
+
+TEMP_UNZIP_PATH = '/mnt/EncryptedData1/Users/javier/ext_datasets/LITS17/temp'
+NII_EXTENSION = '.nii'
+
+
+def save_nifti(data, save_path):
+    data_nifti = nb.Nifti1Image(data, affine=np.eye(4))
+
+    data_nifti.header.get_xyzt_units()
+    try:
+        data_nifti.to_filename(save_path)  # Save as NiBabel file
+        print('Saved {}'.format(save_path))
+    except ValueError:
+        print('Could not save {}'.format(save_path))
+
+
+def unzip_nii_file(file_path):
+    file_dir, file_name = os.path.split(file_path)
+    file_name = file_name.split('.zip')[0]
+
+    dest_path = os.path.join(TEMP_UNZIP_PATH, file_name)
+    zipfile.ZipFile(file_path).extractall(TEMP_UNZIP_PATH)
+
+    if not os.path.exists(dest_path):
+        print('ERR: File {} not unzip-ed!'.format(file_path))
+        dest_path = None
+    return dest_path
+
+
+def delete_temp(file_path, verbose=False):
+    assert NII_EXTENSION in file_path
+    os.remove(file_path)
+    if verbose:
+        print('Deleted file: ', file_path)

DeepDeformationMapRegistration/utils/operators

@@ -0,0 +1,28 @@
+import numpy as np
+import tensorflow as tf
+
+
+def min_max_norm(img: np.ndarray, out_max_val=1.):
+    out_img = img
+    max_val = np.amax(img)
+    min_val = np.amin(img)
+    if (max_val - min_val) != 0:
+        out_img = (img - min_val) / (max_val - min_val)
+    return out_img * out_max_val
+
+
+def soft_threshold(x, threshold, name=None):
+    # https://www.tensorflow.org/probability/api_docs/python/tfp/math/soft_threshold
+    with tf.name_scope(name or 'soft_threshold'):
+        x = tf.convert_to_tensor(x, name='x')
+        threshold = tf.convert_to_tensor(threshold, dtype=x.dtype, name='threshold')
+        return tf.sign(x) * tf.maximum(tf.abs(x) - threshold, 0.)
+
+
+def binary_activation(x):
+    # https://stackoverflow.com/questions/37743574/hard-limiting-threshold-activation-function-in-tensorflow
+    cond = tf.less(x, tf.zeros(tf.shape(x)))
+    out = tf.where(cond, tf.zeros(tf.shape(x)), tf.ones(tf.shape(x)))
+
+    return out
+

DeepDeformationMapRegistration/utils/user_interface.py

@@ -0,0 +1,23 @@
+import numpy as np
+import re
+import os
+
+
+def show_and_select(file_list, msg='Select a file by the number: ', int_if_single=True):
+    # If the selection is a single number, then return that number instead of the list of length 1
+    invalid_selection = True
+    while invalid_selection:
+        for i, f in enumerate(file_list):
+            print('{:03d}) {}'. format(i+1, os.path.split(f)[-1]))
+
+        sel = np.asarray(re.split(',\s|,|\s',input(msg)), np.int) - 1
+
+        if (np.all(sel >= 0)) and (np.all(sel <= len(file_list))):
+            invalid_selection = False
+    sel = [file_list[s] for s in sel]
+    print('Selected: ', ', '.join([os.path.split(f)[-1] for f in sel]))
+
+    if int_if_single:
+        if len(sel) == 1:
+            sel = sel[0]
+    return sel

DeepDeformationMapRegistration/utils/visualization.py

@@ -0,0 +1,1151 @@
+# import matplotlib
+# matplotlib.use('TkAgg')
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+import matplotlib.colors as mcolors
+from matplotlib.lines import Line2D
+from matplotlib import cm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+import tensorflow as tf
+import numpy as np
+import pyVesselRegistration_constants as const
+from skimage.exposure import rescale_intensity
+import scipy.misc as scpmisc
+import os
+
+THRES = 0.9
+
+# COLOR MAPS
+chunks = np.linspace(0, 1, 10)
+cmap1 = plt.get_cmap('hsv', 4)
+# cmaplist = [cmap1(i) for i in range(cmap1.N)]
+cmaplist = [(1, 1, 1, 1), (0, 0, 1, 1), (230 / 255, 97 / 255, 1 / 255, 1), (128 / 255, 0 / 255, 32 / 255, 1)]
+cmaplist[0] = (1, 1, 1, 1.0)
+cmap1 = mcolors.LinearSegmentedColormap.from_list('custom', cmaplist, cmap1.N)
+
+colors = [(0, 0, 1, i) for i in chunks]
+cmap2 = mcolors.LinearSegmentedColormap.from_list('mycmap', colors, N=100)
+
+colors = [(230 / 255, 97 / 255, 1 / 255, i) for i in chunks]
+cmap3 = mcolors.LinearSegmentedColormap.from_list('mycmap', colors, N=100)
+
+colors = [(128 / 255, 0 / 255, 32 / 255, i) for i in chunks]
+cmap4 = mcolors.LinearSegmentedColormap.from_list('mycmap', colors, N=100)
+
+cmap_bin = cm.get_cmap('viridis', 3)  # viridis is the default colormap
+
+
+def view_centerline_sample(sample: np.ndarray, dimensionality: int, ax=None, c=None, name=None):
+    if dimensionality == 2:
+        _plot_2d(sample, ax, c, name=name)
+    elif dimensionality == 3:
+        _plot_3d(sample, ax, c, name=name)
+    else:
+        raise ValueError('Invalid valud for dimensionality. Expected int 2 or 3')
+
+
+def matrix_to_orthographicProjection(matrix: np.ndarray, ret_list=False):
+    """ Given a 3D matrix, returns the three orthographic projections: top, front, right.
+    Top corresponds to dimensions 1 and 2
+    Front corresponds to dimensions 0 and 1
+    Right corresponds to dimensions 0 and 2
+
+    :param matrix: 3D matrix
+    :param ret_list: return a list instead of an array (optional)
+    :return: list or array with the three views [top, front, right]
+    """
+    top = _getProjection(matrix, dim=0)  # YZ
+    front = _getProjection(matrix, dim=2)  # XY
+    right = _getProjection(matrix, dim=1)  # XZ
+
+    if ret_list:
+        return top, front, right
+    else:
+        return np.asarray([top, front, right])
+
+
+def _getProjection(matrix: np.ndarray, dim: int):
+    orth_view = matrix.sum(axis=dim, dtype=float)
+    orth_view = orth_view > 0.0
+    orth_view.astype(np.float)
+
+    return orth_view
+
+
+def orthographicProjection_to_matrix(top: np.ndarray, front: np.ndarray, right: np.ndarray):
+    """ Given the three orthographic projections, it returns a 3D-view of the object based on back projection
+
+    :param top: 2D view top view
+    :param front: 2D front view
+    :param right: 2D right view
+    :return: matrix with the 3D-view
+    """
+    top_mat = np.tile(top, (front.shape[0], 1, 1))
+    front_mat = np.tile(top, (right.shape[1], 1, 1))
+    right_mat = np.tile(top, (top.shape[0], 1, 1))
+
+    reconstruction = np.zeros((front.shape[0], right.shape[1], top.shape[0]))
+    iter = np.nditer([top_mat, front_mat, right_mat, reconstruction], flags=['multi_index'], op_flags=['readwrite'])
+    while not iter.finished:
+        if iter[0] and iter[1] and iter[2]:
+            iter[3] = 1
+        iter.iternext()
+
+    return reconstruction
+
+
+def _plot_2d(sample: np.ndarray, ax=None, c=None, name=None):
+    if isinstance(sample, tf.Tensor):
+        sample = sample.eval(session=tf.Session())
+
+    x_range = list()
+    y_range = list()
+    marker_size = list()
+    for idx, val in np.ndenumerate(sample):
+        if val >= THRES:
+            x_range.append(idx[0])
+            y_range.append(idx[1])
+            marker_size.append(val ** 2)
+
+    if not ax:
+        fig = plt.figure()
+        ax = fig.add_subplot(111)
+
+    if c:
+        ax.scatter(x_range, y_range, c=c, s=marker_size)
+    else:
+        ax.scatter(x_range, y_range, s=marker_size)
+
+    ax.set_xlabel('X')
+    ax.set_ylabel('Y')
+    if name:
+        ax.set_title(name)
+
+    return ax
+
+
+def _plot_3d(sample: np.ndarray, ax=None, c=None, name=None):
+    from mpl_toolkits.mplot3d import Axes3D
+    if isinstance(sample, tf.Tensor):
+        sample = sample.eval(session=tf.Session())
+
+    x_range = list()
+    y_range = list()
+    z_range = list()
+    marker_size = list()
+    for idx, val in np.ndenumerate(sample):
+        if val >= THRES:
+            x_range.append(idx[0])
+            y_range.append(idx[1])
+            z_range.append(idx[2])
+            marker_size.append(val ** 2)
+
+    print('Found ', len(x_range), ' points')
+    # x_range = np.linspace(start=0, stop=sample.shape[0], num=sample.shape[0])
+    # y_range = np.linspace(start=0, stop=sample.shape[1], num=sample.shape[1])
+    # z_range = np.linspace(start=0, stop=sample.shape[2], num=sample.shape[2])
+    #
+    # sample_flat = sample.flatten(order='C')
+
+    if len(x_range):
+        if not ax:
+            fig = plt.figure()
+            ax = fig.add_subplot(111, projection='3d')
+
+        if c:
+            ax.scatter(x_range, y_range, z_range, c=c, s=marker_size)
+        else:
+            ax.scatter(x_range, y_range, z_range, s=marker_size)
+        # ax.scatter(x_range, y_range, z_range, s=marker_size)#, c=sample_flat)
+
+        ax.set_xlabel('X')
+        ax.set_ylabel('Y')
+        ax.set_zlabel('Z')
+        if name:
+            ax.set_title(name)
+
+        return ax
+    else:
+        print('Nothing to plot')
+        return None
+
+
+def plot_training(list_imgs: [np.ndarray], affine_transf=True, filename='img', fig=None):
+    if fig is not None:
+        fig.clear()
+        plt.figure(fig.number)
+    else:
+        fig = plt.figure(dpi=const.DPI)
+
+    ax_fix = fig.add_subplot(231)
+    im_fix = ax_fix.imshow(list_imgs[0][:, :, 0])
+    ax_fix.set_title('Fix image', fontsize=const.FONT_SIZE)
+    ax_fix.tick_params(axis='both',
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+    ax_mov = fig.add_subplot(232)
+    im_mov = ax_mov.imshow(list_imgs[1][:, :, 0])
+    ax_mov.set_title('Moving image', fontsize=const.FONT_SIZE)
+    ax_mov.tick_params(axis='both',
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+
+    ax_pred_im = fig.add_subplot(233)
+    im_pred_im = ax_pred_im.imshow(list_imgs[2][:, :, 0])
+    ax_pred_im.set_title('Prediction', fontsize=const.FONT_SIZE)
+    ax_pred_im.tick_params(axis='both',
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+
+    ax_pred_disp = fig.add_subplot(234)
+    if affine_transf:
+        fake_bg = np.array([[0.0, 0.0, 0.0, 0.0],
+                            [0.0, 0.0, 0.0, 0.0],
+                            [0.0, 0.0, 0.0, 0.0],
+                            [0.0, 0.0, 0.0, 0.0]])
+
+        bottom = np.asarray([0, 0, 0, 1])
+
+        transf_mat = np.reshape(list_imgs[3], (2, 3))
+        transf_mat = np.stack([transf_mat, bottom], axis=0)
+
+        im_pred_disp = ax_pred_disp.imshow(fake_bg)
+        for i in range(4):
+            for j in range(4):
+                ax_pred_disp.text(i, j, transf_mat[i, j], ha="center", va="center", color="b")
+
+        ax_pred_disp.set_title('Affine transformation matrix')
+
+    else:
+        cx, cy, dx, dy, s = _prepare_quiver_map(list_imgs[3])
+        im_pred_disp = ax_pred_disp.imshow(s, interpolation='none', aspect='equal')
+        ax_pred_disp.quiver(cx, cy, dx, dy, scale=const.QUIVER_PARAMS.arrow_scale)
+        ax_pred_disp.set_title('Pred disp map', fontsize=const.FONT_SIZE)
+    ax_pred_disp.tick_params(axis='both',
+                               which='both',
+                               bottom=False,
+                               left=False,
+                               labelleft=False,
+                               labelbottom=False)
+
+    ax_gt_disp = fig.add_subplot(235)
+    if affine_transf:
+        fake_bg = np.array([[0.0, 0.0, 0.0, 0.0],
+                            [0.0, 0.0, 0.0, 0.0],
+                            [0.0, 0.0, 0.0, 0.0],
+                            [0.0, 0.0, 0.0, 0.0]])
+
+        bottom = np.asarray([0, 0, 0, 1])
+
+        transf_mat = np.reshape(list_imgs[4], (2, 3))
+        transf_mat = np.stack([transf_mat, bottom], axis=0)
+
+        im_gt_disp = ax_pred_disp.imshow(fake_bg)
+        for i in range(4):
+            for j in range(4):
+                ax_pred_disp.text(i, j, transf_mat[i, j], ha="center", va="center", color="b")
+
+        ax_pred_disp.set_title('Affine transformation matrix')
+
+    else:
+        cx, cy, dx, dy, s = _prepare_quiver_map(list_imgs[4])
+        im_gt_disp = ax_gt_disp.imshow(s, interpolation='none', aspect='equal')
+        ax_gt_disp.quiver(cx, cy, dx, dy, scale=const.QUIVER_PARAMS.arrow_scale)
+        ax_gt_disp.set_title('GT disp map', fontsize=const.FONT_SIZE)
+    ax_gt_disp.tick_params(axis='both',
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+
+    cb_fix = _set_colorbar(fig, ax_fix, im_fix, False)
+    cb_mov = _set_colorbar(fig, ax_mov, im_mov, False)
+    cb_pred = _set_colorbar(fig, ax_pred_im, im_pred_im, False)
+    cb_pred_disp = _set_colorbar(fig, ax_pred_disp, im_pred_disp, False)
+    cd_gt_disp = _set_colorbar(fig, ax_gt_disp, im_gt_disp, False)
+
+    if filename is not None:
+        plt.savefig(filename, format='png')  # Call before show
+    if not const.REMOTE:
+        plt.show()
+    else:
+        plt.close()
+    return fig
+
+
+def save_centreline_img(img, title, filename, fig=None):
+    if fig is not None:
+        fig.clear()
+        plt.figure(fig.number)
+    else:
+        fig = plt.figure(dpi=const.DPI)
+
+    dim = len(img.shape[:-1])
+
+    if dim == 2:
+        ax = fig.add_subplot(111)
+        fig.suptitle(title)
+        im = ax.imshow(img[..., 0], cmap=cmap_bin)
+        ax.tick_params(axis='both',
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+
+        #cb = _set_colorbar(fig, ax, im, False)
+    else:
+        ax = fig.add_subplot(111, projection='3d')
+        fig.suptitle(title)
+        im = ax.voxels(img[0, ..., 0] > 0.0)
+        _square_3d_plot(np.arange(0, 63), np.arange(0, 63), np.arange(0, 63), ax)
+
+        ax.tick_params(axis='both',
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+
+    plt.savefig(filename, format='png')
+    plt.close()
+
+
+def save_disp_map_img(disp_map, title, filename, affine_transf=False, fig=None):
+    if fig is not None:
+        fig.clear()
+        plt.figure(fig.number)
+    else:
+        fig = plt.figure(dpi=const.DPI)
+
+    dim = disp_map.shape[-1]
+
+    if dim == 2:
+        ax_x = fig.add_subplot(131)
+        ax_x.set_title('H displacement')
+        im_x = ax_x.imshow(disp_map[..., const.H_DISP])
+        ax_x.tick_params(axis='both',
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+        cb_x = _set_colorbar(fig, ax_x, im_x, False)
+
+        ax_y = fig.add_subplot(132)
+        ax_y.set_title('W displacement')
+        im_y = ax_y.imshow(disp_map[..., const.W_DISP])
+        ax_y.tick_params(axis='both',
+                         which='both',
+                         bottom=False,
+                         left=False,
+                         labelleft=False,
+                         labelbottom=False)
+        cb_y = _set_colorbar(fig, ax_y, im_y, False)
+
+        ax = fig.add_subplot(133)
+        if affine_transf:
+            fake_bg = np.array([[0.0, 0.0, 0.0, 0.0],
+                                [0.0, 0.0, 0.0, 0.0],
+                                [0.0, 0.0, 0.0, 0.0],
+                                [0.0, 0.0, 0.0, 0.0]])
+
+            bottom = np.asarray([0, 0, 0, 1])
+
+            transf_mat = np.reshape(disp_map, (2, 3))
+            transf_mat = np.stack([transf_mat, bottom], axis=0)
+
+            im = ax.imshow(fake_bg)
+            for i in range(4):
+                for j in range(4):
+                    ax.text(i, j, transf_mat[i, j], ha="center", va="center", color="b")
+
+        else:
+            c, d, s = _prepare_quiver_map(disp_map, dim=dim)
+            im = ax.imshow(s, interpolation='none', aspect='equal')
+            ax.quiver(c[const.H_DISP], c[const.W_DISP], d[const.H_DISP], d[const.W_DISP],
+                      scale=const.QUIVER_PARAMS.arrow_scale)
+            cb = _set_colorbar(fig, ax, im, False)
+            ax.set_title('Displacement map')
+        ax.tick_params(axis='both',
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+        fig.suptitle(title)
+    else:
+        ax = fig.add_subplot(111, projection='3d')
+        c, d, s = _prepare_quiver_map(disp_map[0, ...], dim=dim)
+        ax.quiver(c[const.H_DISP], c[const.W_DISP], c[const.D_DISP], d[const.H_DISP], d[const.W_DISP], d[const.D_DISP],
+                  norm=True)
+        _square_3d_plot(np.arange(0, 63), np.arange(0, 63), np.arange(0, 63), ax)
+        fig.suptitle('Displacement map')
+        ax.tick_params(axis='both',  # Same parameters as in 2D https://matplotlib.org/api/_as_gen/mpl_toolkits.mplot3d.axes3d.Axes3D.html
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+        fig.suptitle(title)
+
+    plt.savefig(filename, format='png')
+    plt.close()
+
+
+def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, filename='img', fig=None,
+                                 title_first_row='TRAINING', title_second_row='VALIDATION'):
+    if fig is not None:
+        fig.clear()
+        plt.figure(fig.number)
+    else:
+        fig = plt.figure(dpi=const.DPI)
+
+    dim = len(list_imgs[0].shape[:-1])
+
+    if dim == 2:
+        # TRAINING
+        ax_input = fig.add_subplot(241)
+        ax_input.set_ylabel(title_first_row, fontsize=const.FONT_SIZE)
+        im_fix = ax_input.imshow(list_imgs[0][:, :, 0])
+        ax_input.set_title('Fix image', fontsize=const.FONT_SIZE)
+        ax_input.tick_params(axis='both',
+                           which='both',
+                           bottom=False,
+                           left=False,
+                           labelleft=False,
+                           labelbottom=False)
+        ax_mov = fig.add_subplot(242)
+        im_mov = ax_mov.imshow(list_imgs[1][:, :, 0])
+        ax_mov.set_title('Moving image', fontsize=const.FONT_SIZE)
+        ax_mov.tick_params(axis='both',
+                           which='both',
+                           bottom=False,
+                           left=False,
+                           labelleft=False,
+                           labelbottom=False)
+
+        ax_pred_im = fig.add_subplot(244)
+        im_pred_im = ax_pred_im.imshow(list_imgs[2][:, :, 0])
+        ax_pred_im.set_title('Predicted fix image', fontsize=const.FONT_SIZE)
+        ax_pred_im.tick_params(axis='both',
+                           which='both',
+                           bottom=False,
+                           left=False,
+                           labelleft=False,
+                           labelbottom=False)
+
+        ax_pred_disp = fig.add_subplot(243)
+        if affine_transf:
+            fake_bg = np.array([[0.0, 0.0, 0.0, 0.0],
+                                [0.0, 0.0, 0.0, 0.0],
+                                [0.0, 0.0, 0.0, 0.0],
+                                [0.0, 0.0, 0.0, 0.0]])
+
+            bottom = np.asarray([0, 0, 0, 1])
+
+            transf_mat = np.reshape(list_imgs[3], (2, 3))
+            transf_mat = np.stack([transf_mat, bottom], axis=0)
+
+            im_pred_disp = ax_pred_disp.imshow(fake_bg)
+            for i in range(4):
+                for j in range(4):
+                    ax_pred_disp.text(i, j, transf_mat[i, j], ha="center", va="center", color="b")
+
+            ax_pred_disp.set_title('Affine transformation matrix')
+
+        else:
+            cx, cy, dx, dy, s = _prepare_quiver_map(list_imgs[3], dim=dim)
+            im_pred_disp = ax_pred_disp.imshow(s, interpolation='none', aspect='equal')
+            ax_pred_disp.quiver(cx, cy, dx, dy, scale=const.QUIVER_PARAMS.arrow_scale)
+            ax_pred_disp.set_title('Pred disp map', fontsize=const.FONT_SIZE)
+        ax_pred_disp.tick_params(axis='both',
+                                   which='both',
+                                   bottom=False,
+                                   left=False,
+                                   labelleft=False,
+                                   labelbottom=False)
+
+        # VALIDATION
+        axinput_val = fig.add_subplot(245)
+        axinput_val.set_ylabel(title_second_row, fontsize=const.FONT_SIZE)
+        im_fix_val = axinput_val.imshow(list_imgs[4][:, :, 0])
+        axinput_val.set_title('Fix image', fontsize=const.FONT_SIZE)
+        axinput_val.tick_params(axis='both',
+                           which='both',
+                           bottom=False,
+                           left=False,
+                           labelleft=False,
+                           labelbottom=False)
+        ax_mov_val = fig.add_subplot(246)
+        im_mov_val = ax_mov_val.imshow(list_imgs[5][:, :, 0])
+        ax_mov_val.set_title('Moving image', fontsize=const.FONT_SIZE)
+        ax_mov_val.tick_params(axis='both',
+                           which='both',
+                           bottom=False,
+                           left=False,
+                           labelleft=False,
+                           labelbottom=False)
+
+        ax_pred_im_val = fig.add_subplot(248)
+        im_pred_im_val = ax_pred_im_val.imshow(list_imgs[6][:, :, 0])
+        ax_pred_im_val.set_title('Predicted fix image', fontsize=const.FONT_SIZE)
+        ax_pred_im_val.tick_params(axis='both',
+                           which='both',
+                           bottom=False,
+                           left=False,
+                           labelleft=False,
+                           labelbottom=False)
+
+        ax_pred_disp_val = fig.add_subplot(247)
+        if affine_transf:
+            fake_bg = np.array([[0.0, 0.0, 0.0, 0.0],
+                                [0.0, 0.0, 0.0, 0.0],
+                                [0.0, 0.0, 0.0, 0.0],
+                                [0.0, 0.0, 0.0, 0.0]])
+
+            bottom = np.asarray([0, 0, 0, 1])
+
+            transf_mat = np.reshape(list_imgs[7], (2, 3))
+            transf_mat = np.stack([transf_mat, bottom], axis=0)
+
+            im_pred_disp_val = ax_pred_disp_val.imshow(fake_bg)
+            for i in range(4):
+                for j in range(4):
+                    ax_pred_disp_val.text(i, j, transf_mat[i, j], ha="center", va="center", color="b")
+
+            ax_pred_disp_val.set_title('Affine transformation matrix')
+
+        else:
+            c, d, s = _prepare_quiver_map(list_imgs[7], dim=dim)
+            im_pred_disp_val = ax_pred_disp_val.imshow(s, interpolation='none', aspect='equal')
+            ax_pred_disp_val.quiver(c[0], c[1], d[0], d[1], scale=const.QUIVER_PARAMS.arrow_scale)
+            ax_pred_disp_val.set_title('Pred disp map', fontsize=const.FONT_SIZE)
+        ax_pred_disp_val.tick_params(axis='both',
+                                   which='both',
+                                   bottom=False,
+                                   left=False,
+                                   labelleft=False,
+                                   labelbottom=False)
+
+        cb_fix = _set_colorbar(fig, ax_input, im_fix, False)
+        cb_mov = _set_colorbar(fig, ax_mov, im_mov, False)
+        cb_pred = _set_colorbar(fig, ax_pred_im, im_pred_im, False)
+        cb_pred_disp = _set_colorbar(fig, ax_pred_disp, im_pred_disp, False)
+
+        cd_fix_val = _set_colorbar(fig, axinput_val, im_fix_val, False)
+        cb_mov_val = _set_colorbar(fig, ax_mov_val, im_mov_val, False)
+        cb_pred_val = _set_colorbar(fig, ax_pred_im_val, im_pred_im_val, False)
+        cb_pred_disp_val = _set_colorbar(fig, ax_pred_disp_val, im_pred_disp_val, False)
+
+    else:
+        # 3D
+        # TRAINING
+        ax_input = fig.add_subplot(231, projection='3d')
+        ax_input.set_ylabel(title_first_row, fontsize=const.FONT_SIZE)
+        im_fix = ax_input.voxels(list_imgs[0][..., 0] > 0.0, facecolors='red', edgecolors='red', label='Fixed')
+        im_mov = ax_input.voxels(list_imgs[1][..., 0] > 0.0, facecolors='blue', edgecolors='blue', label='Moving')
+        ax_input.set_title('Fix image', fontsize=const.FONT_SIZE)
+        ax_input.tick_params(axis='both',
+                           which='both',
+                           bottom=False,
+                           left=False,
+                           labelleft=False,
+                           labelbottom=False)
+
+        ax_pred_im = fig.add_subplot(232, projection='3d')
+        im_pred_im = ax_input.voxels(list_imgs[2][..., 0] > 0.0, facecolors='green', edgecolors='green', label='Prediction')
+        im_fix = ax_input.voxels(list_imgs[0][..., 0] > 0.0, facecolors='red', edgecolors='red', label='Fixed')
+        ax_pred_im.set_title('Predicted fix image', fontsize=const.FONT_SIZE)
+        ax_pred_im.tick_params(axis='both',
+                               which='both',
+                               bottom=False,
+                               left=False,
+                               labelleft=False,
+                               labelbottom=False)
+
+        ax_pred_disp = fig.add_subplot(233, projection='3d')
+
+        c, d, s = _prepare_quiver_map(list_imgs[3], dim=dim)
+        im_pred_disp = ax_pred_disp.imshow(s, interpolation='none', aspect='equal')
+        ax_pred_disp.quiver(c[const.H_DISP], c[const.W_DISP], c[const.D_DISP],
+                            d[const.H_DISP], d[const.W_DISP], d[const.D_DISP], scale=const.QUIVER_PARAMS.arrow_scale)
+        ax_pred_disp.set_title('Pred disp map', fontsize=const.FONT_SIZE)
+        ax_pred_disp.tick_params(axis='both',
+                                 which='both',
+                                 bottom=False,
+                                 left=False,
+                                 labelleft=False,
+                                 labelbottom=False)
+
+        # VALIDATION
+        axinput_val = fig.add_subplot(234, projection='3d')
+        axinput_val.set_ylabel(title_second_row, fontsize=const.FONT_SIZE)
+        im_fix_val = ax_input.voxels(list_imgs[4][..., 0] > 0.0, facecolors='red', edgecolors='red', label='Fixed (val)')
+        im_mov_val = ax_input.voxels(list_imgs[5][..., 0] > 0.0, facecolors='blue', edgecolors='blue', label='Moving (val)')
+        axinput_val.set_title('Fix image', fontsize=const.FONT_SIZE)
+        axinput_val.tick_params(axis='both',
+                               which='both',
+                               bottom=False,
+                               left=False,
+                               labelleft=False,
+                               labelbottom=False)
+
+        ax_pred_im_val = fig.add_subplot(235, projection='3d')
+        im_pred_im_val = ax_input.voxels(list_imgs[2][..., 0] > 0.0, facecolors='green', edgecolors='green', label='Prediction (val)')
+        im_fix_val = ax_input.voxels(list_imgs[0][..., 0] > 0.0, facecolors='red', edgecolors='red', label='Fixed (val)')
+        ax_pred_im_val.set_title('Predicted fix image', fontsize=const.FONT_SIZE)
+        ax_pred_im_val.tick_params(axis='both',
+                                   which='both',
+                                   bottom=False,
+                                   left=False,
+                                   labelleft=False,
+                                   labelbottom=False)
+
+        ax_pred_disp_val = fig.add_subplot(236, projection='3d')
+        c, d, s = _prepare_quiver_map(list_imgs[7], dim=dim)
+        im_pred_disp_val = ax_pred_disp_val.imshow(s, interpolation='none', aspect='equal')
+        ax_pred_disp_val.quiver(c[const.H_DISP], c[const.W_DISP], c[const.D_DISP],
+                                d[const.H_DISP], d[const.W_DISP], d[const.D_DISP],
+                                scale=const.QUIVER_PARAMS.arrow_scale)
+        ax_pred_disp_val.set_title('Pred disp map', fontsize=const.FONT_SIZE)
+        ax_pred_disp_val.tick_params(axis='both',
+                                     which='both',
+                                     bottom=False,
+                                     left=False,
+                                     labelleft=False,
+                                     labelbottom=False)
+
+    if filename is not None:
+        plt.savefig(filename, format='png')  # Call before show
+    if not const.REMOTE:
+        plt.show()
+    else:
+        plt.close()
+    return fig
+
+
+def _set_colorbar(fig, ax, im, drawedges=True):
+    div = make_axes_locatable(ax)
+    im_cax = div.append_axes('right', size='5%', pad=0.05)
+    im_cb = fig.colorbar(im, cax=im_cax, drawedges=drawedges, shrink=0.5, orientation='vertical')
+    im_cb.ax.tick_params(labelsize=5)
+
+    return im_cb
+
+
+def _prepare_quiver_map(disp_map: np.ndarray, dim=2, spc=const.QUIVER_PARAMS.spacing):
+    if isinstance(disp_map, tf.Tensor):
+        if tf.executing_eagerly():
+            disp_map = disp_map.numpy()
+        else:
+            disp_map = disp_map.eval()
+    dx = disp_map[..., const.H_DISP]
+    dy = disp_map[..., const.W_DISP]
+    if dim > 2:
+        dz = disp_map[..., const.D_DISP]
+
+    img_size_x = disp_map.shape[const.H_DISP]
+    img_size_y = disp_map.shape[const.W_DISP]
+    if dim > 2:
+        img_size_z = disp_map.shape[const.D_DISP]
+
+    if dim > 2:
+        s = np.sqrt(np.square(dx) + np.square(dy) + np.square(dz))
+        s = np.reshape(s, [img_size_x, img_size_y, img_size_z])
+
+        cx, cy, cz = np.meshgrid(list(range(0, img_size_x)), list(range(0, img_size_y)), list(range(0, img_size_z)),
+                                 indexing='ij')
+        c = [cx[::spc, ::spc, ::spc], cy[::spc, ::spc, ::spc], cz[::spc, ::spc, ::spc]]
+        d = [dx[::spc, ::spc, ::spc], dy[::spc, ::spc, ::spc], dz[::spc, ::spc, ::spc]]
+    else:
+        s = np.sqrt(np.square(dx) + np.square(dy))
+        s = np.reshape(s, [img_size_x, img_size_y])
+
+        cx, cy = np.meshgrid(list(range(0, img_size_x)), list(range(0, img_size_y)))
+        c = [cx[::spc, ::spc], cy[::spc, ::spc]]
+        d = [dx[::spc, ::spc], dy[::spc, ::spc]]
+
+    return c, d, s
+
+
+def _prepare_colormap(disp_map: np.ndarray):
+    if isinstance(disp_map, tf.Tensor):
+        disp_map = disp_map.eval()
+    dx = disp_map[:, :, 0]
+    dy = disp_map[:, :, 1]
+
+    mod_img = np.zeros_like(dx)
+
+    for i in range(dx.shape[0]):
+        for j in range(dx.shape[1]):
+            vec = np.asarray([dx[i, j], dy[i, j]])
+            mod_img[i, j] = np.linalg.norm(vec, ord=2)
+
+    p_l, p_h = np.percentile(mod_img, (0, 100))
+    mod_img = rescale_intensity(mod_img, in_range=(p_l, p_h), out_range=(0, 255))
+
+    return mod_img
+
+
+def plot_input_data(fix_img, mov_img, img_size=(64, 64), title=None, filename=None):
+    num_samples = fix_img.shape[0]
+
+    if num_samples != 16 and num_samples != 32:
+        raise ValueError('Only batches of 16 or 32 samples!')
+
+    fig, ax = plt.subplots(nrows=4 if num_samples == 16 else 8, ncols=4)
+    ncol = 0
+    nrow = 0
+    black_col = np.ones([img_size[0], 0])
+    for sample in range(num_samples):
+        combined_img = np.hstack([fix_img[sample, :, :, 0], black_col, mov_img[sample, :, :, 0]])
+        ax[nrow, ncol].imshow(combined_img, cmap='Greys')
+        ax[nrow, ncol].set_ylabel('#{}'.format(sample))
+        ax[nrow, ncol].tick_params(axis='both',
+                                   which='both',
+                                   bottom=False,
+                                   left=False,
+                                   labelleft=False,
+                                   labelbottom=False)
+        ncol += 1
+        if ncol >= 4:
+            ncol = 0
+            nrow += 1
+
+    if title is not None:
+        fig.suptitle(title)
+
+    if filename is not None:
+        plt.savefig(filename, format='png')  # Call before show
+    if not const.REMOTE:
+        plt.show()
+    else:
+        plt.close()
+    return fig
+
+
+def plot_dataset_orthographic_views(view_sets: [[np.ndarray]]):
+    """
+
+    :param views_fix: Expected order: top, front, left
+    :param views_mov: Expected order: top, front, left
+    :return:
+    """
+    nrows = len(view_sets)
+    fig, ax = plt.subplots(nrows=nrows, ncols=3)
+    labels = ['top', 'front', 'left']
+    for nrow in range(nrows):
+        for ncol in range(3):
+            if nrows == 1:
+                ax[ncol].imshow(view_sets[nrow][ncol][:, :, 0])
+                ax[ncol].set_title('Fix ' + labels[ncol])
+                ax[ncol].tick_params(axis='both',
+                                     which='both',
+                                     bottom=False,
+                                     left=False,
+                                     labelleft=False,
+                                     labelbottom=False)
+
+            else:
+                ax[nrow, ncol].imshow(view_sets[nrow][ncol][:, :, 0])
+                ax[nrow, ncol].set_title('Fix ' + labels[ncol])
+                ax[nrow, ncol].tick_params(axis='both',
+                                           which='both',
+                                           bottom=False,
+                                           left=False,
+                                           labelleft=False,
+                                           labelbottom=False)
+
+    plt.show()
+    return fig
+
+
+def plot_compare_2d_images(img1, img2, img1_name='img1', img2_name='img2'):
+    fig, ax = plt.subplots(nrows=1, ncols=2)
+    ax[0].imshow(img1[:, :, 0])
+    ax[0].set_title(img1_name)
+    ax[0].tick_params(axis='both',
+                      which='both',
+                      bottom=False,
+                      left=False,
+                      labelleft=False,
+                      labelbottom=False)
+
+    ax[1].imshow(img2[:, :, 0])
+    ax[1].set_title(img2_name)
+    ax[1].tick_params(axis='both',
+                      which='both',
+                      bottom=False,
+                      left=False,
+                      labelleft=False,
+                      labelbottom=False)
+
+    plt.show()
+    return fig
+
+
+def plot_dataset_3d(img_sets):
+    from mpl_toolkits.mplot3d import Axes3D
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection='3d')
+
+    for idx in range(len(img_sets)):
+        ax = _plot_3d(img_sets[idx], ax=ax, name='Set {}'.format(idx))
+
+    plt.show()
+    return fig
+
+
+def plot_predictions(fix_img_batch, mov_img_batch, disp_map_batch, pred_img_batch, filename='predictions', fig=None):
+    num_rows = fix_img_batch.shape[0]
+    img_size = fix_img_batch.shape[1:3]
+    if fig is not None:
+        fig.clear()
+        plt.figure(fig.number)
+        ax = fig.add_subplot(nrows=num_rows, ncols=4, dpi=const.DPI)
+    else:
+        fig, ax = plt.subplots(nrows=num_rows, ncols=4, dpi=const.DPI)
+
+    for row in range(num_rows):
+        fix_img = fix_img_batch[row, :, :, 0]
+        mov_img = mov_img_batch[row, :, :, 0]
+        disp_map = disp_map_batch[row, :, :, :]
+        pred_img = pred_img_batch[row, :, :, 0]
+        ax[row, 0].imshow(fix_img)
+        ax[row, 0].tick_params(axis='both',
+                               which='both',
+                               bottom=False,
+                               left=False,
+                               labelleft=False,
+                               labelbottom=False)
+        ax[row, 1].imshow(mov_img)
+        ax[row, 1].tick_params(axis='both',
+                               which='both',
+                               bottom=False,
+                               left=False,
+                               labelleft=False,
+                               labelbottom=False)
+
+        cx, cy, dx, dy, s = _prepare_quiver_map(disp_map)
+        disp_map_color = _prepare_colormap(disp_map)
+        ax[row, 2].imshow(disp_map_color, interpolation='none', aspect='equal')
+        ax[row, 2].quiver(cx.eval(), cy.eval(), dx.eval(), dy.eval(), units='xy', scale=const.QUIVER_PARAMS.arrow_scale)
+        ax[row, 2].figure.set_size_inches(img_size)
+        ax[row, 2].tick_params(axis='both',
+                               which='both',
+                               bottom=False,
+                               left=False,
+                               labelleft=False,
+                               labelbottom=False)
+
+        ax[row, 3].tick_params(axis='both',
+                               which='both',
+                               bottom=False,
+                               left=False,
+                               labelleft=False,
+                               labelbottom=False)
+
+    plt.axis('off')
+    ax[0, 0].set_title('Fixed img ($I_f$)', fontsize=const.FONT_SIZE)
+    ax[0, 1].set_title('Moving img ($I_m$)', fontsize=const.FONT_SIZE)
+    ax[0, 2].set_title('Displacement map ($\delta$)', fontsize=const.FONT_SIZE)
+    ax[0, 3].set_title('Updated $I_m$', fontsize=const.FONT_SIZE)
+
+    if filename is not None:
+        plt.savefig(filename, format='png')  # Call before show
+    if not const.REMOTE:
+        plt.show()
+    else:
+        plt.close()
+    return fig
+
+
+def inspect_disp_map_generation(fix_img, mov_img, disp_map, filename=None, fig=None):
+    if fig is not None:
+        fig.clear()
+        plt.figure(fig.number)
+    else:
+        fig = plt.figure(dpi=const.DPI)
+
+    ax0 = fig.add_subplot(221)
+    im0 = ax0.imshow(fix_img[..., 0])
+    ax0.tick_params(axis='both',
+                      which='both',
+                      bottom=False,
+                      left=False,
+                      labelleft=False,
+                      labelbottom=False)
+    ax1 = fig.add_subplot(222)
+    im1 = ax1.imshow(mov_img[..., 0])
+    ax1.tick_params(axis='both',
+                      which='both',
+                      bottom=False,
+                      left=False,
+                      labelleft=False,
+                      labelbottom=False)
+
+    cx, cy, dx, dy, s = _prepare_quiver_map(disp_map)
+    disp_map_color = _prepare_colormap(disp_map)
+    ax2 = fig.add_subplot(223)
+    im2 = ax2.imshow(s, interpolation='none', aspect='equal')
+
+    ax2.quiver(cx, cy, dx, dy, scale=const.QUIVER_PARAMS.arrow_scale)
+    # ax2.figure.set_size_inches(img_size)
+    ax2.tick_params(axis='both',
+                      which='both',
+                      bottom=False,
+                      left=False,
+                      labelleft=False,
+                      labelbottom=False)
+
+    ax3 = fig.add_subplot(224)
+    dif = fix_img[..., 0] - mov_img[..., 0]
+    im3 = ax3.imshow(dif)
+    ax3.quiver(cx, cy, dx, dy, scale=const.QUIVER_PARAMS.arrow_scale)
+    ax3.tick_params(axis='both',
+                    which='both',
+                    bottom=False,
+                    left=False,
+                    labelleft=False,
+                    labelbottom=False)
+
+    plt.axis('off')
+    ax0.set_title('Fixed img ($I_f$)', fontsize=const.FONT_SIZE)
+    ax1.set_title('Moving img ($I_m$)', fontsize=const.FONT_SIZE)
+    ax2.set_title('Displacement map', fontsize=const.FONT_SIZE)
+    ax3.set_title('Fix - Mov', fontsize=const.FONT_SIZE)
+
+    im0_cb = _set_colorbar(fig, ax0, im0, False)
+    im1_cb = _set_colorbar(fig, ax1, im1, False)
+    disp_cb = _set_colorbar(fig, ax2, im2, False)
+    im3_cb = _set_colorbar(fig, ax3, im3, False)
+
+    if filename is not None:
+        plt.savefig(filename, format='png')  # Call before show
+    if not const.REMOTE:
+        plt.show()
+    else:
+        plt.close()
+
+    return fig
+
+
+def inspect_displacement_grid(ctrl_coords, dense_coords, target_coords, disp_coords, disp_map, mask, fix_img, mov_img,
+                              filename=None, fig=None):
+    if fig is not None:
+        fig.clear()
+        plt.figure(fig.number)
+    else:
+        fig = plt.figure()
+
+    ax_grid = fig.add_subplot(231)
+    ax_grid.set_title('Grids', fontsize=const.FONT_SIZE)
+    ax_grid.scatter(ctrl_coords[:, 0], ctrl_coords[:, 1], marker='+', c='r', s=20)
+    ax_grid.scatter(dense_coords[:, 0], dense_coords[:, 1], marker='.', c='r', s=1)
+    ax_grid.tick_params(axis='both',
+                        which='both',
+                        bottom=False,
+                        left=False,
+                        labelleft=False,
+                        labelbottom=False)
+
+    ax_grid.scatter(target_coords[:, 0], target_coords[:, 1], marker='+', c='b', s=20)
+    ax_grid.scatter(disp_coords[:, 0], disp_coords[:, 1], marker='.', c='b', s=1)
+
+    ax_grid.set_aspect('equal')
+
+    ax_disp = fig.add_subplot(232)
+    ax_disp.set_title('Displacement map', fontsize=const.FONT_SIZE)
+    cx, cy, dx, dy, s = _prepare_quiver_map(disp_map)
+    ax_disp.imshow(s, interpolation='none', aspect='equal')
+    ax_disp.tick_params(axis='both',
+                        which='both',
+                        bottom=False,
+                        left=False,
+                        labelleft=False,
+                        labelbottom=False)
+
+    ax_mask = fig.add_subplot(233)
+    ax_mask.set_title('Mask', fontsize=const.FONT_SIZE)
+    ax_mask.imshow(mask)
+    ax_mask.tick_params(axis='both',
+                        which='both',
+                        bottom=False,
+                        left=False,
+                        labelleft=False,
+                        labelbottom=False)
+
+    ax_fix = fig.add_subplot(234)
+    ax_fix.set_title('Fix image', fontsize=const.FONT_SIZE)
+    ax_fix.imshow(fix_img[..., 0])
+    ax_fix.tick_params(axis='both',
+                        which='both',
+                        bottom=False,
+                        left=False,
+                        labelleft=False,
+                        labelbottom=False)
+
+    ax_mov = fig.add_subplot(235)
+    ax_mov.set_title('Moving image', fontsize=const.FONT_SIZE)
+    ax_mov.imshow(mov_img[..., 0])
+    ax_mov.tick_params(axis='both',
+                        which='both',
+                        bottom=False,
+                        left=False,
+                        labelleft=False,
+                        labelbottom=False)
+
+    ax_dif = fig.add_subplot(236)
+    ax_dif.set_title('Fix - Moving image', fontsize=const.FONT_SIZE)
+    ax_dif.imshow(fix_img[..., 0] - mov_img[..., 0], cmap=cmap_bin)
+    ax_dif.tick_params(axis='both',
+                        which='both',
+                        bottom=False,
+                        left=False,
+                        labelleft=False,
+                        labelbottom=False)
+    legend_elems = [Line2D([0], [0], color=cmap_bin(0), lw=2),
+                    Line2D([0], [0], color=cmap_bin(2), lw=2)]
+
+    ax_dif.legend(legend_elems, ['Mov', 'Fix'], loc='upper left', bbox_to_anchor=(0., 0., 1., 0.),
+                  ncol=2, mode='expand')
+
+    if filename is not None:
+        plt.savefig(filename, format='png')  # Call before show
+    if not const.REMOTE:
+        plt.show()
+
+    return fig
+
+
+def compare_disp_maps(disp_m_f, disp_f_m, fix_img, mov_img, filename=None, fig=None):
+    if fig is not None:
+        fig.clear()
+        plt.figure(fig.number)
+    else:
+        fig = plt.figure()
+
+    ax_d_m_f = fig.add_subplot(131)
+    ax_d_m_f.set_title('Disp M->F', fontsize=const.FONT_SIZE)
+    cx, cy, dx, dy, s = _prepare_quiver_map(disp_m_f)
+    ax_d_m_f.imshow(s, interpolation='none', aspect='equal')
+    ax_d_m_f.quiver(cx, cy, dx, dy, scale=const.QUIVER_PARAMS.arrow_scale)
+    ax_d_m_f.tick_params(axis='both',
+                        which='both',
+                        bottom=False,
+                        left=False,
+                        labelleft=False,
+                        labelbottom=False)
+
+    ax_d_f_m = fig.add_subplot(132)
+    ax_d_f_m.set_title('Disp F->M', fontsize=const.FONT_SIZE)
+    cx, cy, dx, dy, s = _prepare_quiver_map(disp_f_m)
+    ax_d_f_m.quiver(cx, cy, dx, dy, scale=const.QUIVER_PARAMS.arrow_scale)
+    ax_d_f_m.imshow(s, interpolation='none', aspect='equal')
+    ax_d_f_m.tick_params(axis='both',
+                         which='both',
+                         bottom=False,
+                         left=False,
+                         labelleft=False,
+                         labelbottom=False)
+
+    ax_dif = fig.add_subplot(133)
+    ax_dif.set_title('Fix - Moving image', fontsize=const.FONT_SIZE)
+    ax_dif.imshow(fix_img[..., 0] - mov_img[..., 0], cmap=cmap_bin)
+    ax_dif.tick_params(axis='both',
+                       which='both',
+                       bottom=False,
+                       left=False,
+                       labelleft=False,
+                       labelbottom=False)
+
+    legend_elems = [Line2D([0], [0], color=cmap_bin(0), lw=2),
+                    Line2D([0], [0], color=cmap_bin(2), lw=2)]
+
+    ax_dif.legend(legend_elems, ['Mov', 'Fix'], loc='upper left', bbox_to_anchor=(0., 0., 1., 0.),
+                  ncol=2, mode='expand')
+
+    if filename is not None:
+        plt.savefig(filename, format='png')  # Call before show
+    if not const.REMOTE:
+        plt.show()
+    else:
+        plt.close()
+
+    return fig
+
+
+def plot_train_step(list_imgs: [np.ndarray], fig_title='TRAINING', dest_folder='.', save_file=True):
+    # list_imgs[0]: fix image
+    # list_imgs[1]: moving image
+    # list_imgs[2]: prediction scale 1
+    # list_imgs[3]: prediction scale 2
+    # list_imgs[4]: prediction scale 3
+    # list_imgs[5]: disp map scale 1
+    # list_imgs[6]: disp map scale 2
+    # list_imgs[7]: disp map scale 3
+    num_imgs = len(list_imgs)
+    num_preds = (num_imgs - 2) // 2
+    num_cols = num_preds + 1
+    # 3D
+    # TRAINING
+    fig = plt.figure(figsize=(12.8, 10.24))
+    fig.tight_layout(pad=5.0)
+    ax = fig.add_subplot(2, num_cols, 1, projection='3d')
+    ax.voxels(list_imgs[0][0, ..., 0] > 0.0, facecolors='red', edgecolors='red', label='Fixed')
+    ax.set_title('Fix image', fontsize=const.FONT_SIZE)
+    _square_3d_plot(np.arange(0, 63), np.arange(0, 63), np.arange(0, 63), ax)
+
+    for i in range(2, num_preds+2):
+        ax = fig.add_subplot(2, num_cols, i, projection='3d')
+        ax.voxels(list_imgs[0][0, ..., 0] > 0.0, facecolors='red',  edgecolors='red', label='Fixed')
+        ax.voxels(list_imgs[i][0, ..., 0] > 0.0, facecolors='green', edgecolors='green', label='Pred_{}'.format(i - 1))
+        ax.set_title('Pred. #{}'.format(i - 1))
+        _square_3d_plot(np.arange(0, 63), np.arange(0, 63), np.arange(0, 63), ax)
+
+    ax = fig.add_subplot(2, num_cols, num_preds+2, projection='3d')
+    ax.voxels(list_imgs[1][0, ..., 0] > 0.0, facecolors='blue', edgecolors='blue', label='Moving')
+    ax.set_title('Fix image', fontsize=const.FONT_SIZE)
+    _square_3d_plot(np.arange(0, 63), np.arange(0, 63), np.arange(0, 63), ax)
+
+    for i in range(num_preds+2, 2 * num_preds + 2):
+        ax = fig.add_subplot(2, num_cols, i + 1, projection='3d')
+        c, d, s = _prepare_quiver_map(list_imgs[i][0, ...], dim=3)
+        ax.quiver(c[const.H_DISP], c[const.W_DISP], c[const.D_DISP],
+                  d[const.H_DISP], d[const.W_DISP], d[const.D_DISP],
+                  norm=True)
+        ax.set_title('Disp. #{}'.format(i - 5))
+        _square_3d_plot(np.arange(0, 63), np.arange(0, 63), np.arange(0, 63), ax)
+
+    fig.suptitle(fig_title, fontsize=const.FONT_SIZE)
+
+    if save_file:
+        plt.savefig(os.path.join(dest_folder, fig_title+'.png'), format='png')  # Call before show
+    if not const.REMOTE:
+        plt.show()
+    else:
+        plt.close()
+    return fig
+
+
+def _square_3d_plot(X, Y, Z, ax):
+    max_range = np.array([X.max() - X.min(), Y.max() - Y.min(), Z.max() - Z.min()]).max() / 2.0
+
+    mid_x = (X.max() + X.min()) * 0.5
+    mid_y = (Y.max() + Y.min()) * 0.5
+    mid_z = (Z.max() + Z.min()) * 0.5
+    ax.set_xlim(mid_x - max_range, mid_x + max_range)
+    ax.set_ylim(mid_y - max_range, mid_y + max_range)
+    ax.set_zlim(mid_z - max_range, mid_z + max_range)
+

EvaluationScripts/evaluation.py

@@ -0,0 +1,84 @@
+import os, sys
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard
+import voxelmorph as vxm
+import neurite as ne
+import h5py
+from datetime import datetime
+
+if PYCHARM_EXEC:
+    import scripts.tf.myScript_constants as const
+    from scripts.tf.myScript_data_generator import DataGeneratorManager
+    from scripts.tf.myScript_utils import save_nifti, try_mkdir
+else:
+    import myScript_constants as const
+    from myScript_data_generator import DataGeneratorManager
+    from myScript_utils import save_nifti, try_mkdir
+
+os.environ['CUDA_DEVICE_ORDER'] = const.DEV_ORDER
+os.environ['CUDA_VISIBLE_DEVICES'] = const.GPU_NUM  # Check availability before running using 'nvidia-smi'
+
+const.TRAINING_DATASET = '/mnt/EncryptedData1/Users/javier/vessel_registration/sanity_dataset_LITS'
+const.BATCH_SIZE = 8
+const.LIMIT_NUM_SAMPLES = None
+const.EPOCHS = 1000
+
+if PYCHARM_EXEC:
+    path_prefix = os.path.join('scripts', 'tf')
+else:
+    path_prefix = ''
+
+# Load data
+# Build data generator
+data_generator = DataGeneratorManager(const.TRAINING_DATASET, const.BATCH_SIZE, True, const.LIMIT_NUM_SAMPLES,
+                                      1 - const.TRAINING_PERC, voxelmorph=True)
+
+test_generator = data_generator.get_generator('validation')
+test_fix_img, test_mov_img, *_ = test_generator.get_random_sample(1)
+
+# Build model
+in_shape = test_generator.get_input_shape()[1:-1]
+enc_features = [16, 32, 32, 32]# const.ENCODER_FILTERS
+dec_features = [32, 32, 32, 32, 32, 16, 16]# const.ENCODER_FILTERS[::-1]
+nb_features = [enc_features, dec_features]
+vxm_model = vxm.networks.VxmDense(inshape=in_shape, nb_unet_features=nb_features, int_steps=0)
+
+weight_files = [os.path.join(path_prefix, 'checkpoints', f) for f in os.listdir(os.path.join(path_prefix, 'checkpoints')) if 'weights' in f]
+weight_files.sort()
+pred_folder = os.path.join(path_prefix, 'predictions')
+try_mkdir(pred_folder)
+
+# Prepare the images
+fix_img = test_fix_img.squeeze()
+mid_slice_fix = [np.take(fix_img, fix_img.shape[d]//2, axis=d) for d in range(3)]
+mid_slice_fix[1] = np.rot90(mid_slice_fix[1], 1)
+mid_slice_fix[2] = np.rot90(mid_slice_fix[2], -1)
+
+mid_mov_slice = list()
+mid_disp_slice = list()
+# Due to slicing, it can happen that the last file is not tested. So include it always
+slice = 5
+for f in weight_files[:-1:slice] + [weight_files[-1]]:
+    name = os.path.split(f)[-1].split('.h5')[0]
+    vxm_model.load_weights(f)
+    pred_img, pred_disp = vxm_model.predict([test_mov_img, test_fix_img])
+    pred_img = pred_img.squeeze()
+
+    mov_slices = [np.take(pred_img, pred_img.shape[d]//2, axis=d) for d in range(3)]
+    mov_slices[1] = np.rot90(mov_slices[1], 1)
+    mov_slices[2] = np.rot90(mov_slices[2], -1)
+    mid_mov_slice.append(mov_slices)
+
+
+
+
+
+# Get sample for testing
+test_sample = test_generator.get_single_sample()

TrainingScripts/Train_2d.py

@@ -0,0 +1,86 @@
+import os, sys
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+
+import tensorflow as tf
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+import voxelmorph as vxm
+from datetime import datetime
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.data_generator import DataGeneratorManager2D
+from DeepDeformationMapRegistration.utils.misc import try_mkdir
+from DeepDeformationMapRegistration.losses import HausdorffDistance
+
+
+os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
+os.environ['CUDA_VISIBLE_DEVICES'] = C.GPU_NUM  # Check availability before running using 'nvidia-smi'
+
+C.TRAINING_DATASET = '/mnt/EncryptedData1/Users/javier/vessel_registration/ov_dataset/training'
+C.BATCH_SIZE = 256
+C.LIMIT_NUM_SAMPLES = None
+C.EPOCHS = 10000
+
+if PYCHARM_EXEC:
+    path_prefix = os.path.join('scripts', 'tf')
+else:
+    path_prefix = ''
+
+# Load data
+# Build data generator
+sample_list = [os.path.join(C.TRAINING_DATASET, f) for f in os.listdir(C.TRAINING_DATASET) if
+               f.startswith('sample')]
+sample_list.sort()
+
+data_generator = DataGeneratorManager2D(sample_list[:C.LIMIT_NUM_SAMPLES],
+                                        C.BATCH_SIZE, C.TRAINING_PERC,
+                                        (64, 64, 1),
+                                        fix_img_tag='dilated/input/fix',
+                                        mov_img_tag='dilated/input/mov'
+                                        )
+
+# Build model
+in_shape = data_generator.train_generator.input_shape[:-1]
+enc_features = [32, 32, 32, 32, 32, 32]  # const.ENCODER_FILTERS
+dec_features = [32, 32, 32, 32, 32, 32, 32, 16]  # const.ENCODER_FILTERS[::-1]
+nb_features = [enc_features, dec_features]
+vxm_model = vxm.networks.VxmDense(inshape=in_shape, nb_unet_features=nb_features, int_steps=0)
+
+# Losses and loss weights
+def comb_loss(y_true, y_pred):
+    return 1e-3 * HausdorffDistance(ndim=2, nerosion=2).loss(y_true, y_pred) + vxm.losses.Dice().loss(y_true, y_pred)
+
+
+losses = [comb_loss, vxm.losses.Grad('l2').loss]
+loss_weights = [1, 0.01]
+
+# Compile the model
+vxm_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=losses, loss_weights=loss_weights)
+
+# Train
+output_folder = os.path.join('train_2d_dice_hausdorff_grad_'+datetime.now().strftime("%H%M%S-%d%m%Y"))
+try_mkdir(output_folder)
+try_mkdir(os.path.join(output_folder, 'checkpoints'))
+try_mkdir(os.path.join(output_folder, 'tensorboard'))
+my_callbacks = [
+    # EarlyStopping(patience=const.EARLY_STOP_PATIENCE, monitor='dice', mode='max', verbose=1),
+    ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'best_model.h5'),
+                    save_best_only=True, monitor='val_loss', verbose=0, mode='min'),
+    ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'weights.{epoch:05d}-{val_loss:.2f}.h5'),
+                    save_best_only=True, save_weights_only=True, monitor='val_loss', verbose=0, mode='min'),
+    # CSVLogger(train_log_name, ';'),
+    # UpdateLossweights([haus_weight, dice_weight], [const.MODEL+'_resampler_seg', const.MODEL+'_resampler_seg'])
+    TensorBoard(log_dir=os.path.join(output_folder, 'tensorboard'),
+                batch_size=C.BATCH_SIZE, write_images=True, histogram_freq=10, update_freq='epoch',
+                write_grads=True),
+    EarlyStopping(monitor='val_loss', verbose=1, patience=50, min_delta=0.0001)
+]
+hist = vxm_model.fit_generator(data_generator.train_generator,
+                               epochs=C.EPOCHS,
+                               validation_data=data_generator.validation_generator,
+                               verbose=2,
+                               callbacks=my_callbacks)

TrainingScripts/Train_2d_uncertaintyWeighting.py

@@ -0,0 +1,103 @@
+import os, sys
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+import voxelmorph as vxm
+import neurite as ne
+import h5py
+from datetime import datetime
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.data_generator import DataGeneratorManager2D
+from DeepDeformationMapRegistration.utils.misc import try_mkdir
+from DeepDeformationMapRegistration.losses import HausdorffDistance
+from DeepDeformationMapRegistration.layers import UncertaintyWeighting
+
+
+os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
+os.environ['CUDA_VISIBLE_DEVICES'] = '1' # const.GPU_NUM  # Check availability before running using 'nvidia-smi'
+
+C.TRAINING_DATASET = '/mnt/EncryptedData1/Users/javier/vessel_registration/ov_dataset/training'
+C.BATCH_SIZE = 256
+C.LIMIT_NUM_SAMPLES = None
+C.EPOCHS = 10000
+
+if PYCHARM_EXEC:
+    path_prefix = os.path.join('scripts', 'tf')
+else:
+    path_prefix = ''
+
+# Load data
+# Build data generator
+sample_list = [os.path.join(C.TRAINING_DATASET, f) for f in os.listdir(C.TRAINING_DATASET) if
+               f.startswith('sample')]
+sample_list.sort()
+
+data_generator = DataGeneratorManager2D(sample_list[:C.LIMIT_NUM_SAMPLES],
+                                        C.BATCH_SIZE, C.TRAINING_PERC,
+                                        (64, 64, 1),
+                                        fix_img_tag='dilated/input/fix',
+                                        mov_img_tag='dilated/input/mov',
+                                        multi_loss=True,
+                                        )
+
+# Build model
+in_shape_img, in_shape_grad = data_generator.train_generator.input_shape
+enc_features = [32, 32, 32, 32, 32, 32]  # const.ENCODER_FILTERS
+dec_features = [32, 32, 32, 32, 32, 32, 32, 16]  # const.ENCODER_FILTERS[::-1]
+nb_features = [enc_features, dec_features]
+vxm_model = vxm.networks.VxmDense(inshape=in_shape_img[:-1], nb_unet_features=nb_features, int_steps=0)
+
+#moving = tf.keras.Input(shape=in_shape_img, name='multiLoss_moving_input', dtype=tf.float32)
+#fixed = tf.keras.Input(shape=in_shape_img, name='multiLoss_fixed_input', dtype=tf.float32)
+grad = tf.keras.Input(shape=(*in_shape_img[:-1], 2), name='multiLoss_grad_input', dtype=tf.float32)
+
+def dice_loss(y_true, y_pred):
+    # Dice().loss returns -Dice score
+    return 1 + vxm.losses.Dice().loss(y_true, y_pred)
+
+#fixed_pred, dm_pred = vxm_model([moving, fixed])
+multiLoss = UncertaintyWeighting(num_loss_fns=2,
+                                 num_reg_fns=1,
+                                 loss_fns=[HausdorffDistance(2, 2).loss, dice_loss],
+                                 reg_fns=[vxm.losses.Grad('l2').loss],
+                                 prior_loss_w=[1., 1.],
+                                 prior_reg_w=[0.01],
+                                 name='MultiLossLayer')
+loss = multiLoss([vxm_model.inputs[1], vxm_model.inputs[1], vxm_model.references.y_source, vxm_model.references.y_source, grad, vxm_model.references.pos_flow])
+
+full_model = tf.keras.Model(inputs=vxm_model.inputs + [grad], outputs=vxm_model.outputs + [loss])
+
+# Compile the model
+full_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=None)
+
+# Train
+output_folder = os.path.join('train_2d_multiloss_haussdorf_dice_grad' + datetime.now().strftime("%H%M%S-%d%m%Y"))
+try_mkdir(output_folder)
+try_mkdir(os.path.join(output_folder, 'checkpoints'))
+try_mkdir(os.path.join(output_folder, 'tensorboard'))
+my_callbacks = [
+    # EarlyStopping(patience=const.EARLY_STOP_PATIENCE, monitor='dice', mode='max', verbose=1),
+    ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'best_model.h5'),
+                    save_best_only=True, monitor='val_loss', verbose=0, mode='min'),
+    ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'weights.{epoch:05d}-{val_loss:.2f}.h5'),
+                    save_best_only=True, save_weights_only=True, monitor='val_loss', verbose=0, mode='min'),
+    # CSVLogger(train_log_name, ';'),
+    # UpdateLossweights([haus_weight, dice_weight], [const.MODEL+'_resampler_seg', const.MODEL+'_resampler_seg'])
+    TensorBoard(log_dir=os.path.join(output_folder, 'tensorboard'),
+                batch_size=C.BATCH_SIZE, write_images=True, histogram_freq=10, update_freq='epoch',
+                write_grads=True),
+    EarlyStopping(monitor='val_loss', verbose=1, patience=50, min_delta=0.0001)
+]
+hist = full_model.fit_generator(data_generator.train_generator,
+                                epochs=C.EPOCHS,
+                                validation_data=data_generator.validation_generator,
+                                verbose=2,
+                                callbacks=my_callbacks)

TrainingScripts/Train_3d.py

@@ -0,0 +1,76 @@
+import os, sys
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+import voxelmorph as vxm
+import neurite as ne
+import h5py
+from datetime import datetime
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.data_generator import DataGeneratorManager
+from DeepDeformationMapRegistration.utils.misc import try_mkdir
+
+
+os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
+os.environ['CUDA_VISIBLE_DEVICES'] = '1'  # Check availability before running using 'nvidia-smi'
+
+C.TRAINING_DATASET = '/mnt/EncryptedData1/Users/javier/vessel_registration/sanity_dataset_LITS'
+C.BATCH_SIZE = 2
+C.LIMIT_NUM_SAMPLES = None
+C.EPOCHS = 10000
+
+# Load data
+# Build data generator
+data_generator = DataGeneratorManager(C.TRAINING_DATASET, C.BATCH_SIZE, True, C.LIMIT_NUM_SAMPLES,
+                                      1 - C.TRAINING_PERC, voxelmorph=True)
+
+train_generator = data_generator.get_generator('train')
+validation_generator = data_generator.get_generator('validation')
+
+
+# Build model
+in_shape = train_generator.get_input_shape()[1:-1]
+enc_features = [16, 32, 32, 32, 32, 32]# const.ENCODER_FILTERS
+dec_features = [32, 32, 32, 32, 32, 32, 32, 16, 16]# const.ENCODER_FILTERS[::-1]
+nb_features = [enc_features, dec_features]
+vxm_model = vxm.networks.VxmDense(inshape=in_shape, nb_unet_features=nb_features, int_steps=7)
+
+
+# Losses and loss weights
+
+def comb_loss(y_true, y_pred):
+    return vxm.losses.MSE().loss(y_true, y_pred) + vxm.losses.NCC().loss(y_true, y_pred)
+
+
+losses = [comb_loss, vxm.losses.Grad('l2').loss]
+loss_weights = [1., 0.01]
+
+# Compile the model
+vxm_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=losses, loss_weights=loss_weights)
+
+# Train
+output_folder = os.path.join('train_3d_mse_ncc_grad_'+datetime.now().strftime("%H%M%S-%d%m%Y"))
+try_mkdir(output_folder)
+try_mkdir(os.path.join(output_folder, 'checkpoints'))
+try_mkdir(os.path.join(output_folder, 'tensorboard'))
+my_callbacks = [
+        #EarlyStopping(patience=const.EARLY_STOP_PATIENCE, monitor='dice', mode='max', verbose=1),
+        ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'best_model.h5'),
+                        save_best_only=True, monitor='val_loss', verbose=0, mode='min'),
+        ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'weights.{epoch:05d}-{val_loss:.2f}.h5'),
+                        save_best_only=True, save_weights_only=True, monitor='val_loss', verbose=0, mode='min'),
+        # CSVLogger(train_log_name, ';'),
+        # UpdateLossweights([haus_weight, dice_weight], [const.MODEL+'_resampler_seg', const.MODEL+'_resampler_seg'])
+        TensorBoard(log_dir=os.path.join(output_folder, 'tensorboard'),
+                    batch_size=C.BATCH_SIZE, write_images=False, histogram_freq=10, update_freq='epoch',
+                    write_grads=True),
+        EarlyStopping(monitor='val_loss', verbose=1, patience=50, min_delta=0.0001)
+    ]
+hist = vxm_model.fit(train_generator, epochs=C.EPOCHS, validation_data=validation_generator, verbose=2, callbacks=my_callbacks)

TrainingScripts/Train_3d_weaklySupervised.py

@@ -0,0 +1,92 @@
+import os, sys
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping
+import voxelmorph as vxm
+import neurite as ne
+import h5py
+from datetime import datetime
+
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.data_generator import DataGeneratorManager
+from DeepDeformationMapRegistration.utils.misc import try_mkdir
+from DeepDeformationMapRegistration.networks import VxmWeaklySupervised
+from DeepDeformationMapRegistration.losses import HausdorffDistance
+from DeepDeformationMapRegistration.layers import UncertaintyWeighting
+
+
+os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # Check availability before running using 'nvidia-smi'
+
+C.TRAINING_DATASET = '/mnt/EncryptedData1/Users/javier/vessel_registration/sanity_dataset_vessels'
+C.BATCH_SIZE = 2
+C.LIMIT_NUM_SAMPLES = None
+C.EPOCHS = 10000
+
+# Load data
+# Build data generator
+
+data_generator = DataGeneratorManager(C.TRAINING_DATASET, C.BATCH_SIZE, True, C.LIMIT_NUM_SAMPLES,
+                                      1 - C.TRAINING_PERC, voxelmorph=True, segmentations=True)
+
+train_generator = data_generator.get_generator('train')
+validation_generator = data_generator.get_generator('validation')
+
+
+# Build model
+in_shape = train_generator.get_input_shape()[1:-1]
+enc_features = [16, 32, 32, 32, 32, 32]# const.ENCODER_FILTERS
+dec_features = [32, 32, 32, 32, 32, 32, 32, 16, 16]# const.ENCODER_FILTERS[::-1]
+nb_features = [enc_features, dec_features]
+vxm_model = VxmWeaklySupervised(inshape=in_shape, all_labels=[1], nb_unet_features=nb_features, int_steps=5)
+
+# Losses and loss weights
+
+grad = tf.keras.Input(shape=(*in_shape, 3), name='multiLoss_grad_input', dtype=tf.float32)
+fix_img = tf.keras.Input(shape=(*in_shape, 1), name='multiLoss_fix_img_input', dtype=tf.float32)
+def dice_loss(y_true, y_pred):
+    # Dice().loss returns -Dice score
+    return 1 + vxm.losses.Dice().loss(y_true, y_pred)
+
+multiLoss = UncertaintyWeighting(num_loss_fns=3,
+                                 num_reg_fns=1,
+                                 loss_fns=[HausdorffDistance(3, 5).loss, dice_loss, vxm.losses.NCC().loss],
+                                 reg_fns=[vxm.losses.Grad('l2').loss],
+                                 prior_loss_w=[1., 1., 1.],
+                                 prior_reg_w=[0.01],
+                                 name='MultiLossLayer')
+loss = multiLoss([vxm_model.inputs[1], vxm_model.inputs[1], fix_img,
+                  vxm_model.references.pred_segm, vxm_model.references.pred_segm, vxm_model.references.pred_img,
+                  grad,
+                  vxm_model.references.pos_flow])
+
+full_model = tf.keras.Model(inputs=vxm_model.inputs + [fix_img, grad], outputs=vxm_model.outputs + [loss])
+
+# Compile the model
+full_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=None)
+
+# Train
+output_folder = os.path.join('train_3d_multiloss_segm_haus_dice_ncc_grad_'+datetime.now().strftime("%H%M%S-%d%m%Y"))
+try_mkdir(output_folder)
+try_mkdir(os.path.join(output_folder, 'checkpoints'))
+try_mkdir(os.path.join(output_folder, 'tensorboard'))
+my_callbacks = [
+        #EarlyStopping(patience=const.EARLY_STOP_PATIENCE, monitor='dice', mode='max', verbose=1),
+        ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'best_model.h5'),
+                        save_best_only=True, monitor='val_loss', verbose=0, mode='min'),
+        ModelCheckpoint(os.path.join(output_folder, 'checkpoints', 'weights.{epoch:05d}-{val_loss:.2f}.h5'),
+                        save_best_only=True, save_weights_only=True, monitor='val_loss', verbose=0, mode='min'),
+        # CSVLogger(train_log_name, ';'),
+        # UpdateLossweights([haus_weight, dice_weight], [const.MODEL+'_resampler_seg', const.MODEL+'_resampler_seg'])
+        TensorBoard(log_dir=os.path.join(output_folder, 'tensorboard'),
+                    batch_size=C.BATCH_SIZE, write_images=False, histogram_freq=10, update_freq='epoch',
+                    write_grads=True),
+        EarlyStopping(monitor='val_loss', verbose=1, patience=50, min_delta=0.0001)
+    ]
+hist = full_model.fit(train_generator, epochs=C.EPOCHS, validation_data=validation_generator, verbose=2, callbacks=my_callbacks)