Update DeepDeformationMapRegistration package

DeepDeformationMapRegistration/callbacks.py

@@ -1,5 +1,7 @@
 import tensorflow as tf
 import tensorflow.keras.backend as K
+import logging
+import numpy as np
 
 
 class RollingAverageWeighting(tf.keras.callbacks.Callback):
@@ -43,3 +45,16 @@ class RollingAverageWeighting(tf.keras.callbacks.Callback):
         for name, val in zip(self.loss_weights.keys(), new_weights):
             out_str += '{}: {:7.2f}\t'.format(name, val)
         print('WEIGHTS UPDATE: ' + out_str)
+
+
+class UncertaintyWeightingRollingAverageCallback(tf.keras.callbacks.Callback):
+    def __init__(self, method, epoch_update):
+        super(UncertaintyWeightingRollingAverageCallback, self).__init__()
+        self.method = method
+        self.epoch_update = epoch_update
+
+    def on_epoch_end(self, epoch, logs=None):
+        if epoch > self.epoch_update:
+            self.method()
+            print('Calling method: '+self.method.__name__)
+

DeepDeformationMapRegistration/data_generator.py

@@ -4,9 +4,16 @@ import os
 import h5py
 import random
 from PIL import Image
+import nibabel as nib
+from nilearn.image import resample_img
+from skimage.exposure import equalize_adapthist
+from scipy.ndimage import zoom
+import tensorflow as tf
 
 import DeepDeformationMapRegistration.utils.constants as C
 from DeepDeformationMapRegistration.utils.operators import min_max_norm
+from DeepDeformationMapRegistration.utils.thin_plate_splines import ThinPlateSplines
+from voxelmorph.tf.layers import SpatialTransformer
 
 
 class DataGeneratorManager(keras.utils.Sequence):
@@ -113,7 +120,7 @@ class DataGeneratorManager(keras.utils.Sequence):
             file_list.sort()
             for data_file in files:
                 file_name, extension = os.path.splitext(data_file)
-                if extension.lower() == '.hd5':
+                if extension.lower() == '.hd5' or '.h5':
                     file_list.append(os.path.join(root, data_file))
 
         if not file_list:
@@ -230,12 +237,7 @@ class DataGenerator(DataGeneratorManager):
                 ret_list.append(np.zeros([data_dict['BATCH_SIZE'], *C.DISP_MAP_SHAPE]))
         return ret_list
 
-    def __getitem__(self, index):
-        """
-        Generate one batch of data
-        :param index: epoch index
-        :return:
-        """
+    def __getitem1(self, index):
         idxs = self.__internal_idxs[index * self.__batch_size:(index + 1) * self.__batch_size]
 
         data_dict = self.__load_data(idxs)
@@ -248,6 +250,14 @@ class DataGenerator(DataGeneratorManager):
 
         return (inputs, outputs)
 
+    def __getitem__(self, index):
+        """
+        Generate one batch of data
+        :param index: epoch index
+        :return:
+        """
+        return self.__getitem2(index)
+
     def next_batch(self):
         if self.__last_batch > self.__batches_per_epoch:
             raise ValueError('No more batches for this epoch')
@@ -259,15 +269,15 @@ class DataGenerator(DataGeneratorManager):
         if label in self.__input_labels or label in self.__output_labels:
             # To avoid extra overhead
             try:
-                retVal = data_file[label][:]
+                retVal = data_file[label][:][np.newaxis, ...]
             except KeyError:
                 # That particular label is not found in the file. But this should be known by the user by now
                 retVal = None
 
             if append_array is not None and retVal is not None:
-                return np.append(append_array, [data_file[C.H5_FIX_IMG][:]], axis=0)
+                return np.append(append_array, retVal, axis=0)
             elif append_array is None:
-                return retVal[np.newaxis, ...]
+                return retVal
             else:
                 return retVal  # None
         else:
@@ -280,19 +290,22 @@ class DataGenerator(DataGeneratorManager):
         :return:
         """
         if isinstance(idx_list, (list, np.ndarray)):
-            fix_img = np.empty((0, ) + C.IMG_SHAPE)
-            mov_img = np.empty((0, ) + C.IMG_SHAPE)
+            fix_img = np.empty((0, ) + C.IMG_SHAPE, np.float32)
+            mov_img = np.empty((0, ) + C.IMG_SHAPE, np.float32)
+
+            fix_parench = np.empty((0, ) + C.IMG_SHAPE, np.float32)
+            mov_parench = np.empty((0, ) + C.IMG_SHAPE, np.float32)
 
-            fix_parench = np.empty((0, ) + C.IMG_SHAPE)
-            mov_parench = np.empty((0, ) + C.IMG_SHAPE)
+            fix_vessels = np.empty((0, ) + C.IMG_SHAPE, np.float32)
+            mov_vessels = np.empty((0, ) + C.IMG_SHAPE, np.float32)
 
-            fix_vessels = np.empty((0, ) + C.IMG_SHAPE)
-            mov_vessels = np.empty((0, ) + C.IMG_SHAPE)
+            fix_tumors = np.empty((0, ) + C.IMG_SHAPE, np.float32)
+            mov_tumors = np.empty((0, ) + C.IMG_SHAPE, np.float32)
 
-            fix_tumors = np.empty((0, ) + C.IMG_SHAPE)
-            mov_tumors = np.empty((0, ) + C.IMG_SHAPE)
+            disp_map = np.empty((0, ) + C.DISP_MAP_SHAPE, np.float32)
 
-            disp_map = np.empty((0, ) + C.DISP_MAP_SHAPE)
+            fix_centroid = np.empty((0, 3))
+            mov_centroid = np.empty((0, 3))
 
             for idx in idx_list:
                 data_file = h5py.File(self.__list_files[idx], 'r')
@@ -306,11 +319,14 @@ class DataGenerator(DataGeneratorManager):
                 fix_vessels = self.__try_load(data_file, C.H5_FIX_VESSELS_MASK, fix_vessels)
                 mov_vessels = self.__try_load(data_file, C.H5_MOV_VESSELS_MASK, mov_vessels)
 
-                fix_tumors = self.__try_load(data_file, C.H5_FIX_TUMORS_MASK, mov_parench)
-                mov_tumors = self.__try_load(data_file, C.H5_MOV_TUMORS_MASK, mov_parench)
+                fix_tumors = self.__try_load(data_file, C.H5_FIX_TUMORS_MASK, fix_tumors)
+                mov_tumors = self.__try_load(data_file, C.H5_MOV_TUMORS_MASK, mov_tumors)
 
                 disp_map = self.__try_load(data_file, C.H5_GT_DISP, disp_map)
 
+                fix_centroid = self.__try_load(data_file, C.H5_FIX_CENTROID, fix_centroid)
+                mov_centroid = self.__try_load(data_file, C.H5_MOV_CENTROID, mov_centroid)
+
                 data_file.close()
             batch_size = len(idx_list)
         else:
@@ -330,6 +346,9 @@ class DataGenerator(DataGeneratorManager):
 
             disp_map = self.__try_load(data_file, C.H5_GT_DISP)
 
+            fix_centroid = self.__try_load(data_file, C.H5_FIX_CENTROID)
+            mov_centroid = self.__try_load(data_file, C.H5_MOV_CENTROID)
+
             data_file.close()
             batch_size = 1
 
@@ -342,11 +361,61 @@ class DataGenerator(DataGeneratorManager):
                      C.H5_MOV_VESSELS_MASK: mov_vessels,
                      C.H5_MOV_PARENCHYMA_MASK: mov_parench,
                      C.H5_GT_DISP: disp_map,
+                     C.H5_FIX_CENTROID: fix_centroid,
+                     C.H5_MOV_CENTROID: mov_centroid,
                      'BATCH_SIZE': batch_size
                      }
 
         return data_dict
 
+    @staticmethod
+    def __get_data_shape(file_path, label):
+        f = h5py.File(file_path, 'r')
+        shape = f[label][:].shape
+        f.close()
+        return shape
+
+    def __load_data_by_label(self, label, idx_list):
+        if isinstance(idx_list, (list, np.ndarray)):
+            data_shape = self.__get_data_shape(self.__list_files[idx_list[0]], label)
+            container = np.empty((0, *data_shape), np.float32)
+            # if label == C.H5_GT_DISP:
+            #     container = np.empty((0, ) + C.DISP_MAP_SHAPE, np.float32)
+            # elif label == C.H5_MOV_CENTROID or label == C.H5_FIX_CENTROID:
+            #     container = np.empty((0, 3), np.float32)
+            # else:
+            #     container = np.empty((0, ) + C.IMG_SHAPE, np.float32)
+
+            for idx in idx_list:
+                data_file = h5py.File(self.__list_files[idx], 'r')
+                container = self.__try_load(data_file, label, container)
+                data_file.close()
+        else:
+            data_file = h5py.File(self.__list_files[idx_list], 'r')
+            container = self.__try_load(data_file, label)
+            data_file.close()
+
+        return container
+
+    def __build_list2(self, label_list, file_idxs):
+        ret_list = list()
+        for label in label_list:
+            if label is C.DG_LBL_ZERO_GRADS:
+                aux = np.zeros([len(file_idxs), *C.DISP_MAP_SHAPE])
+            else:
+                aux = self.__load_data_by_label(label, file_idxs)
+
+                if label in [C.DG_LBL_MOV_IMG, C.DG_LBL_FIX_IMG]:
+                    aux = min_max_norm(aux).astype(np.float32)
+            ret_list.append(aux)
+        return ret_list
+
+    def __getitem2(self, index):
+        f_indices = self.__internal_idxs[index * self.__batch_size:(index + 1) * self.__batch_size]
+
+        return self.__build_list2(self.__input_labels, f_indices), self.__build_list2(self.__output_labels, f_indices)
+
+
     def get_samples(self, num_samples, random=False):
         if random:
             idxs = np.random.randint(0, self.__num_samples, num_samples)
@@ -509,3 +578,345 @@ class DataGenerator2D(keras.utils.Sequence):
         return mov, fix
 
 
+FILE_EXT = {'nifti': '.nii.gz',
+            'h5': '.h5'}
+CTRL_GRID = C.CoordinatesGrid()
+CTRL_GRID.set_coords_grid([128]*3, [C.TPS_NUM_CTRL_PTS_PER_AXIS]*3, batches=False, norm=False, img_type=tf.float32)
+
+FINE_GRID = C.CoordinatesGrid()
+FINE_GRID.set_coords_grid([128]*3, [128]*3, batches=FINE_GRID, norm=False)
+
+class DataGeneratorAugment(DataGeneratorManager):
+    def __init__(self, GeneratorManager: DataGeneratorManager, file_type='nifti', 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.__augm_per_sample = 10
+        self.__samples_per_batch = np.ceil(self.__batch_size / (self.__augm_per_sample + 1))  # B = S + S*A
+        self.__total_samples = len(self.__list_files)
+        self.__clip_range = GeneratorManager.clip_rage
+        self.__manager = GeneratorManager
+        self.__shuffle = GeneratorManager.shuffle
+        self.__file_extension = FILE_EXT[file_type]
+
+        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.__input_labels = GeneratorManager.input_labels
+        self.__output_labels = GeneratorManager.output_labels
+
+
+    def __get_dataset_files(self, 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() == self.__file_extension:
+                    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:
+        """
+        return self.__getitem(index)
+
+    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 __try_load(self, data_file, label, append_array=None):
+        if label in self.__input_labels or label in self.__output_labels:
+            # To avoid extra overhead
+            try:
+                retVal = data_file[label][:][np.newaxis, ...]
+            except KeyError:
+                # That particular label is not found in the file. But this should be known by the user by now
+                retVal = None
+
+            if append_array is not None and retVal is not None:
+                return np.append(append_array, retVal, axis=0)
+            elif append_array is None:
+                return retVal
+            else:
+                return retVal  # None
+        else:
+            return None
+
+    @staticmethod
+    def __get_data_shape(file_path, label):
+        f = h5py.File(file_path, 'r')
+        shape = f[label][:].shape
+        f.close()
+        return shape
+
+    def __load_data_by_label(self, label, idx_list):
+        if isinstance(idx_list, (list, np.ndarray)):
+            data_shape = self.__get_data_shape(self.__list_files[idx_list[0]], label)
+            container = np.empty((0, *data_shape), np.float32)
+            # if label == C.H5_GT_DISP:
+            #     container = np.empty((0, ) + C.DISP_MAP_SHAPE, np.float32)
+            # elif label == C.H5_MOV_CENTROID or label == C.H5_FIX_CENTROID:
+            #     container = np.empty((0, 3), np.float32)
+            # else:
+            #     container = np.empty((0, ) + C.IMG_SHAPE, np.float32)
+
+            for idx in idx_list:
+                data_file = h5py.File(self.__list_files[idx], 'r')
+                container = self.__try_load(data_file, label, container)
+                data_file.close()
+        else:
+            data_file = h5py.File(self.__list_files[idx_list], 'r')
+            container = self.__try_load(data_file, label)
+            data_file.close()
+
+        return container
+
+    def __build_list(self, label_list, file_idxs):
+        ret_list = list()
+        for label in label_list:
+            if label is C.DG_LBL_ZERO_GRADS:
+                aux = np.zeros([len(file_idxs), *C.DISP_MAP_SHAPE])
+            else:
+                aux = self.__load_data_by_label(label, file_idxs)
+
+                if label in [C.DG_LBL_MOV_IMG, C.DG_LBL_FIX_IMG]:
+                    aux = min_max_norm(aux).astype(np.float32)
+            ret_list.append(aux)
+        return ret_list
+
+    def __getitem(self, index):
+        f_indices = self.__internal_idxs[index * self.__samples_per_batch:(index + 1) * self.__samples_per_batch]
+        # 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 'h5' in self.__file_extension:
+            return self.__build_list(self.__input_labels, f_indices), self.__build_list(self.__output_labels, f_indices)
+        else:
+            f_list = [self.__list_files[i] for i in f_indices]
+            return self.__augment(f_list, 'fixed', C.H5_FIX_IMG), self.__augment(f_list, 'moving', C.H5_FIX_IMG)
+
+
+    def __intensity_preprocessing(self, img_data):
+        # Histogram normalization
+        processed_img = equalize_adapthist(img_data, clip_limit=0.03)
+        processed_img = min_max_norm(processed_img)
+
+        return processed_img
+
+
+    def __resize_img(self, img, output_shape):
+        if isinstance(output_shape, int):
+            output_shape = [output_shape] * len(img.shape)
+        # Resize
+        zoom_vals = np.asarray(output_shape) / np.asarray(img.shape)
+        return zoom(img, zoom_vals)
+
+
+    def __build_augmented_batch(self, f_list, mode):
+        for f_path in f_list:
+            h5_file = h5py.File(f_path, 'r')
+            img_nib = nib.load(h5_file[C.H5_FIX_IMG][:])
+            img_nib = resample_img(img_nib, np.eye(3))
+            try:
+                seg_nib = nib.load(h5_file[C.H5_FIX_SEGMENTATIONS][:])
+                seg_nib = resample_img(seg_nib, np.eye(3))
+            except FileNotFoundError:
+                seg_nib = None
+
+            img_nib = self.__intensity_preprocessing(img_nib)
+            img_nib = self.__resize_img(img_nib, 128)
+
+
+
+
+
+
+
+
+    def get_samples(self, num_samples, random=False):
+        return
+
+    def get_input_shape(self):
+        input_batch, _ = self.__getitem__(0)
+        data_dict = self.__load_data(0)
+
+        ret_val = data_dict[self.__input_labels[0]].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
+
+
+def tf_graph_deform():
+    # Place holders
+    fix_img = tf.placeholder(tf.float32, [128]*3, 'fix_img')
+    fix_segmentations = tf.placeholder_with_default(np.zeros([128]*3), shape=[128]*3, name='fix_segmentations')
+    max_deformation = tf.placeholder(tf.float32, shape=(), name='max_deformation')
+    max_displacement = tf.placeholder(tf.float32, shape=(), name='max_displacement')
+    max_rotation = tf.placeholder(tf.float32, shape=(), name='max_rotation')
+    num_moved_points = tf.placeholder_with_default(50, shape=(), name='num_moved_points')
+    only_image = tf.placeholder_with_default(True, shape=(), name='only_image')
+
+    search_voxels = tf.cond(only_image,
+                            lambda: fix_img,
+                            lambda: fix_segmentations)
+
+    # Apply TPS deformation
+    # Get points in the segmentation or image, and add it to the control grid and target grid
+    # Indices of the points in the seaerch image with intensity greater than 0  (It would be bad if we only move the bg)
+    idx_points_in_label = tf.where(tf.greater(search_voxels, 0.0))
+
+    # Randomly select one of the points
+    random_idx = tf.random.uniform((num_moved_points,), minval=0, maxval=tf.shape(idx_points_in_label)[0], dtype=tf.int32)
+
+    disp_location = tf.gather_nd(idx_points_in_label, random_idx)  # And get the coordinates
+    disp_location = tf.cast(disp_location, tf.float32)
+    # Get the coordinates of the control point displaces
+    rand_disp = tf.random.uniform((num_moved_points, 3), minval=-1, maxval=1, dtype=tf.float32) * max_deformation
+    warped_location = disp_location + rand_disp
+
+    # Add the selected locations to the control grid and the warped locations to the target grid
+    control_grid = tf.concat([CTRL_GRID.grid_flat(), disp_location], axis=0)
+    trg_grid = tf.concat([CTRL_GRID.grid_flat(), warped_location], axis=0)
+
+    # Add global affine transformation
+    trg_grid, aff = transform_points(trg_grid, max_displacement=max_displacement, max_rotation=max_rotation)
+
+    tps = ThinPlateSplines(control_grid, trg_grid)
+    def_grid = tps.interpolate(FINE_GRID.grid_flat())
+
+    disp_map = FINE_GRID.grid_flat() - def_grid
+    disp_map = tf.reshape(disp_map, (*FINE_GRID.shape, -1))
+    # disp_map = interpn(disp_map, FULL_FINE_GRID.grid)
+
+    # add the batch and channel dimensions
+    fix_img = tf.expand_dims(tf.expand_dims(fix_img, -1), 0)
+    fix_segmentations = tf.expand_dims(tf.expand_dims(fix_img, -1), 0)
+    disp_map = tf.cast(tf.expand_dims(disp_map, 0), tf.float32)
+
+    mov_img = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([fix_img, disp_map])
+    mov_segmentations = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([fix_segmentations, disp_map])
+
+    return tf.squeeze(mov_img),\
+           tf.squeeze(mov_segmentations),\
+           tf.squeeze(disp_map),\
+           disp_location,\
+           rand_disp,\
+           aff #, w, trg_grid, def_grid
+
+
+def transform_points(points: tf.Tensor, max_displacement, max_rotation):
+    axis = tf.random.uniform((), 0, 3)
+
+    alpha = tf.cond(tf.less_equal(axis, 0.),
+                    lambda: tf.random.uniform((1,), -max_rotation, max_rotation),
+                    lambda: tf.zeros((1,), tf.float32))
+    beta = tf.cond(tf.less_equal(axis, 1.),
+                   lambda: tf.random.uniform((1,), -max_rotation, max_rotation),
+                   lambda: tf.zeros((1,), tf.float32))
+    gamma = tf.cond(tf.less_equal(axis, 2.),
+                    lambda: tf.random.uniform((1,), -max_rotation, max_rotation),
+                    lambda: tf.zeros((1,), tf.float32))
+
+    ti = tf.random.uniform((), minval=-1, maxval=1, dtype=tf.float32) * max_displacement
+    tj = tf.random.uniform((), minval=-1, maxval=1, dtype=tf.float32) * max_displacement
+    tk = tf.random.uniform((), minval=-1, maxval=1, dtype=tf.float32) * max_displacement
+
+    M = build_affine_trf(tf.convert_to_tensor(FINE_GRID.shape, tf.float32), alpha, beta, gamma, ti, tj, tk)
+    if points.shape.as_list()[-1] == 3:
+        points = tf.transpose(points)
+    new_pts = tf.matmul(M[:3, :3], points)
+    new_pts = tf.expand_dims(M[:3, -1], -1) + new_pts
+    return tf.transpose(new_pts), M  # Remove the last row of ones
+
+
+def build_affine_trf(img_size, alpha, beta, gamma, ti, tj, tk):
+    img_centre = tf.expand_dims(tf.divide(img_size, 2.), -1)
+
+    # Rotation matrix around the image centre
+    # R* = T(p) R(ang) T(-p)
+    # tf.cos and tf.sin expect radians
+    zero = tf.zeros((1,))
+    one = tf.ones((1,))
+
+    T = tf.convert_to_tensor([[one, zero, zero, ti],
+                              [zero, one, zero, tj],
+                              [zero, zero, one, tk],
+                              [zero, zero, zero, one]], tf.float32)
+    T = tf.squeeze(T)
+
+    R = tf.convert_to_tensor([[tf.math.cos(gamma) * tf.math.cos(beta),
+                               tf.math.cos(gamma) * tf.math.sin(beta) * tf.math.sin(alpha) - tf.math.sin(gamma) * tf.math.cos(alpha),
+                               tf.math.cos(gamma) * tf.math.sin(beta) * tf.math.cos(alpha) + tf.math.sin(gamma) * tf.math.sin(alpha),
+                               zero],
+                              [tf.math.sin(gamma) * tf.math.cos(beta),
+                               tf.math.sin(gamma) * tf.math.sin(beta) * tf.math.sin(gamma) + tf.math.cos(gamma) * tf.math.cos(alpha),
+                               tf.math.sin(gamma) * tf.math.sin(beta) * tf.math.cos(gamma) - tf.math.cos(gamma) * tf.math.sin(gamma),
+                               zero],
+                              [-tf.math.sin(beta),
+                               tf.math.cos(beta) * tf.math.sin(alpha),
+                               tf.math.cos(beta) * tf.math.cos(alpha),
+                               zero],
+                              [zero, zero, zero, one]], tf.float32)
+
+    R = tf.squeeze(R)
+
+    Tc = tf.convert_to_tensor([[one, zero, zero, img_centre[0]],
+                               [zero, one, zero, img_centre[1]],
+                               [zero, zero, one, img_centre[2]],
+                               [zero, zero, zero, one]], tf.float32)
+    Tc = tf.squeeze(Tc)
+    Tc_ = tf.convert_to_tensor([[one, zero, zero, -img_centre[0]],
+                                [zero, one, zero, -img_centre[1]],
+                                [zero, zero, one, -img_centre[2]],
+                                [zero, zero, zero, one]], tf.float32)
+    Tc_ = tf.squeeze(Tc_)
+
+    return tf.matmul(T, tf.matmul(Tc, tf.matmul(R, Tc_)))

DeepDeformationMapRegistration/layers/__init__.py

@@ -0,0 +1,5 @@
+from .augmentation import AugmentationLayer
+from .uncertainty_weighting import UncertaintyWeighting, UncertaintyWeightingWithRollingAverage
+from .upsampling import UpSampling1D, UpSampling2D, UpSampling3D
+from .depthwise_conv_3d import DepthwiseConv3D
+from .b_splines import interpolate_spline

DeepDeformationMapRegistration/layers/augmentation.py

@@ -0,0 +1,326 @@
+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 as tf
+from tensorflow.python.framework.errors import InvalidArgumentError
+
+from DeepDeformationMapRegistration.utils.operators import soft_threshold, gaussian_kernel, sample_unique
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.thin_plate_splines import ThinPlateSplines
+from voxelmorph.tf.layers import SpatialTransformer
+
+
+class AugmentationLayer(kl.Layer):
+    def __init__(self,
+                 max_deformation,
+                 max_displacement,
+                 max_rotation,
+                 num_control_points,
+                 in_img_shape,
+                 out_img_shape,
+                 num_augmentations=1,
+                 gamma_augmentation=True,
+                 brightness_augmentation=True,
+                 only_image=False,
+                 only_resize=True,
+                 return_displacement_map=False,
+                 **kwargs):
+        super(AugmentationLayer, self).__init__(**kwargs)
+
+        self.max_deformation = max_deformation
+        self.max_displacement = max_displacement
+        self.max_rotation = max_rotation
+        self.num_control_points = num_control_points
+        self.num_augmentations = num_augmentations
+        self.in_img_shape = in_img_shape
+        self.out_img_shape = out_img_shape
+        self.only_image = only_image
+        self.return_disp_map = return_displacement_map
+
+        self.do_gamma_augm = gamma_augmentation
+        self.do_brightness_augm = brightness_augmentation
+
+        grid = C.CoordinatesGrid()
+        grid.set_coords_grid(in_img_shape, [C.TPS_NUM_CTRL_PTS_PER_AXIS] * 3)
+        self.control_grid = tf.identity(grid.grid_flat(), name='control_grid')
+        self.target_grid = tf.identity(grid.grid_flat(), name='target_grid')
+
+        grid.set_coords_grid(in_img_shape, in_img_shape)
+        self.fine_grid = tf.identity(grid.grid_flat(), 'fine_grid')
+
+        if out_img_shape is not None:
+            self.downsample_factor = [i // o for o, i in zip(out_img_shape, in_img_shape)]
+            self.img_gauss_filter = gaussian_kernel(3, 0.001, 1, 1, 3)
+            # self.resize_transf = tf.diag([*self.downsample_factor, 1])[:-1, :]
+            # self.resize_transf = tf.expand_dims(tf.reshape(self.resize_transf, [-1]), 0, name='resize_transformation')  # ST expects a (12,) vector
+
+        self.augment = not only_resize
+
+    def compute_output_shape(self, input_shape):
+        input_shape = tf.TensorShape(input_shape).as_list()
+        img_shape = (input_shape[0], *self.out_img_shape, 1)
+        seg_shape = (input_shape[0], *self.out_img_shape, input_shape[-1] - 1)
+        disp_shape = (input_shape[0], *self.out_img_shape, 3)
+        # Expect the input to have the image and segmentations in the same tensor
+        if self.return_disp_map:
+            return (img_shape, img_shape, seg_shape, seg_shape, disp_shape)
+        else:
+            return (img_shape, img_shape, seg_shape, seg_shape)
+
+    #@tf.custom_gradient
+    def call(self, in_data, training=None):
+        # def custom_grad(in_grad):
+        #     return tf.ones_like(in_grad)
+        if training is not None:
+            self.augment = training
+        return self.build_batch(in_data)# , custom_grad
+
+    def build_batch(self, fix_data: tf.Tensor):
+        if len(fix_data.get_shape().as_list()) < 5:
+            fix_data = tf.expand_dims(fix_data, axis=0)  # Add Batch dimension
+        # fix_data = tf.tile(fix_data, (self.num_augmentations, *(1,)*4))
+        fix_img_batch, mov_img_batch, fix_seg_batch, mov_seg_batch, disp_map = tf.map_fn(lambda x: self.augment_sample(x),
+                                                                                         fix_data,
+                                                                                         dtype=(tf.float32, tf.float32, tf.float32, tf.float32, tf.float32))
+        # map_fn unstacks elems on axis 0
+        if self.return_disp_map:
+            return fix_img_batch, mov_img_batch, fix_seg_batch, mov_seg_batch, disp_map
+        else:
+            return fix_img_batch, mov_img_batch, fix_seg_batch, mov_seg_batch
+
+    def augment_sample(self, fix_data: tf.Tensor):
+        if self.only_image or not self.augment:
+            fix_img = fix_data
+            fix_segm = tf.zeros_like(fix_data, dtype=tf.float32)
+        else:
+            fix_img = fix_data[..., 0]
+            fix_img = tf.expand_dims(fix_img, -1)
+            fix_segm = fix_data[..., 1:]        # We expect several segmentation masks
+
+        if self.augment:
+            # If we are training, do the full-fledged augmentation
+            fix_img = self.min_max_normalization(fix_img)
+
+            mov_img, mov_segm, disp_map = self.deform_image(tf.squeeze(fix_img), fix_segm)
+            mov_img = tf.expand_dims(mov_img, -1)  # Add the removed channel axis
+
+        # Resample to output_shape
+            if self.out_img_shape is not None:
+                fix_img = self.downsize_image(fix_img)
+                mov_img = self.downsize_image(mov_img)
+
+                fix_segm = self.downsize_segmentation(fix_segm)
+                mov_segm = self.downsize_segmentation(mov_segm)
+
+                disp_map = self.downsize_displacement_map(disp_map)
+
+            if self.do_gamma_augm:
+                fix_img = self.gamma_augmentation(fix_img)
+                mov_img = self.gamma_augmentation(mov_img)
+
+            if self.do_brightness_augm:
+                fix_img = self.brightness_augmentation(fix_img)
+                mov_img = self.brightness_augmentation(mov_img)
+
+        else:
+            # During inference, just resize the input images
+            mov_img = tf.zeros_like(fix_img)
+            mov_segm = tf.zeros_like(fix_segm)
+
+            disp_map = tf.tile(tf.zeros_like(fix_img), [1, 1, 1, 1, 3])
+
+            if self.out_img_shape is not None:
+                fix_img = self.downsize_image(fix_img)
+                mov_img = self.downsize_image(mov_img)
+
+                fix_segm = self.downsize_segmentation(fix_segm)
+                mov_segm = self.downsize_segmentation(mov_segm)
+
+                disp_map = self.downsize_displacement_map(disp_map)
+
+        fix_img = self.min_max_normalization(fix_img)
+        mov_img = self.min_max_normalization(mov_img)
+        return fix_img, mov_img, fix_segm, mov_segm, disp_map
+
+    def downsize_image(self, img):
+        img = tf.expand_dims(img, axis=0)
+        # The filter is symmetrical along the three axes, hence there is no need for transposing the H and D dims
+        img = tf.nn.conv3d(img, self.img_gauss_filter, strides=[1, ] * 5, padding='SAME', data_format='NDHWC')
+        img = tf.layers.MaxPooling3D([1]*3, self.downsample_factor, padding='valid', data_format='channels_last')(img)
+
+        return tf.squeeze(img, axis=0)
+
+    def downsize_segmentation(self, segm):
+        segm = tf.expand_dims(segm, axis=0)
+        segm = tf.layers.MaxPooling3D([1]*3, self.downsample_factor, padding='valid', data_format='channels_last')(segm)
+
+        segm = tf.cast(segm, tf.float32)
+        return tf.squeeze(segm, axis=0)
+
+    def downsize_displacement_map(self, disp_map):
+        disp_map = tf.expand_dims(disp_map, axis=0)
+        # The filter is symmetrical along the three axes, hence there is no need for transposing the H and D dims
+        disp_map = tf.layers.AveragePooling3D([1]*3, self.downsample_factor, padding='valid', data_format='channels_last')(disp_map)
+
+        # self.downsample_factor = in_shape / out_shape, but here we need out_shape / in_shape. Hence, 1 / factor
+        if self.downsample_factor[0] != self.downsample_factor[1] != self.downsample_factor[2]:
+            # Downsize the displacement magnitude along the different axes
+            disp_map_x = disp_map[..., 0] * 1 / self.downsample_factor[0]
+            disp_map_y = disp_map[..., 1] * 1 / self.downsample_factor[1]
+            disp_map_z = disp_map[..., 2] * 1 / self.downsample_factor[2]
+
+            disp_map = tf.stack([disp_map_x, disp_map_y, disp_map_z], axis=-1)
+        else:
+            disp_map = disp_map * 1 / self.downsample_factor[0]
+
+        return tf.squeeze(disp_map, axis=0)
+
+    def gamma_augmentation(self, in_img: tf.Tensor):
+        in_img += 1e-5  # To prvent NaNs
+        gamma = tf.random.uniform((), 0.5, 2, tf.float32)
+
+        return tf.clip_by_value(tf.pow(in_img, gamma), 0, 1)
+
+    def brightness_augmentation(self, in_img: tf.Tensor):
+        c = tf.random.uniform((), 0.5, 2, tf.float32)
+        return tf.clip_by_value(c*in_img, 0, 1)
+
+    def min_max_normalization(self, in_img: tf.Tensor):
+        return tf.div(tf.subtract(in_img, tf.reduce_min(in_img)),
+                      tf.subtract(tf.reduce_max(in_img), tf.reduce_min(in_img)))
+
+    def deform_image(self, fix_img: tf.Tensor, fix_segm: tf.Tensor):
+        # Get locations where the intensity > 0.0
+        idx_points_in_label = tf.where(tf.greater(fix_img, 0.0))
+
+        # Randomly select N points
+        # random_idx = tf.random.uniform((self.num_control_points,),
+        #                                          minval=0, maxval=tf.shape(idx_points_in_label)[0],
+        #                                          dtype=tf.int32)
+        #
+        # disp_location = tf.gather(idx_points_in_label, random_idx)  # And get the coordinates
+        # disp_location = tf.cast(disp_location, tf.float32)
+        disp_location = sample_unique(idx_points_in_label, self.num_control_points, tf.float32)
+
+        # Get the coordinates of the control point displaces
+        rand_disp = tf.random.uniform((self.num_control_points, 3), minval=-1, maxval=1, dtype=tf.float32) * self.max_deformation
+        warped_location = disp_location + rand_disp
+
+        # Add the selected locations to the control grid and the warped locations to the target grid
+        control_grid = tf.concat([self.control_grid, disp_location], axis=0)
+        trg_grid = tf.concat([self.control_grid, warped_location], axis=0)
+
+        # Apply global transformation
+        valid_trf = False
+        while not valid_trf:
+            trg_grid, aff = self.global_transformation(trg_grid)
+
+            # Interpolate the displacement map
+            try:
+                tps = ThinPlateSplines(control_grid, trg_grid)
+                def_grid = tps.interpolate(self.fine_grid)
+            except InvalidArgumentError as err:
+                # If the transformation raises a non-invertible error,
+                # try again until we get a valid transformation
+                continue
+            else:
+                valid_trf = True
+
+        disp_map = self.fine_grid - def_grid
+        disp_map = tf.reshape(disp_map, (*self.in_img_shape, -1))
+
+        # Apply the displacement map
+        fix_img = tf.expand_dims(tf.expand_dims(fix_img, -1), 0)
+        fix_segm = tf.expand_dims(fix_segm, 0)
+        disp_map = tf.cast(tf.expand_dims(disp_map, 0), tf.float32)
+
+        mov_img = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([fix_img, disp_map])
+        mov_segm = SpatialTransformer(interp_method='nearest', indexing='ij', single_transform=False)([fix_segm, disp_map])
+
+        mov_img = tf.where(tf.is_nan(mov_img), tf.zeros_like(mov_img), mov_img)
+        mov_img = tf.where(tf.is_inf(mov_img), tf.zeros_like(mov_img), mov_img)
+
+        mov_segm = tf.where(tf.is_nan(mov_segm), tf.zeros_like(mov_segm), mov_segm)
+        mov_segm = tf.where(tf.is_inf(mov_segm), tf.zeros_like(mov_segm), mov_segm)
+
+        return tf.squeeze(mov_img), tf.squeeze(mov_segm, axis=0), tf.squeeze(disp_map, axis=0)
+
+    def global_transformation(self, points: tf.Tensor):
+        axis = tf.random.uniform((), 0, 3)
+
+        alpha = C.DEG_TO_RAD * tf.cond(tf.logical_and(tf.greater(axis, 0.), tf.less_equal(axis, 1.)),
+                        lambda: tf.random.uniform((), -self.max_rotation, self.max_rotation),
+                        lambda: tf.zeros((), tf.float32))
+        beta = C.DEG_TO_RAD * tf.cond(tf.logical_and(tf.greater(axis, 1.), tf.less_equal(axis, 2.)),
+                       lambda: tf.random.uniform((), -self.max_rotation, self.max_rotation),
+                       lambda: tf.zeros((), tf.float32))
+        gamma = C.DEG_TO_RAD * tf.cond(tf.logical_and(tf.greater(axis, 2.), tf.less_equal(axis, 3.)),
+                        lambda: tf.random.uniform((), -self.max_rotation, self.max_rotation),
+                        lambda: tf.zeros((), tf.float32))
+
+        ti = tf.random.uniform((), minval=-1, maxval=1, dtype=tf.float32) * self.max_displacement
+        tj = tf.random.uniform((), minval=-1, maxval=1, dtype=tf.float32) * self.max_displacement
+        tk = tf.random.uniform((), minval=-1, maxval=1, dtype=tf.float32) * self.max_displacement
+
+        M = self.build_affine_transformation(tf.convert_to_tensor(self.in_img_shape, tf.float32),
+                                             alpha, beta, gamma, ti, tj, tk)
+
+        points = tf.transpose(points)
+        new_pts = tf.matmul(M[:3, :3], points)
+        new_pts = tf.expand_dims(M[:3, -1], -1) + new_pts
+        return tf.transpose(new_pts), M
+
+    @staticmethod
+    def build_affine_transformation(img_shape, alpha, beta, gamma, ti, tj, tk):
+        img_centre = tf.divide(img_shape, 2.)
+
+        # Rotation matrix around the image centre
+        # R* = T(p) R(ang) T(-p)
+        # tf.cos and tf.sin expect radians
+
+        T = tf.convert_to_tensor([[1, 0, 0, ti],
+                                  [0, 1, 0, tj],
+                                  [0, 0, 1, tk],
+                                  [0, 0, 0, 1]], tf.float32)
+
+        Ri = tf.convert_to_tensor([[1, 0, 0, 0],
+                                   [0, tf.math.cos(alpha), -tf.math.sin(alpha), 0],
+                                   [0, tf.math.sin(alpha),  tf.math.cos(alpha), 0],
+                                   [0, 0, 0, 1]], tf.float32)
+
+        Rj = tf.convert_to_tensor([[ tf.math.cos(beta), 0, tf.math.sin(beta), 0],
+                                   [0, 1, 0, 0],
+                                   [-tf.math.sin(beta), 0, tf.math.cos(beta), 0],
+                                   [0, 0, 0, 1]], tf.float32)
+
+        Rk = tf.convert_to_tensor([[tf.math.cos(gamma), -tf.math.sin(gamma), 0, 0],
+                                   [tf.math.sin(gamma),  tf.math.cos(gamma), 0, 0],
+                                   [0, 0, 1, 0],
+                                   [0, 0, 0, 1]], tf.float32)
+
+        R = tf.matmul(tf.matmul(Ri, Rj), Rk)
+
+        Tc = tf.convert_to_tensor([[1, 0, 0, img_centre[0]],
+                                   [0, 1, 0, img_centre[1]],
+                                   [0, 0, 1, img_centre[2]],
+                                   [0, 0, 0, 1]], tf.float32)
+
+        Tc_ = tf.convert_to_tensor([[1, 0, 0, -img_centre[0]],
+                                    [0, 1, 0, -img_centre[1]],
+                                    [0, 0, 1, -img_centre[2]],
+                                    [0, 0, 0, 1]], tf.float32)
+
+        return tf.matmul(T, tf.matmul(Tc, tf.matmul(R, Tc_)))
+
+    def get_config(self):
+        config = super(AugmentationLayer, self).get_config()
+        return config
+
+
+

DeepDeformationMapRegistration/layers/b_splines.py

@@ -0,0 +1,320 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Polyharmonic spline interpolation."""
+
+import tensorflow as tf
+from typing import Union, List
+import numpy as np
+
+Number = Union[
+    float,
+    int,
+    np.float16,
+    np.float32,
+    np.float64,
+    np.int8,
+    np.int16,
+    np.int32,
+    np.int64,
+    np.uint8,
+    np.uint16,
+    np.uint32,
+    np.uint64,
+]
+
+TensorLike = Union[
+    List[Union[Number, list]],
+    tuple,
+    Number,
+    np.ndarray,
+    tf.Tensor,
+    tf.SparseTensor,
+    tf.Variable,
+]
+FloatTensorLike = Union[tf.Tensor, float, np.float16, np.float32, np.float64]
+
+EPSILON = 0.0000000001
+
+
+def _cross_squared_distance_matrix(x: TensorLike, y: TensorLike) -> tf.Tensor:
+    """Pairwise squared distance between two (batch) matrices' rows (2nd dim).
+    Computes the pairwise distances between rows of x and rows of y.
+    Args:
+      x: `[batch_size, n, d]` float `Tensor`.
+      y: `[batch_size, m, d]` float `Tensor`.
+    Returns:
+      squared_dists: `[batch_size, n, m]` float `Tensor`, where
+      `squared_dists[b,i,j] = ||x[b,i,:] - y[b,j,:]||^2`.
+    """
+    x_norm_squared = tf.reduce_sum(tf.square(x), 2)
+    y_norm_squared = tf.reduce_sum(tf.square(y), 2)
+
+    # Expand so that we can broadcast.
+    x_norm_squared_tile = tf.expand_dims(x_norm_squared, 2)
+    y_norm_squared_tile = tf.expand_dims(y_norm_squared, 1)
+
+    x_y_transpose = tf.matmul(x, y, adjoint_b=True)
+
+    # squared_dists[b,i,j] = ||x_bi - y_bj||^2 =
+    # x_bi'x_bi- 2x_bi'x_bj + x_bj'x_bj
+    squared_dists = x_norm_squared_tile - 2 * x_y_transpose + y_norm_squared_tile
+
+    return squared_dists
+
+
+def _pairwise_squared_distance_matrix(x: TensorLike) -> tf.Tensor:
+    """Pairwise squared distance among a (batch) matrix's rows (2nd dim).
+    This saves a bit of computation vs. using
+    `_cross_squared_distance_matrix(x, x)`
+    Args:
+      x: `[batch_size, n, d]` float `Tensor`.
+    Returns:
+      squared_dists: `[batch_size, n, n]` float `Tensor`, where
+      `squared_dists[b,i,j] = ||x[b,i,:] - x[b,j,:]||^2`.
+    """
+
+    x_x_transpose = tf.matmul(x, x, adjoint_b=True)
+    x_norm_squared = tf.linalg.diag_part(x_x_transpose)
+    x_norm_squared_tile = tf.expand_dims(x_norm_squared, 2)
+
+    # squared_dists[b,i,j] = ||x_bi - x_bj||^2 =
+    # = x_bi'x_bi- 2x_bi'x_bj + x_bj'x_bj
+    squared_dists = (
+            x_norm_squared_tile
+            - 2 * x_x_transpose
+            + tf.transpose(x_norm_squared_tile, [0, 2, 1])
+    )
+
+    return squared_dists
+
+
+def _solve_interpolation(
+        train_points: TensorLike,
+        train_values: TensorLike,
+        order: int,
+        regularization_weight: FloatTensorLike,
+) -> TensorLike:
+    r"""Solve for interpolation coefficients.
+    Computes the coefficients of the polyharmonic interpolant for the
+    'training' data defined by `(train_points, train_values)` using the kernel
+    $\phi$.
+    Args:
+      train_points: `[b, n, d]` interpolation centers.
+      train_values: `[b, n, k]` function values.
+      order: order of the interpolation.
+      regularization_weight: weight to place on smoothness regularization term.
+    Returns:
+      w: `[b, n, k]` weights on each interpolation center
+      v: `[b, d, k]` weights on each input dimension
+    Raises:
+      ValueError: if d or k is not fully specified.
+    """
+
+    # These dimensions are set dynamically at runtime.
+    b, n, _ = tf.unstack(tf.shape(train_points), num=3)
+
+    d = train_points.shape[-1]
+    if d is None:
+        raise ValueError(
+            "The dimensionality of the input points (d) must be "
+            "statically-inferrable."
+        )
+
+    k = train_values.shape[-1]
+    if k is None:
+        raise ValueError(
+            "The dimensionality of the output values (k) must be "
+            "statically-inferrable."
+        )
+
+    # First, rename variables so that the notation (c, f, w, v, A, B, etc.)
+    # follows https://en.wikipedia.org/wiki/Polyharmonic_spline.
+    # To account for python style guidelines we use
+    # matrix_a for A and matrix_b for B.
+
+    c = train_points
+    f = train_values
+
+    # Next, construct the linear system.
+    with tf.name_scope("construct_linear_system"):
+
+        matrix_a = _phi(_pairwise_squared_distance_matrix(c), order)  # [b, n, n]
+        if regularization_weight > 0:
+            batch_identity_matrix = tf.expand_dims(tf.eye(n, dtype=c.dtype), 0)
+            matrix_a += regularization_weight * batch_identity_matrix
+
+        # Append ones to the feature values for the bias term
+        # in the linear model.
+        ones = tf.ones_like(c[..., :1], dtype=c.dtype)
+        matrix_b = tf.concat([c, ones], 2)  # [b, n, d + 1]
+
+        # [b, n + d + 1, n]
+        left_block = tf.concat([matrix_a, tf.transpose(matrix_b, [0, 2, 1])], 1)
+
+        num_b_cols = matrix_b.get_shape()[2]  # d + 1
+        lhs_zeros = tf.zeros([b, num_b_cols, num_b_cols], train_points.dtype)
+        right_block = tf.concat([matrix_b, lhs_zeros], 1)  # [b, n + d + 1, d + 1]
+        lhs = tf.concat([left_block, right_block], 2)  # [b, n + d + 1, n + d + 1]
+
+        rhs_zeros = tf.zeros([b, d + 1, k], train_points.dtype)
+        rhs = tf.concat([f, rhs_zeros], 1)  # [b, n + d + 1, k]
+
+    # Then, solve the linear system and unpack the results.
+    with tf.name_scope("solve_linear_system"):
+        w_v = tf.linalg.solve(lhs, rhs)
+        w = w_v[:, :n, :]
+        v = w_v[:, n:, :]
+
+    return w, v
+
+
+def _apply_interpolation(
+        query_points: TensorLike,
+        train_points: TensorLike,
+        w: TensorLike,
+        v: TensorLike,
+        order: int,
+) -> TensorLike:
+    """Apply polyharmonic interpolation model to data.
+    Given coefficients w and v for the interpolation model, we evaluate
+    interpolated function values at query_points.
+    Args:
+      query_points: `[b, m, d]` x values to evaluate the interpolation at.
+      train_points: `[b, n, d]` x values that act as the interpolation centers
+          (the c variables in the wikipedia article).
+      w: `[b, n, k]` weights on each interpolation center.
+      v: `[b, d, k]` weights on each input dimension.
+      order: order of the interpolation.
+    Returns:
+      Polyharmonic interpolation evaluated at points defined in `query_points`.
+    """
+
+    # First, compute the contribution from the rbf term.
+    pairwise_dists = _cross_squared_distance_matrix(query_points, train_points)
+    phi_pairwise_dists = _phi(pairwise_dists, order)
+
+    rbf_term = tf.matmul(phi_pairwise_dists, w)
+
+    # Then, compute the contribution from the linear term.
+    # Pad query_points with ones, for the bias term in the linear model.
+    query_points_pad = tf.concat(
+        [query_points, tf.ones_like(query_points[..., :1], train_points.dtype)], 2
+    )
+    linear_term = tf.matmul(query_points_pad, v)
+
+    return rbf_term + linear_term
+
+
+def _phi(r: FloatTensorLike, order: int) -> FloatTensorLike:
+    """Coordinate-wise nonlinearity used to define the order of the
+    interpolation.
+    See https://en.wikipedia.org/wiki/Polyharmonic_spline for the definition.
+    Args:
+      r: input op.
+      order: interpolation order.
+    Returns:
+      `phi_k` evaluated coordinate-wise on `r`, for `k = r`.
+    """
+
+    # using EPSILON prevents log(0), sqrt0), etc.
+    # sqrt(0) is well-defined, but its gradient is not
+    with tf.name_scope("phi"):
+        if order == 1:
+            r = tf.maximum(r, EPSILON)
+            r = tf.sqrt(r)
+            return r
+        elif order == 2:
+            return 0.5 * r * tf.math.log(tf.maximum(r, EPSILON))
+        elif order == 4:
+            return 0.5 * tf.square(r) * tf.math.log(tf.maximum(r, EPSILON))
+        elif order % 2 == 0:
+            r = tf.maximum(r, EPSILON)
+            return 0.5 * tf.pow(r, 0.5 * order) * tf.math.log(r)
+        else:
+            r = tf.maximum(r, EPSILON)
+            return tf.pow(r, 0.5 * order)
+
+
+def interpolate_spline(
+        train_points: TensorLike,
+        train_values: TensorLike,
+        query_points: TensorLike,
+        order: int,
+        regularization_weight: FloatTensorLike = 0.0,
+        name: str = "interpolate_spline",
+) -> tf.Tensor:
+    r"""Interpolate signal using polyharmonic interpolation.
+    The interpolant has the form
+    $$f(x) = \sum_{i = 1}^n w_i \phi(||x - c_i||) + v^T x + b.$$
+    This is a sum of two terms: (1) a weighted sum of radial basis function
+    (RBF) terms, with the centers \\(c_1, ... c_n\\), and (2) a linear term
+    with a bias. The \\(c_i\\) vectors are 'training' points.
+    In the code, b is absorbed into v
+    by appending 1 as a final dimension to x. The coefficients w and v are
+    estimated such that the interpolant exactly fits the value of the function
+    at the \\(c_i\\) points, the vector w is orthogonal to each \\(c_i\\),
+    and the vector w sums to 0. With these constraints, the coefficients
+    can be obtained by solving a linear system.
+    \\(\phi\\) is an RBF, parametrized by an interpolation
+    order. Using order=2 produces the well-known thin-plate spline.
+    We also provide the option to perform regularized interpolation. Here, the
+    interpolant is selected to trade off between the squared loss on the
+    training data and a certain measure of its curvature
+    ([details](https://en.wikipedia.org/wiki/Polyharmonic_spline)).
+    Using a regularization weight greater than zero has the effect that the
+    interpolant will no longer exactly fit the training data. However, it may
+    be less vulnerable to overfitting, particularly for high-order
+    interpolation.
+    Note the interpolation procedure is differentiable with respect to all
+    inputs besides the order parameter.
+    We support dynamically-shaped inputs, where batch_size, n, and m are None
+    at graph construction time. However, d and k must be known.
+    Args:
+      train_points: `[batch_size, n, d]` float `Tensor` of n d-dimensional
+        locations. These do not need to be regularly-spaced.
+      train_values: `[batch_size, n, k]` float `Tensor` of n c-dimensional
+        values evaluated at train_points.
+      query_points: `[batch_size, m, d]` `Tensor` of m d-dimensional locations
+        where we will output the interpolant's values.
+      order: order of the interpolation. Common values are 1 for
+        \\(\phi(r) = r\\), 2 for \\(\phi(r) = r^2 * log(r)\\)
+        (thin-plate spline), or 3 for \\(\phi(r) = r^3\\).
+      regularization_weight: weight placed on the regularization term.
+        This will depend substantially on the problem, and it should always be
+        tuned. For many problems, it is reasonable to use no regularization.
+        If using a non-zero value, we recommend a small value like 0.001.
+      name: name prefix for ops created by this function
+    Returns:
+      `[b, m, k]` float `Tensor` of query values. We use train_points and
+      train_values to perform polyharmonic interpolation. The query values are
+      the values of the interpolant evaluated at the locations specified in
+      query_points.
+    """
+    with tf.name_scope(name or "interpolate_spline"):
+        train_points = tf.convert_to_tensor(train_points)
+        train_values = tf.convert_to_tensor(train_values)
+        query_points = tf.convert_to_tensor(query_points)
+
+        # First, fit the spline to the observed data.
+        with tf.name_scope("solve"):
+            w, v = _solve_interpolation(
+                train_points, train_values, order, regularization_weight
+            )
+
+        # Then, evaluate the spline at the query locations.
+        with tf.name_scope("predict"):
+            query_values = _apply_interpolation(query_points, train_points, w, v, order)
+
+    return query_values

DeepDeformationMapRegistration/layers/depthwise_conv_3d.py

@@ -0,0 +1,264 @@
+# SRC: https://github.com/alexandrosstergiou/keras-DepthwiseConv3D
+
+'''
+This is a modification of the SeparableConv3D code in Keras,
+to perform just the Depthwise Convolution (1st step) of the
+Depthwise Separable Convolution layer.
+'''
+from __future__ import absolute_import
+
+from tensorflow.keras import backend as K
+from tensorflow.keras import initializers
+from tensorflow.keras import regularizers
+from tensorflow.keras import constraints
+import tensorflow.keras.utils as conv_utils
+from tensorflow.keras.layers import Conv3D, InputSpec
+from tensorflow.python.keras.backend import _preprocess_padding, _preprocess_conv3d_input
+
+import tensorflow as tf
+
+
+class DepthwiseConv3D(Conv3D):
+    """Depthwise 3D convolution.
+    Depth-wise part of separable convolutions consist in performing
+    just the first step/operation
+    (which acts on each input channel separately).
+    It does not perform the pointwise convolution (second step).
+    The `depth_multiplier` argument controls how many
+    output channels are generated per input channel in the depthwise step.
+    # Arguments
+        kernel_size: An integer or tuple/list of 3 integers, specifying the
+            depth, width and height of the 3D convolution window.
+            Can be a single integer to specify the same value for
+            all spatial dimensions.
+        strides: An integer or tuple/list of 3 integers,
+            specifying the strides of the convolution along the depth, width and height.
+            Can be a single integer to specify the same value for
+            all spatial dimensions.
+        padding: one of `"valid"` or `"same"` (case-insensitive).
+        depth_multiplier: The number of depthwise convolution output channels
+            for each input channel.
+            The total number of depthwise convolution output
+            channels will be equal to `filterss_in * depth_multiplier`.
+        groups: The depth size of the convolution (as a variant of the original Depthwise conv)
+        data_format: A string,
+            one of `channels_last` (default) or `channels_first`.
+            The ordering of the dimensions in the inputs.
+            `channels_last` corresponds to inputs with shape
+            `(batch, height, width, channels)` while `channels_first`
+            corresponds to inputs with shape
+            `(batch, channels, height, width)`.
+            It defaults to the `image_data_format` value found in your
+            Keras config file at `~/.keras/keras.json`.
+            If you never set it, then it will be "channels_last".
+        activation: Activation function to use
+            (see [activations](../activations.md)).
+            If you don't specify anything, no activation is applied
+            (ie. "linear" activation: `a(x) = x`).
+        use_bias: Boolean, whether the layer uses a bias vector.
+        depthwise_initializer: Initializer for the depthwise kernel matrix
+            (see [initializers](../initializers.md)).
+        bias_initializer: Initializer for the bias vector
+            (see [initializers](../initializers.md)).
+        depthwise_regularizer: Regularizer function applied to
+            the depthwise kernel matrix
+            (see [regularizer](../regularizers.md)).
+        bias_regularizer: Regularizer function applied to the bias vector
+            (see [regularizer](../regularizers.md)).
+        dialation_rate: List of ints.
+                        Defines the dilation factor for each dimension in the
+                        input. Defaults to (1,1,1)
+        activity_regularizer: Regularizer function applied to
+            the output of the layer (its "activation").
+            (see [regularizer](../regularizers.md)).
+        depthwise_constraint: Constraint function applied to
+            the depthwise kernel matrix
+            (see [constraints](../constraints.md)).
+        bias_constraint: Constraint function applied to the bias vector
+            (see [constraints](../constraints.md)).
+    # Input shape
+        5D tensor with shape:
+        `(batch, depth, channels, rows, cols)` if data_format='channels_first'
+        or 5D tensor with shape:
+        `(batch, depth, rows, cols, channels)` if data_format='channels_last'.
+    # Output shape
+        5D tensor with shape:
+        `(batch, filters * depth, new_depth, new_rows, new_cols)` if data_format='channels_first'
+        or 4D tensor with shape:
+        `(batch, new_depth, new_rows, new_cols, filters * depth)` if data_format='channels_last'.
+        `rows` and `cols` values might have changed due to padding.
+    """
+
+    #@legacy_depthwise_conv3d_support
+    def __init__(self,
+                 kernel_size,
+                 strides=(1, 1, 1),
+                 padding='valid',
+                 depth_multiplier=1,
+                 groups=None,
+                 data_format=None,
+                 activation=None,
+                 use_bias=True,
+                 depthwise_initializer='glorot_uniform',
+                 bias_initializer='zeros',
+                 dilation_rate = (1, 1, 1),
+                 depthwise_regularizer=None,
+                 bias_regularizer=None,
+                 activity_regularizer=None,
+                 depthwise_constraint=None,
+                 bias_constraint=None,
+                 **kwargs):
+        super(DepthwiseConv3D, self).__init__(
+            filters=None,
+            kernel_size=kernel_size,
+            strides=strides,
+            padding=padding,
+            data_format=data_format,
+            activation=activation,
+            use_bias=use_bias,
+            bias_regularizer=bias_regularizer,
+            dilation_rate=dilation_rate,
+            activity_regularizer=activity_regularizer,
+            bias_constraint=bias_constraint,
+            **kwargs)
+        self.depth_multiplier = depth_multiplier
+        self.groups = groups
+        self.depthwise_initializer = initializers.get(depthwise_initializer)
+        self.depthwise_regularizer = regularizers.get(depthwise_regularizer)
+        self.depthwise_constraint = constraints.get(depthwise_constraint)
+        self.bias_initializer = initializers.get(bias_initializer)
+        self.dilation_rate = dilation_rate
+        self._padding = _preprocess_padding(self.padding)
+        self._strides = (1,) + self.strides + (1,)
+        self._data_format = "NDHWC"
+        self.input_dim = None
+
+    def build(self, input_shape):
+        if len(input_shape) < 5:
+            raise ValueError('Inputs to `DepthwiseConv3D` should have rank 5. '
+                             'Received input shape:', str(input_shape))
+        if self.data_format == 'channels_first':
+            channel_axis = 1
+        else:
+            channel_axis = -1
+        if input_shape[channel_axis] is None:
+            raise ValueError('The channel dimension of the inputs to '
+                             '`DepthwiseConv3D` '
+                             'should be defined. Found `None`.')
+        self.input_dim = int(input_shape[channel_axis])
+
+        if self.groups is None:
+            self.groups = self.input_dim
+
+        if self.groups > self.input_dim:
+            raise ValueError('The number of groups cannot exceed the number of channels')
+
+        if self.input_dim % self.groups != 0:
+            raise ValueError('Warning! The channels dimension is not divisible by the group size chosen')
+
+        depthwise_kernel_shape = (self.kernel_size[0],
+                                  self.kernel_size[1],
+                                  self.kernel_size[2],
+                                  self.input_dim,
+                                  self.depth_multiplier)
+
+        self.depthwise_kernel = self.add_weight(
+            shape=depthwise_kernel_shape,
+            initializer=self.depthwise_initializer,
+            name='depthwise_kernel',
+            regularizer=self.depthwise_regularizer,
+            constraint=self.depthwise_constraint)
+
+        if self.use_bias:
+            self.bias = self.add_weight(shape=(self.groups * self.depth_multiplier,),
+                                        initializer=self.bias_initializer,
+                                        name='bias',
+                                        regularizer=self.bias_regularizer,
+                                        constraint=self.bias_constraint)
+        else:
+            self.bias = None
+        # Set input spec.
+        self.input_spec = InputSpec(ndim=5, axes={channel_axis: self.input_dim})
+        self.built = True
+
+    def call(self, inputs, training=None):
+        inputs = _preprocess_conv3d_input(inputs, self.data_format)
+
+        if self.data_format == 'channels_last':
+            dilation = (1,) + self.dilation_rate + (1,)
+        else:
+            dilation = self.dilation_rate + (1,) + (1,)
+
+        if self._data_format == 'NCDHW':
+            outputs = tf.concat(
+                [tf.nn.conv3d(inputs[0][:, i:i+self.input_dim//self.groups, :, :, :], self.depthwise_kernel[:, :, :, i:i+self.input_dim//self.groups, :],
+                              strides=self._strides,
+                              padding=self._padding,
+                              dilations=dilation,
+                              data_format=self._data_format) for i in range(0, self.input_dim, self.input_dim//self.groups)], axis=1)
+
+        else:
+            outputs = tf.concat(
+                [tf.nn.conv3d(inputs[0][:, :, :, :, i:i+self.input_dim//self.groups], self.depthwise_kernel[:, :, :, i:i+self.input_dim//self.groups, :],
+                              strides=self._strides,
+                              padding=self._padding,
+                              dilations=dilation,
+                              data_format=self._data_format) for i in range(0, self.input_dim, self.input_dim//self.groups)], axis=-1)
+
+        if self.bias is not None:
+            outputs = K.bias_add(
+                outputs,
+                self.bias,
+                data_format=self.data_format)
+
+        if self.activation is not None:
+            return self.activation(outputs)
+
+        return outputs
+
+    def compute_output_shape(self, input_shape):
+        if self.data_format == 'channels_first':
+            depth = input_shape[2]
+            rows = input_shape[3]
+            cols = input_shape[4]
+            out_filters = self.groups * self.depth_multiplier
+        elif self.data_format == 'channels_last':
+            depth = input_shape[1]
+            rows = input_shape[2]
+            cols = input_shape[3]
+            out_filters = self.groups * self.depth_multiplier
+
+        depth = conv_utils.conv_output_length(depth, self.kernel_size[0],
+                                              self.padding,
+                                              self.strides[0])
+
+        rows = conv_utils.conv_output_length(rows, self.kernel_size[1],
+                                             self.padding,
+                                             self.strides[1])
+
+        cols = conv_utils.conv_output_length(cols, self.kernel_size[2],
+                                             self.padding,
+                                             self.strides[2])
+
+        if self.data_format == 'channels_first':
+            return (input_shape[0], out_filters, depth, rows, cols)
+
+        elif self.data_format == 'channels_last':
+            return (input_shape[0], depth, rows, cols, out_filters)
+
+    def get_config(self):
+        config = super(DepthwiseConv3D, self).get_config()
+        config.pop('filters')
+        config.pop('kernel_initializer')
+        config.pop('kernel_regularizer')
+        config.pop('kernel_constraint')
+        config['depth_multiplier'] = self.depth_multiplier
+        config['depthwise_initializer'] = initializers.serialize(self.depthwise_initializer)
+        config['depthwise_regularizer'] = regularizers.serialize(self.depthwise_regularizer)
+        config['depthwise_constraint'] = constraints.serialize(self.depthwise_constraint)
+        return config
+
+    def __call__(self, inputs, training=True):
+        return self.call(inputs, training)
+
+DepthwiseConvolution3D = DepthwiseConv3D

DeepDeformationMapRegistration/layers/uncertainty_weighting.py

@@ -0,0 +1,347 @@
+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
+import random
+
+from DeepDeformationMapRegistration.utils.operators import soft_threshold, gaussian_kernel, sample_unique
+import DeepDeformationMapRegistration.utils.constants as C
+from DeepDeformationMapRegistration.utils.thin_plate_splines import ThinPlateSplines
+from voxelmorph.tf.layers import SpatialTransformer
+from neurite.tf.utils import resize
+#from cupyx.scipy.ndimage import zoom
+#import cupy
+
+
+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
+
+
+class UncertaintyWeightingWithRollingAverage(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.],
+                 roll_avg_reference=0,  # position in loss_fns of the reference loss function for the rolling avg
+                 **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))
+        # Rolling average attributes
+        self.ref_loss = roll_avg_reference
+        self.compute_roll_avg = False    # Toogle between computing the average of the losses or updating a know average
+        self.scale_factor = [1.] * num_loss_fns
+        self.n = 0  # Number of viewed samples
+        self.temp_storage = [0.] * num_loss_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(UncertaintyWeightingWithRollingAverage, 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(UncertaintyWeightingWithRollingAverage, self).build(input_shape)
+
+    def store_values(self, new_loss_values):
+        for i, (t, v) in enumerate(zip(self.temp_storage, new_loss_values)):
+            self.temp_storage[i] = t + v
+        self.n += 1
+
+    def compute_scale_factors(self):
+        for i, val in enumerate(self.temp_storage):
+            self.scale_factor[i] = self.n / val  # 1/avg
+
+        self.scale_factor[self.ref_loss] = 1.
+
+        self.temp_storage = [0.] * self.num_loss
+        self.n = 0
+
+    @property
+    def ref_on_epoch_end_function(self):
+        return self.compute_scale_factors
+
+    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, sf in zip(ys_true, ys_pred, self.loss_fns, self.manual_loss_w, self.scale_factor):
+            loss_values.append(sf * tf.keras.backend.mean(man_w * loss_fn(y_true, y_pred)))
+            loss_names_loss.append(loss_fn.__name__)
+
+        self.store_values(loss_values)
+        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')
+                sc_tf = tf.convert_to_tensor(self.scale_factor, dtype=tf.float32, name='scale_factors_tf')
+                self.add_metric(tf.slice(sc_tf, [i], [1]), name='SCALE_FACTOR_{}_{}'.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/layers/upsampling.py

@@ -0,0 +1,761 @@
+import tensorflow as tf
+import numpy as np
+from numpy import (zeros, where, diff, floor, minimum, maximum, array, concatenate, logical_or, logical_xor,
+                   sqrt)
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.keras.utils import conv_utils
+from tensorflow.python.keras.engine.base_layer import Layer
+from tensorflow.python.keras.engine.input_spec import InputSpec
+from tensorflow.python.util.tf_export import keras_export # api_export
+# SRC: https://github.com/tensorflow/tensorflow/issues/46609
+# import functools
+
+# keras_export = functools.partial(api_export, 'keras')  # keras_export is not defined in 1.13 but in 1.15 --> https://github.com/tensorflow/tensorflow/blob/3d6e4f24e32b5dbe0a83aaa6e9d0f6671ba41da8/tensorflow/python/util/tf_export.py
+
+def linear_interpolate(x_fix, y_fix, x_var):
+    '''
+        Functionality:
+            1D linear interpolation
+        Author:
+            Michael Osthege
+        Link:
+            https://gist.github.com/michaelosthege/e20d242bc62a434843b586c78ebce6cc
+    '''
+
+    x_repeat = tf.tile(x_var[:, None], (len(x_fix), ))
+    distances = tf.abs(x_repeat - x_fix)
+
+    x_indices = tf.searchsorted(x_fix, x_var)
+
+    weights = tf.zeros_like(distances)
+    idx = tf.arange(len(x_indices))
+    weights[idx, x_indices] = distances[idx, x_indices - 1]
+    weights[idx, x_indices - 1] = distances[idx, x_indices]
+    weights /= np.sum(weights, axis=1)[:, None]
+
+    y_var = np.dot(weights, y_fix.T)
+
+    return y_var
+
+
+def cubic_interpolate(x, y, x0):
+    '''
+        Functionliaty:
+            1D cubic spline interpolation
+        Author:
+            Raphael Valentin
+        Link:
+            https://stackoverflow.com/questions/31543775/how-to-perform-cubic-spline-interpolation-in-python
+    '''
+
+    x = np.asfarray(x)
+    y = np.asfarray(y)
+
+    # remove non finite values
+    # indexes = np.isfinite(x)
+    # x = x[indexes]
+    # y = y[indexes]
+
+    # check if sorted
+    if np.any(np.diff(x) < 0):
+        indexes = np.argsort(x)
+        x = x[indexes]
+        y = y[indexes]
+
+    size = len(x)
+
+    xdiff = np.diff(x)
+    ydiff = np.diff(y)
+
+    # allocate buffer matrices
+    Li = np.empty(size)
+    Li_1 = np.empty(size - 1)
+    z = np.empty(size)
+
+    # fill diagonals Li and Li-1 and solve [L][y] = [B]
+    Li[0] = sqrt(2 * xdiff[0])
+    Li_1[0] = 0.0
+    B0 = 0.0  # natural boundary
+    z[0] = B0 / Li[0]
+
+    for i in range(1, size - 1, 1):
+        Li_1[i] = xdiff[i - 1] / Li[i - 1]
+        Li[i] = sqrt(2 * (xdiff[i - 1] + xdiff[i]) - Li_1[i - 1] * Li_1[i - 1])
+        Bi = 6 * (ydiff[i] / xdiff[i] - ydiff[i - 1] / xdiff[i - 1])
+        z[i] = (Bi - Li_1[i - 1] * z[i - 1]) / Li[i]
+
+    i = size - 1
+    Li_1[i - 1] = xdiff[-1] / Li[i - 1]
+    Li[i] = sqrt(2 * xdiff[-1] - Li_1[i - 1] * Li_1[i - 1])
+    Bi = 0.0  # natural boundary
+    z[i] = (Bi - Li_1[i - 1] * z[i - 1]) / Li[i]
+
+    # solve [L.T][x] = [y]
+    i = size - 1
+    z[i] = z[i] / Li[i]
+    for i in range(size - 2, -1, -1):
+        z[i] = (z[i] - Li_1[i - 1] * z[i + 1]) / Li[i]
+
+    # find index
+    index = x.searchsorted(x0)
+    np.clip(index, 1, size - 1, index)
+
+    xi1, xi0 = x[index], x[index - 1]
+    yi1, yi0 = y[index], y[index - 1]
+    zi1, zi0 = z[index], z[index - 1]
+    hi1 = xi1 - xi0
+
+    # calculate cubic
+    f0 = zi0/(6*hi1)*(xi1-x0)**3 + \
+         zi1/(6*hi1)*(x0-xi0)**3 + \
+         (yi1/hi1 - zi1*hi1/6)*(x0-xi0) + \
+         (yi0/hi1 - zi0*hi1/6)*(xi1-x0)
+
+    return f0
+
+
+def pchip_interpolate(xi, yi, x, mode="mono", verbose=False):
+    '''
+        Functionality:
+            1D PCHP interpolation
+        Authors:
+            Michael Taylor <[email protected]>
+            Mathieu Virbel <[email protected]>
+        Link:
+            https://gist.github.com/tito/553f1135959921ce6699652bf656150d
+    '''
+
+    if mode not in ("mono", "quad"):
+        raise ValueError("Unrecognized mode string")
+
+    # Search for [xi,xi+1] interval for each x
+    xi = xi.astype("double")
+    yi = yi.astype("double")
+
+    x_index = zeros(len(x), dtype="int")
+    xi_steps = diff(xi)
+    if not all(xi_steps > 0):
+        raise ValueError("x-coordinates are not in increasing order.")
+
+    x_steps = diff(x)
+    if xi_steps.max() / xi_steps.min() < 1.000001:
+        # uniform input grid
+        if verbose:
+            print("pchip: uniform input grid")
+        xi_start = xi[0]
+        xi_step = (xi[-1] - xi[0]) / (len(xi) - 1)
+        x_index = minimum(maximum(floor((x - xi_start) / xi_step).astype(int), 0), len(xi) - 2)
+
+        # Calculate gradients d
+        h = (xi[-1] - xi[0]) / (len(xi) - 1)
+        d = zeros(len(xi), dtype="double")
+        if mode == "quad":
+            # quadratic polynomial fit
+            d[[0]] = (yi[1] - yi[0]) / h
+            d[[-1]] = (yi[-1] - yi[-2]) / h
+            d[1:-1] = (yi[2:] - yi[0:-2]) / 2 / h
+        else:
+            # mode=='mono', Fritsch-Carlson algorithm from fortran numerical
+            # recipe
+            delta = diff(yi) / h
+            d = concatenate((delta[0:1], 2 / (1 / delta[0:-1] + 1 / delta[1:]), delta[-1:]))
+            d[concatenate((array([False]), logical_xor(delta[0:-1] > 0, delta[1:] > 0), array([False])))] = 0
+            d[logical_or(concatenate((array([False]), delta == 0)), concatenate(
+                (delta == 0, array([False]))))] = 0
+        # Calculate output values y
+        dxxi = x - xi[x_index]
+        dxxid = x - xi[1 + x_index]
+        dxxi2 = pow(dxxi, 2)
+        dxxid2 = pow(dxxid, 2)
+        y = (2 / pow(h, 3) * (yi[x_index] * dxxid2 * (dxxi + h / 2) - yi[1 + x_index] * dxxi2 *
+                              (dxxid - h / 2)) + 1 / pow(h, 2) *
+             (d[x_index] * dxxid2 * dxxi + d[1 + x_index] * dxxi2 * dxxid))
+    else:
+        # not uniform input grid
+        if (x_steps.max() / x_steps.min() < 1.000001 and x_steps.max() / x_steps.min() > 0.999999):
+            # non-uniform input grid, uniform output grid
+            if verbose:
+                print("pchip: non-uniform input grid, uniform output grid")
+            x_decreasing = x[-1] < x[0]
+            if x_decreasing:
+                x = x[::-1]
+            x_start = x[0]
+            x_step = (x[-1] - x[0]) / (len(x) - 1)
+            x_indexprev = -1
+            for xi_loop in range(len(xi) - 2):
+                x_indexcur = max(int(floor((xi[1 + xi_loop] - x_start) / x_step)), -1)
+                x_index[1 + x_indexprev:1 + x_indexcur] = xi_loop
+                x_indexprev = x_indexcur
+            x_index[1 + x_indexprev:] = len(xi) - 2
+            if x_decreasing:
+                x = x[::-1]
+                x_index = x_index[::-1]
+        elif all(x_steps > 0) or all(x_steps < 0):
+            # non-uniform input/output grids, output grid monotonic
+            if verbose:
+                print("pchip: non-uniform in/out grid, output grid monotonic")
+            x_decreasing = x[-1] < x[0]
+            if x_decreasing:
+                x = x[::-1]
+            x_len = len(x)
+            x_loop = 0
+            for xi_loop in range(len(xi) - 1):
+                while x_loop < x_len and x[x_loop] < xi[1 + xi_loop]:
+                    x_index[x_loop] = xi_loop
+                    x_loop += 1
+            x_index[x_loop:] = len(xi) - 2
+            if x_decreasing:
+                x = x[::-1]
+                x_index = x_index[::-1]
+        else:
+            # non-uniform input/output grids, output grid not monotonic
+            if verbose:
+                print("pchip: non-uniform in/out grids, " "output grid not monotonic")
+            for index in range(len(x)):
+                loc = where(x[index] < xi)[0]
+                if loc.size == 0:
+                    x_index[index] = len(xi) - 2
+                elif loc[0] == 0:
+                    x_index[index] = 0
+                else:
+                    x_index[index] = loc[0] - 1
+        # Calculate gradients d
+        h = diff(xi)
+        d = zeros(len(xi), dtype="double")
+        delta = diff(yi) / h
+        if mode == "quad":
+            # quadratic polynomial fit
+            d[[0, -1]] = delta[[0, -1]]
+            d[1:-1] = (delta[1:] * h[0:-1] + delta[0:-1] * h[1:]) / (h[0:-1] + h[1:])
+        else:
+            # mode=='mono', Fritsch-Carlson algorithm from fortran numerical
+            # recipe
+            d = concatenate(
+                (delta[0:1], 3 * (h[0:-1] + h[1:]) / ((h[0:-1] + 2 * h[1:]) / delta[0:-1] +
+                                                      (2 * h[0:-1] + h[1:]) / delta[1:]), delta[-1:]))
+            d[concatenate((array([False]), logical_xor(delta[0:-1] > 0, delta[1:] > 0), array([False])))] = 0
+            d[logical_or(concatenate((array([False]), delta == 0)), concatenate(
+                (delta == 0, array([False]))))] = 0
+        dxxi = x - xi[x_index]
+        dxxid = x - xi[1 + x_index]
+        dxxi2 = pow(dxxi, 2)
+        dxxid2 = pow(dxxid, 2)
+        y = (2 / pow(h[x_index], 3) *
+             (yi[x_index] * dxxid2 * (dxxi + h[x_index] / 2) - yi[1 + x_index] * dxxi2 *
+              (dxxid - h[x_index] / 2)) + 1 / pow(h[x_index], 2) *
+             (d[x_index] * dxxid2 * dxxi + d[1 + x_index] * dxxi2 * dxxid))
+    return y
+
+
+def Interpolate1D(x, y, xx, method='nearest'):
+    '''
+        Functionality:
+            1D interpolation with various methods
+        Author:
+            Kai Gao <[email protected]>
+    '''
+
+    n = len(x)
+    nn = len(xx)
+    yy = np.zeros(nn)
+
+    # Nearest neighbour interpolation
+    if method == 'nearest':
+        for i in range(0, nn):
+            xi = tf.argmin(tf.abs(xx[i] - x))
+            yy[i] = y[xi]
+
+    # Linear interpolation
+    elif method == 'linear':
+
+        # # slower version
+        # if n == 1:
+        #     yy[:-1] = y[0]
+
+        # else:
+        #     for i in range(0, nn):
+
+        #         if xx[i] < x[0]:
+        #             t = (xx[i] - x[0]) / (x[1] - x[0])
+        #             yy[i] = (1.0 - t) * y[0] + t * y[1]
+
+        #         elif x[n - 1] <= xx[i]:
+        #             t = (xx[i] - x[n - 2]) / (x[n - 1] - x[n - 2])
+        #             yy[i] = (1.0 - t) * y[n - 2] + t * y[n - 1]
+
+        #         else:
+        #             for k in range(1, n):
+        #                 if x[k - 1] <= xx[i] and xx[i] < x[k]:
+        #                     t = (xx[i] - x[k - 1]) / (x[k] - x[k - 1])
+        #                     yy[i] = (1.0 - t) * y[k - 1] + t * y[k]
+        #                     break
+
+        # # faster version
+        yy = linear_interpolate(x, y, xx)
+
+    # Cubic interpolation
+    elif method == 'cubic':
+        yy = cubic_interpolate(x, y, xx)
+
+    # Piecewise cubic Hermite interpolating polynomial (PCHIP)
+    elif method == 'pchip':
+        yy = pchip_interpolate(x, y, xx, mode='mono')
+
+    return yy
+
+
+def Interpolate2D(x, y, f, xx, yy, method='nearest'):
+    '''
+        Functionality:
+            2D interpolation implemented in a separable fashion
+            There are methods that do real 2D non-separable interpolation, which are
+                more difficult to implement.
+        Author:
+            Kai Gao <[email protected]>
+    '''
+
+    n1 = len(x)
+    n2 = len(y)
+    nn1 = len(xx)
+    nn2 = len(yy)
+
+    w = np.zeros((nn1, n2))
+    ff = np.zeros((nn1, nn2))
+
+    # Interpolate along the 1st dimension
+    for j in range(0, n2):
+        w[:, j] = Interpolate1D(x, f[:, j], xx, method)
+
+    # Interpolate along the 2nd dimension
+    for i in range(0, nn1):
+        ff[i, :] = Interpolate1D(y, w[i, :], yy, method)
+
+    return ff
+
+
+def Interpolate3D(x, y, z, f, xx, yy, zz, method='nearest'):
+    '''
+        Functionality:
+            3D interpolation implemented in a separable fashion
+            There are methods that do real 3D non-separable interpolation, which are
+                more difficult to implement.
+        Author:
+            Kai Gao <[email protected]>
+    '''
+
+    n1 = len(x)
+    n2 = len(y)
+    n3 = len(z)
+    nn1 = len(xx)
+    nn2 = len(yy)
+    nn3 = len(zz)
+
+    w1 = tf.zeros((nn1, n2, n3))
+    w2 = tf.zeros((nn1, nn2, n3))
+    ff = tf.zeros((nn1, nn2, nn3))
+
+    # Interpolate along the 1st dimension
+    for k in range(0, n3):
+        for j in range(0, n2):
+            w1[:, j, k] = Interpolate1D(x, f[:, j, k], xx, method)
+
+    # Interpolate along the 2nd dimension
+    for k in range(0, n3):
+        for i in range(0, nn1):
+            w2[i, :, k] = Interpolate1D(y, w1[i, :, k], yy, method)
+
+    # Interpolate along the 3rd dimension
+    for j in range(0, nn2):
+        for i in range(0, nn1):
+            ff[i, j, :] = Interpolate1D(z, w2[i, j, :], zz, method)
+
+    return ff
+
+
+def UpInterpolate1D(x, size=2, interpolation='nearest', data_format='channels_first', align_corners=True):
+    '''
+        Functionality:
+            1D upsampling interpolation for tf
+        Author:
+            Kai Gao <[email protected]>
+    '''
+
+    x = x.numpy()
+
+    if data_format == 'channels_last':
+        nb, nr, nh = x.shape
+    elif data_format == 'channels_first':
+        nb, nh, nr = x.shape
+
+    r = size
+    ir = np.linspace(0.0, nr - 1.0, num=nr)
+
+    if align_corners:
+        # align_corners=True assumes that values are sampled at discrete points
+        iir = np.linspace(0.0, nr - 1.0, num=nr * r)
+    else:
+        # aling_corners=False assumes that values are sampled at centers of discrete blocks
+        iir = np.linspace(0.0 - 0.5 + 0.5 / r, nr - 1.0 + 0.5 - 0.5 / r, num=nr * r)
+        iir = np.clip(iir, 0.0, nr - 1.0)
+
+    if data_format == 'channels_last':
+        xx = np.zeros((nb, nr * r, nh))
+        for i in range(0, nb):
+            for j in range(0, nh):
+                t = np.reshape(x[i, :, j], (nr))
+                xx[i, :, j] = Interpolate1D(ir, t, iir, interpolation)
+
+    elif data_format == 'channels_first':
+        xx = np.zeros((nb, nh, nr * r))
+        for i in range(0, nb):
+            for j in range(0, nh):
+                t = np.reshape(x[i, j, :], (nr))
+                xx[i, j, :] = Interpolate1D(ir, t, iir, interpolation)
+
+    return tf.convert_to_tensor(xx, dtype=x.dtype)
+
+
+def UpInterpolate2D(x,
+                    size=(2, 2),
+                    interpolation='nearest',
+                    data_format='channels_first',
+                    align_corners=True):
+    '''
+        Functionality:
+            2D upsampling interpolation for tf
+        Author:
+            Kai Gao <[email protected]>
+    '''
+
+    x = x.numpy()
+
+    if data_format == 'channels_last':
+        nb, nr, nc, nh = x.shape
+    elif data_format == 'channels_first':
+        nb, nh, nr, nc = x.shape
+
+    r = size[0]
+    c = size[1]
+    ir = np.linspace(0.0, nr - 1.0, num=nr)
+    ic = np.linspace(0.0, nc - 1.0, num=nc)
+
+    if align_corners:
+        # align_corners=True assumes that values are sampled at discrete points
+        iir = np.linspace(0.0, nr - 1.0, num=nr * r)
+        iic = np.linspace(0.0, nc - 1.0, num=nc * c)
+    else:
+        # aling_corners=False assumes that values are sampled at centers of discrete blocks
+        iir = np.linspace(0.0 - 0.5 + 0.5 / r, nr - 1.0 + 0.5 - 0.5 / r, num=nr * r)
+        iic = np.linspace(0.0 - 0.5 + 0.5 / c, nc - 1.0 + 0.5 - 0.5 / c, num=nc * c)
+        iir = np.clip(iir, 0.0, nr - 1.0)
+        iic = np.clip(iic, 0.0, nc - 1.0)
+
+    if data_format == 'channels_last':
+        xx = np.zeros((nb, nr * r, nc * c, nh))
+        for i in range(0, nb):
+            for j in range(0, nh):
+                t = np.reshape(x[i, :, :, j], (nr, nc))
+                xx[i, :, :, j] = Interpolate2D(ir, ic, t, iir, iic, interpolation)
+
+    elif data_format == 'channels_first':
+        xx = np.zeros((nb, nh, nr * r, nc * c))
+        for i in range(0, nb):
+            for j in range(0, nh):
+                t = np.reshape(x[i, j, :, :], (nr, nc))
+                xx[i, j, :, :] = Interpolate2D(ir, ic, t, iir, iic, interpolation)
+
+    return tf.convert_to_tensor(xx, dtype=x.dtype)
+
+
+def UpInterpolate3D(x,
+                    size=(2, 2, 2),
+                    interpolation='nearest',
+                    data_format='channels_first',
+                    align_corners=True):
+    '''
+        Functionality:
+            3D upsampling interpolation for tf
+        Author:
+            Kai Gao <[email protected]>
+    '''
+
+    # x = x.numpy()
+
+    if data_format == 'channels_last':
+        nb, nr, nc, nd, nh = tf.TensorShape(x).as_list()
+    elif data_format == 'channels_first':
+        nb, nh, nr, nc, nd = tf.TensorShape(x).as_list()
+
+    r = size[0]
+    c = size[1]
+    d = size[2]
+    ir = tf.linspace(0.0, nr - 1.0, num=nr)
+    ic = tf.linspace(0.0, nc - 1.0, num=nc)
+    id = tf.linspace(0.0, nd - 1.0, num=nd)
+
+    if align_corners:
+        # align_corners=True assumes that values are sampled at discrete points
+        iir = tf.linspace(0.0, nr - 1.0, num=nr * r)
+        iic = tf.linspace(0.0, nc - 1.0, num=nc * c)
+        iid = tf.linspace(0.0, nd - 1.0, num=nd * d)
+    else:
+        # aling_corners=False assumes that values are sampled at centers of discrete blocks
+        iir = tf.linspace(0.0 - 0.5 + 0.5 / r, nr - 1.0 + 0.5 - 0.5 / r, num=nr * r)
+        iic = tf.linspace(0.0 - 0.5 + 0.5 / c, nc - 1.0 + 0.5 - 0.5 / c, num=nc * c)
+        iid = tf.linspace(0.0 - 0.5 + 0.5 / d, nd - 1.0 + 0.5 - 0.5 / d, num=nd * d)
+        iir = tf.clip_by_value(iir, 0.0, nr - 1.0)
+        iic = tf.clip_by_value(iic, 0.0, nc - 1.0)
+        iid = tf.clip_by_value(iid, 0.0, nd - 1.0)
+
+    if data_format == 'channels_last':
+        xx = tf.zeros((nb, nr * r, nc * c, nd * d, nh))
+        for i in range(0, nb):
+            for j in range(0, nh):
+                t = tf.reshape(x[i, :, :, :, j], (nr, nc, nd))
+                xx[i, :, :, :, j] = Interpolate3D(ir, ic, id, t, iir, iic, iid, interpolation)
+
+    elif data_format == 'channels_first':
+        xx = tf.zeros((nb, nh, nr * r, nc * c, nd * d))
+        for i in range(0, nb):
+            for j in range(0, nh):
+                t = tf.reshape(x[i, j, :, :, :], (nr, nc, nd))
+                xx[i, j, :, :, :] = Interpolate3D(ir, ic, id, t, iir, iic, iid, interpolation)
+
+    return tf.convert_to_tensor(xx, dtype=x.dtype)
+
+
+# ################################################################################
+@keras_export('keras.layers.UpSampling1D')
+class UpSampling1D(Layer):
+    """Upsampling layer for 1D inputs.
+  Repeats each temporal step `size` times along the time axis.
+  Examples:
+  >>> input_shape = (2, 2, 3)
+  >>> x = np.arange(np.prod(input_shape)).reshape(input_shape)
+  >>> print(x)
+  [[[ 0  1  2]
+    [ 3  4  5]]
+    [[ 6  7  8]
+    [ 9 10 11]]]
+  >>> y = tf.keras.layers.UpSampling1D(size=2)(x)
+  >>> print(y)
+  tf.Tensor(
+    [[[ 0  1  2]
+      [ 0  1  2]
+      [ 3  4  5]
+      [ 3  4  5]]
+      [[ 6  7  8]
+      [ 6  7  8]
+      [ 9 10 11]
+      [ 9 10 11]]], shape=(2, 4, 3), dtype=int64)
+  Args:
+    size: Integer. Upsampling factor.
+  Input shape:
+    3D tensor with shape: `(batch_size, steps, features)`.
+  Output shape:
+    3D tensor with shape: `(batch_size, upsampled_steps, features)`.
+  """
+    def __init__(self, size=2, data_format='None', interpolation='nearest', align_corners=True, **kwargs):
+        super(UpSampling1D, self).__init__(**kwargs)
+        self.data_format = conv_utils.normalize_data_format(data_format)
+        self.size = int(size)
+        self.input_spec = InputSpec(ndim=3)
+        self.interpolation = interpolation
+        if self.interpolation not in {'nearest', 'linear', 'cubic', 'pchip'}:
+            raise ValueError('`interpolation` argument should be one of `"nearest"` '
+                             'or `"linear"` '
+                             'or `"cubic"` '
+                             'or `"pchip"`.')
+        self.align_corners = align_corners
+
+    def compute_output_shape(self, input_shape):
+        input_shape = tf.TensorShape(input_shape).as_list()
+        size = self.size * input_shape[1] if input_shape[1] is not None else None
+        return tf.TensorShape([input_shape[0], size, input_shape[2]])
+
+    def call(self, inputs):
+        return UpInterpolate1D(inputs,
+                               self.size,
+                               data_format=self.data_format,
+                               interpolation=self.interpolation,
+                               align_corners=self.align_corners)
+
+    def get_config(self):
+        config = {'size': self.size}
+        base_config = super(UpSampling1D, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@keras_export('keras.layers.UpSampling2D')
+class UpSampling2D(Layer):
+    """Upsampling layer for 2D inputs.
+  Repeats the rows and columns of the data
+  by `size[0]` and `size[1]` respectively.
+  Examples:
+  >>> input_shape = (2, 2, 1, 3)
+  >>> x = np.arange(np.prod(input_shape)).reshape(input_shape)
+  >>> print(x)
+  [[[[ 0  1  2]]
+    [[ 3  4  5]]]
+    [[[ 6  7  8]]
+    [[ 9 10 11]]]]
+  >>> y = tf.keras.layers.UpSampling2D(size=(1, 2))(x)
+  >>> print(y)
+  tf.Tensor(
+    [[[[ 0  1  2]
+        [ 0  1  2]]
+      [[ 3  4  5]
+        [ 3  4  5]]]
+      [[[ 6  7  8]
+        [ 6  7  8]]
+      [[ 9 10 11]
+        [ 9 10 11]]]], shape=(2, 2, 2, 3), dtype=int64)
+  Args:
+    size: Int, or tuple of 2 integers.
+      The upsampling factors for rows and columns.
+    data_format: A string,
+      one of `channels_last` (default) or `channels_first`.
+      The ordering of the dimensions in the inputs.
+      `channels_last` corresponds to inputs with shape
+      `(batch_size, height, width, channels)` while `channels_first`
+      corresponds to inputs with shape
+      `(batch_size, channels, height, width)`.
+      It defaults to the `image_data_format` value found in your
+      Keras config file at `~/.keras/keras.json`.
+      If you never set it, then it will be "channels_last".
+    interpolation: A string, one of `nearest` or `bilinear`.
+  Input shape:
+    4D tensor with shape:
+    - If `data_format` is `"channels_last"`:
+        `(batch_size, rows, cols, channels)`
+    - If `data_format` is `"channels_first"`:
+        `(batch_size, channels, rows, cols)`
+  Output shape:
+    4D tensor with shape:
+    - If `data_format` is `"channels_last"`:
+        `(batch_size, upsampled_rows, upsampled_cols, channels)`
+    - If `data_format` is `"channels_first"`:
+        `(batch_size, channels, upsampled_rows, upsampled_cols)`
+  """
+    def __init__(self, size=(2, 2), data_format=None, interpolation='nearest', align_corners=True, **kwargs):
+        super(UpSampling2D, self).__init__(**kwargs)
+        self.data_format = conv_utils.normalize_data_format(data_format)
+        self.size = conv_utils.normalize_tuple(size, 2, 'size')
+        self.input_spec = InputSpec(ndim=4)
+        self.interpolation = interpolation
+        if self.interpolation not in {'nearest', 'bilinear', 'linear', 'cubic', 'pchip'}:
+            raise ValueError('`interpolation` argument should be one of `"nearest"` '
+                             'or `"bilinear"` '
+                             'or `"linear"` '
+                             'or `"cubic"` '
+                             'or `"pchip"`.')
+        if self.interpolation == 'bilinear':
+            self.interpolation = 'linear'
+        self.align_corners = align_corners
+
+    def compute_output_shape(self, input_shape):
+        input_shape = tensor_shape.TensorShape(input_shape).as_list()
+        if self.data_format == 'channels_first':
+            height = self.size[0] * input_shape[2] if input_shape[2] is not None else None
+            width = self.size[1] * input_shape[3] if input_shape[3] is not None else None
+            return tensor_shape.TensorShape([input_shape[0], input_shape[1], height, width])
+        else:
+            height = self.size[0] * input_shape[1] if input_shape[1] is not None else None
+            width = self.size[1] * input_shape[2] if input_shape[2] is not None else None
+            return tensor_shape.TensorShape([input_shape[0], height, width, input_shape[3]])
+
+    def call(self, inputs):
+        return UpInterpolate2D(inputs,
+                               self.size,
+                               data_format=self.data_format,
+                               interpolation=self.interpolation,
+                               align_corners=self.align_corners)
+
+    def get_config(self):
+        config = {'size': self.size, 'data_format': self.data_format, 'interpolation': self.interpolation}
+        base_config = super(UpSampling2D, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@keras_export('keras.layers.UpSampling3D')
+class UpSampling3D(Layer):
+    """Upsampling layer for 3D inputs.
+  Repeats the 1st, 2nd and 3rd dimensions
+  of the data by `size[0]`, `size[1]` and `size[2]` respectively.
+  Examples:
+  >>> input_shape = (2, 1, 2, 1, 3)
+  >>> x = tf.constant(1, shape=input_shape)
+  >>> y = tf.keras.layers.UpSampling3D(size=2)(x)
+  >>> print(y.shape)
+  (2, 2, 4, 2, 3)
+  Args:
+    size: Int, or tuple of 3 integers.
+      The upsampling factors for dim1, dim2 and dim3.
+    data_format: A string,
+      one of `channels_last` (default) or `channels_first`.
+      The ordering of the dimensions in the inputs.
+      `channels_last` corresponds to inputs with shape
+      `(batch_size, spatial_dim1, spatial_dim2, spatial_dim3, channels)`
+      while `channels_first` corresponds to inputs with shape
+      `(batch_size, channels, spatial_dim1, spatial_dim2, spatial_dim3)`.
+      It defaults to the `image_data_format` value found in your
+      Keras config file at `~/.keras/keras.json`.
+      If you never set it, then it will be "channels_last".
+  Input shape:
+    5D tensor with shape:
+    - If `data_format` is `"channels_last"`:
+        `(batch_size, dim1, dim2, dim3, channels)`
+    - If `data_format` is `"channels_first"`:
+        `(batch_size, channels, dim1, dim2, dim3)`
+  Output shape:
+    5D tensor with shape:
+    - If `data_format` is `"channels_last"`:
+        `(batch_size, upsampled_dim1, upsampled_dim2, upsampled_dim3, channels)`
+    - If `data_format` is `"channels_first"`:
+        `(batch_size, channels, upsampled_dim1, upsampled_dim2, upsampled_dim3)`
+  """
+    def __init__(self,
+                 size=(2, 2, 2),
+                 data_format=None,
+                 interpolation='nearest',
+                 align_corners=True,
+                 **kwargs):
+        super(UpSampling3D, self).__init__(**kwargs)
+        self.data_format = conv_utils.normalize_data_format(data_format)
+        self.size = conv_utils.normalize_tuple(size, 3, 'size')
+        self.input_spec = InputSpec(ndim=5)
+        self.interpolation = interpolation
+        if interpolation not in {'nearest', 'trilinear', 'linear', 'cubic', 'pchip'}:
+            raise ValueError('`interpolation` argument should be one of `"nearest"` '
+                             'or `"trilinear"` '
+                             'or `"linear"` '
+                             'or `"cubic"` '
+                             'or `"pchip"`.')
+        if self.interpolation == 'trilinear':
+            self.interpolation = 'linear'
+        self.align_corners = align_corners
+
+    def compute_output_shape(self, input_shape):
+        input_shape = tensor_shape.TensorShape(input_shape).as_list()
+        if self.data_format == 'channels_first':
+            dim1 = self.size[0] * input_shape[2] if input_shape[2] is not None else None
+            dim2 = self.size[1] * input_shape[3] if input_shape[3] is not None else None
+            dim3 = self.size[2] * input_shape[4] if input_shape[4] is not None else None
+            return tensor_shape.TensorShape([input_shape[0], input_shape[1], dim1, dim2, dim3])
+        else:
+            dim1 = self.size[0] * input_shape[1] if input_shape[1] is not None else None
+            dim2 = self.size[1] * input_shape[2] if input_shape[2] is not None else None
+            dim3 = self.size[2] * input_shape[3] if input_shape[3] is not None else None
+            return tensor_shape.TensorShape([input_shape[0], dim1, dim2, dim3, input_shape[4]])
+
+    def call(self, inputs):
+        return UpInterpolate3D(inputs,
+                               self.size,
+                               data_format=self.data_format,
+                               interpolation=self.interpolation,
+                               align_corners=self.align_corners)
+
+    def get_config(self):
+        config = {'size': self.size, 'data_format': self.data_format}
+        base_config = super(UpSampling3D, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))

DeepDeformationMapRegistration/losses.py

@@ -1,12 +1,17 @@
 import tensorflow as tf
+import tensorflow.keras.backend as K
 from scipy.ndimage import generate_binary_structure
+from sklearn.utils.extmath import cartesian
 
-from DeepDeformationMapRegistration.utils.operators import soft_threshold
+from DeepDeformationMapRegistration.utils.operators import soft_threshold, min_max_norm, hard_threshold
 from DeepDeformationMapRegistration.utils.constants import EPS_tf
+from DeepDeformationMapRegistration.utils.misc import function_decorator
 
+import numpy as np
+import warnings
 
 class HausdorffDistanceErosion:
-    def __init__(self, ndim=3, nerosion=10, value_per_channel=False):
+    def __init__(self, ndim=3, nerosion=10, im_shape: [list, tuple] = (64, 64, 64, 1), alpha=2):
         """
         Approximation of the Hausdorff distance based on erosion operations based on the work done by Karimi D., et al.
             Karimi D., et al., "Reducing the Hausdorff Distance in Medical Image Segmentation with Convolutional Neural
@@ -14,50 +19,222 @@ class HausdorffDistanceErosion:
 
         :param ndim: Dimensionality of the images
         :param nerosion: Number of erosion steps. Defaults to 10.
-        :param value_per_channel: Return an array with the HD distance computed on each channel independently or the sum
+        :param alpha: Parameter to penalize large segmentations. Defaults to 2
         """
+        assert len(im_shape) == ndim + 1, "im_shape does not match with ndim. Missing channel dimension?"
         self.ndims = ndim
-        self.conv = getattr(tf.nn, 'conv%dd' % self.ndims)
+        axes = np.arange(0, self.ndims).tolist()
+        self.before_erosion_transp = [axes[-1], *axes[:-1]]  # [H, W, ..., C] -> [C, H, W, ...]
+        self.after_erosion_transp = [*axes[1:], axes[0]]  # [C, H, W, ...] -> [H, W, ..., C]
         self.nerosions = nerosion
-        self.sum_range = tf.range(0, self.ndims) if value_per_channel else None
-
-    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')
+        self.sum_range = tf.range(0, self.ndims)
 
-    def _erode_per_channel(self, in_tensor, kernel):
-        # In the lambda function we add a fictitious channel and then remove it, so the final shape is [1, H, W, D]
-        er_tensor = tf.map_fn(lambda tens: tf.squeeze(self._erode(tf.expand_dims(tens, -1), kernel)),
-                         tf.transpose(in_tensor, [3, 0, 1, 2]), tf.float32)  # Iterate along the channel dimension (3)
+        self.im_shape = im_shape
+        self.im_vol = np.prod(im_shape[:-1])
+        kernel = generate_binary_structure(self.ndims, 1).astype(int)
+        kernel = kernel / np.sum(kernel)
+        kernel = kernel[..., np.newaxis, np.newaxis]
+        self.kernel = tf.convert_to_tensor(kernel, tf.float32)
+        self.k_alpha = [np.power(k, alpha).astype(float) for k in range(1, nerosion + 1)]
+        self.conv = getattr(tf.nn, 'conv%dd' % self.ndims)
 
-        return tf.transpose(er_tensor, [1, 2, 3, 0])  # move the channels back to the end
+    def _erode(self, in_tensor):
+        indiv_channels = tf.split(in_tensor, self.im_shape[-1], -1)
+        res = list()
+        with tf.variable_scope('erode', reuse=tf.AUTO_REUSE):
+            for ch in indiv_channels:
+                res.append(self.conv(tf.expand_dims(ch, 0), self.kernel, [1] * (self.ndims + 2), 'SAME'))
+        # out = -tf.nn.max_pool3d(-tf.expand_dims(in_tensor, 0), [3]*self.ndims, [1]*self.ndims, 'SAME', name='HDE_erosion')
+        out = tf.concat(res, -1)
+        out = tf.squeeze(out, axis=0)
+        out = hard_threshold(out, 0.5, name='thresholding')  # soft_threshold(out, 0.5, name='thresholding')
+        return out
 
     def _erosion_distance_single(self, y_true, y_pred):
-        diff = tf.math.pow(y_pred - y_true, 2)
+        diff = tf.math.pow(y_pred - y_true, 2, name='HDE_diff')
         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)   # [H, W, D, C_in, C_out]
-
         ret = 0.
         for k in range(1, self.nerosions+1):
             er = diff
             # k successive erosions
             for j in range(k):
-                er = self._erode_per_channel(er, kernel)
-            ret += tf.reduce_sum(tf.multiply(er, tf.cast(tf.pow(k, alpha), tf.float32)), self.sum_range)
+                er = self._erode(er)       # er contains the eroded version along the channels
+            ret += tf.reduce_sum(tf.multiply(er, self.k_alpha[k - 1]), self.sum_range, name='HDE_ret')
 
-        img_vol = tf.cast(tf.reduce_prod(tf.shape(y_true)[:-1]), tf.float32)  # Volume of each channel
-        return tf.divide(ret, img_vol)  # Divide by the image size
+        return tf.divide(ret, self.im_vol)  # Divide by the image size
 
-    def loss(self, y_true, y_pred):
+    @function_decorator('Hausdorff_erosion__loss')
+    def loss(self, y_true, y_pred, name='HDE_loss'):
         batched_dist = tf.map_fn(lambda x: self._erosion_distance_single(x[0], x[1]), (y_true, y_pred),
-                                 dtype=tf.float32)
+                                 dtype=tf.float32, name=name+'_map_fn')
 
         return tf.reduce_mean(batched_dist)
 
+    @function_decorator('Hausdorff_erosion__metric')
+    def metric(self, y_true, y_pred):
+        return self.loss(y_true, y_pred, name='HDE_metric')
+
+    def debug(self, y_true, y_pred):
+        return tf.map_fn(lambda x: self._erosion_distance_single(x[0], x[1]), (y_true, y_pred),
+                         dtype=tf.float32, name='HDE_loss_map_fn')
+
+
+# class HausdorffDiatanceConvolution:
+#     def __init__(self, ndim=3, im_shape: tuple = (64, 64, 64, 1), max_kernel_size=9, step_kernel_size=3, alpha=2):
+#         """
+#         Approximation of the Hausdorff distance based on erosion operations based on the work done by Karimi D., et al.
+#             Karimi D., et al., "Reducing the Hausdorff Distance in Medical Image Segmentation with Convolutional Neural
+#             Networks". IEEE Transactions on Medical Imaging, 39, 2020. DOI 10.1109/TMI.2019.2930068
+#
+#         :param ndim: Dimensionality of the images
+#         :param nerosion: Number of erosion steps. Defaults to 10.
+#         :param alpha: Parameter to penalize large segmentations. Defaults to 2
+#         """
+#         assert len(im_shape) == ndim + 1, "im_shape does not match with ndim. Missing channel dimension?"
+#         self.ndims = ndim
+#         axes = np.arange(0, self.ndims).tolist()
+#         self.before_erosion_transp = [axes[-1], *axes[:-1]]  # [H, W, ..., C] -> [C, H, W, ...]
+#         self.after_erosion_transp = [*axes[1:], axes[0]]  # [C, H, W, ...] -> [H, W, ..., C]
+#         self.conv = getattr(tf.nn, 'conv%dd' % self.ndims)
+#         self.sum_range = tf.range(0, self.ndims)
+#
+#         self.im_shape = im_shape
+#         self.im_vol = np.prod(im_shape[:-1])
+#         kernel = generate_binary_structure(self.ndims, 1).astype(int)
+#         self.kernel = tf.constant(kernel / np.sum(kernel), tf.float32)
+#         self.kernel = tf.expand_dims(tf.expand_dims(self.kernel, -1), -1)   # [H, W, D, C_in, C_out]
+#         self.kernel = tf.tile(self.kernel, [*[1]*self.ndims, self.im_shape[-1], self.im_shape[-1]])
+#         self.alpha = int(alpha)
+#         self.radii = np.arange(1, max_kernel_size, step=step_kernel_size)
+#         self.radii_alpha = [np.pow(r, alpha).astype(float) for r in self.radii]
+#
+#     def soft_diff(self, p, q):
+#         return tf.multiply(tf.pow(p - q, 2.), q)
+#
+#     def body(self, y_true, y_pred):
+#
+
+"""
+class HausdorffDistanceErosion_2:
+    def __init__(self, im_shape, num_erosions, num_dimensions=3, alpha=2., loop_max_iterations=20):
+        self.alpha = alpha
+        self.ndims = num_dimensions
+        self.conv = getattr(tf.nn, 'conv%dd' % self.ndims)
+
+        self.iterator = tf.constant(num_erosions, name='num_erosions')
+        self.norm = 1 / np.prod(im_shape)
+        self.erosion_kernel = generate_binary_structure(self.ndims, 1).astype(float)
+        self.erosion_kernel /= np.sum(self.erosion_kernel)
+        self.erosion_kernel = tf.constant( self.erosion_kernel, tf.float32)
+
+        self.loop_max_iterations = loop_max_iterations
+
+    def erosion_sum(self, p, q, k):
+        er_tensor = p - q
+        er_tensor = tf.pow(er_tensor, 2.)
+
+        def erode(in_tensor):
+            # Erosion of in_tensor = Dilation of (1 - in_tensor)
+            return self.conv(tf.expand_dims(1. - in_tensor, 0), self.erosion_kernel, [1] * (self.ndims + 2), 'SAME')
+
+        def while_loop_body(i, in_tensor):
+            in_tensor = erode(in_tensor)
+            i -= 1
+            return i, in_tensor
+
+        def while_loop_condition(i, in_tensor):
+            return tf.less_equal(i, 1), in_tensor
+
+        er_iterator = tf.constant(k)
+        _, er_tensor = tf.while_loop(while_loop_condition, while_loop_body, loop_vars=[er_iterator, er_tensor],
+                                     maximum_iterations=self.loop_max_iterations)
+
+        er_tensor *= tf.pow(k, self.alpha)
+        return tf.reduce_sum(er_tensor)
+
+    def loss(self, y_true, y_pred):
+        hd_distance = tf.constant(0, name='hausdroff_distance')
+
+        def while_loop_body(i, p, q, ret):
+            i -= 1
+            return i, p, q, ret + self.erosion_sum(p, q, i)
+
+        _, _, _, hd_distance = tf.while_loop(lambda i, p, q, ret: tf.less_equal(i, 1),
+                                             while_loop_body,
+                                             loop_vars=[self.iterator, y_pred, y_true, hd_distance])
+        hd_distance /= self.norm
+        return hd_distance
+
+"""
+
+
+class WeightedHausdorffDistance:
+    def __init__(self, input_shape, alpha=-1, threshold=0.5):
+        """
+        WARNING: Requires a insane amount of memory
+        :param input_shape: [H, W, D, C] or [H, W, C]
+        :param alpha: Parameter of the generalized mean. Ideally -inf, but then the function becomes less smooth.
+        :param threshold: Threshold of segmentations, used in tf.where function
+        """
+        warnings.warn("This function requires an insane amount of memory")
+        self.input_shape = input_shape
+        self.dim = len(input_shape[:-1])
+        self.ohe_segm = bool(input_shape[-1] > 1)   # One-Hot Encoded segmentations on the channel axis
+        aux = np.arange(len(self.input_shape)).tolist()
+        self.ohe_transpose = [aux[-1], *aux[:-1]]
+        self.alpha = alpha
+        self.threshold = threshold
+        list_coords = [np.arange(c) for c in self.input_shape[:-1]]
+        self.img_loc = tf.convert_to_tensor(cartesian(list_coords), dtype=tf.float32)
+        self.max_dist = np.sqrt(np.sum(np.square(self.input_shape[:-1])))      # Largest diagonal
+
+    def pairwise_distance(self, A, B):
+        sq_norm_a = tf.reduce_sum(tf.square(A), 1)
+        sq_norm_b = tf.reduce_sum(tf.square(B), 1)
+
+        sq_norm_a = tf.reshape(sq_norm_a, [-1, 1])
+        sq_norm_b = tf.reshape(sq_norm_b, [1, -1])
+
+        return tf.sqrt(tf.maximum(sq_norm_a - 2 * tf.matmul(A, B, transpose_a=False, transpose_b=True) + sq_norm_b, 0.))
+
+    def hausdorff(self, y_true, y_pred):
+        if self.ohe_segm:
+            y_true = tf.transpose(y_true, self.ohe_transpose)
+            y_pred = tf.transpose(y_pred, self.ohe_transpose)
+            hausdorff_per_ch = tf.map_fn(lambda x: self.hausdorff_per_channel(x[0], x[1]), (y_true, y_pred), tf.float32)
+            return tf.reduce_mean(hausdorff_per_ch)
+        else:
+            return self.hausdorff_per_channel(y_true, y_pred)
+
+    def hausdorff_per_channel(self, y_true, y_pred):
+        Y = tf.cast(tf.where(y_true > self.threshold), dtype=tf.float32)
+        p = K.flatten(y_pred)   # Flatten the predicted segmentation (activation map 'p' in d_WH)
+
+        size_Y = tf.shape(Y)[0]
+        S = tf.reduce_sum(p)
+
+        p = tf.squeeze(K.repeat(tf.expand_dims(p, -1), size_Y))
+        dist_mat = self.pairwise_distance(self.img_loc, Y)
+
+        term_1 = tf.reduce_sum(p * tf.minimum(dist_mat, 1)) / (S + EPS_tf)
+
+        term_2 = tf.minimum((dist_mat + EPS_tf) / (tf.pow(p, self.alpha) + (EPS_tf / self.max_dist)), 0.)
+        term_2 = tf.clip_by_value(term_2, 0., self.max_dist)
+        term_2 = tf.reduce_mean(term_2, axis=0)
+
+        return term_1 + term_2
+
+    @function_decorator('Weighted_Hausdorff__loss')
+    def loss(self, y_true, y_pred):
+        batch_hdist = tf.map_fn(lambda x: self.hausdorff(x[0], x[1]), (y_true, y_pred), dtype=tf.float32)
+
+        return tf.reduce_mean(batch_hdist)
+
+    @function_decorator('Weighted_Hausdorff__metric')
+    def metric(self, y_true, y_pred):
+        return self.loss(y_true, y_pred)
+
 
 class NCC:
     def __init__(self, in_shape, eps=EPS_tf):
@@ -69,48 +246,105 @@ class NCC:
         f_yp = tf.reshape(y_pred, [-1])
         mean_yt = tf.reduce_mean(f_yt)
         mean_yp = tf.reduce_mean(f_yp)
-        std_yt = tf.math.reduce_std(f_yt)
-        std_yp = tf.math.reduce_std(f_yp)
 
         n_f_yt = f_yt - mean_yt
         n_f_yp = f_yp - mean_yp
-        numerator = tf.reduce_sum(n_f_yt * n_f_yp)
-        denominator = std_yt * std_yp * self.__shape_size + self.__eps
+        norm_yt = tf.norm(f_yt, ord='euclidean')
+        norm_yp = tf.norm(f_yp, ord='euclidean')
+        numerator = tf.reduce_sum(tf.multiply(n_f_yt, n_f_yp))
+        denominator = norm_yt * norm_yp + self.__eps
         return tf.math.divide_no_nan(numerator, denominator)
 
+    @function_decorator('NCC__loss')
     def loss(self, y_true, y_pred):
         # According to the documentation, the loss returns a scalar
         # Ref: https://www.tensorflow.org/api_docs/python/tf/keras/Model#compile
         return tf.reduce_mean(tf.map_fn(lambda x: 1 - self.ncc(x[0], x[1]), (y_true, y_pred), tf.float32))
 
+    @function_decorator('NCC__metric')
+    def metric(self, y_true, y_pred):
+        return tf.reduce_mean(tf.map_fn(lambda x: self.ncc(x[0], x[1]), (y_true, y_pred), tf.float32))
+
+
+def ncc(y_true, y_pred):
+    y_true = K.flatten(K.cast(y_true, 'float32'))
+    y_pred = K.flatten(K.cast(y_pred, 'float32'))
+
+    mean_true = K.mean(y_true)
+    mean_pred = K.mean(y_pred)
+
+    std_true = K.std(y_true)
+    std_pred = K.std(y_pred)
+
+    num = K.mean((y_true - mean_true) * (y_pred - mean_pred))
+    den = std_true * std_pred + EPS_tf
+    batch_ncc = num / den
+
+    return K.mean(batch_ncc)
+
 
 class StructuralSimilarity:
     # Based on https://github.com/keras-team/keras-contrib/blob/master/keras_contrib/losses/dssim.py
-    def __init__(self, k1=0.01, k2=0.03, patch_size=3, dynamic_range=1., overlap=0.0):
+    def __init__(self, k1=0.01, k2=0.03,
+                 patch_size=32, dynamic_range=1., overlap=0.0, dim=3,
+                 alpha=1., beta=1., gamma=1.,
+                 **kwargs):
         """
         Structural (Di)Similarity Index Measure:
 
         :param k1: Internal parameter. Defaults to 0.01
         :param k2: Internal parameter. Defaults to 0.02
-        :param patch_size: Size of the extracted patches
+        :param patch_size: Size of the extracted patches. Defaults to 32. Recommendation: half the image size.
         :param dynamic_range: Maximum numerical intensity value (typ. 2^bits_per_pixel - 1). Defaults to 1.
+        :param overlap: Patch overlap ratio. Must be in the range [0., 1.). Defaults to 0.
+        :param dim: Data dimensionality. Must be {1, 2, 3}. Defaults to 3.
+        :param alpha, beta, gamma: Exponential parameters to balance the contribution of the luminance, contrast and
+                                    structure measures. Default to 1.
         """
-        self.__c1 = (k1 * dynamic_range) ** 2
-        self.__c2 = (k2 * dynamic_range) ** 2
-        self.__kernel_shape = [1] + [patch_size] * 3 + [1]
+        assert (dim > 0) and (dim < 4), 'Invalid dimension. It must be 1, 2, or 3'
+        assert overlap < 1., 'Invalid overlap. It must be in the range [0., 1.)'
+        self.c1 = (k1 * dynamic_range) ** 2
+        self.c2 = (k2 * dynamic_range) ** 2
+        self.c3 = self.c2 / 2
+        self.alpha = tf.cast(alpha, tf.float32)
+        self.beta = tf.cast(beta, tf.float32)
+        self.gamma = tf.cast(gamma, tf.float32)
+
+        self.kernel_shape = [1] + [patch_size] * dim + [1]
         stride = int(patch_size * (1 - overlap))
-        self.__stride = [1] + [stride if stride else 1] * 3 + [1]
-        self.__max_val = dynamic_range
+        self.stride = [1] + [stride if stride else 1] * dim + [1]
+        self.dim = dim
+        self.patch_extractor = None
+        self.reduce_axis = list()
+        if dim == 2:
+            self.patch_extractor = tf.extract_image_patches
+            self.reduce_axis = [1, 2]
+        elif dim == 3:
+            self.patch_extractor = tf.extract_volume_patches
+            self.reduce_axis = [1, 2, 3]
+        else:
+            raise ValueError('Invalid dimension value. Expected 2 or 3')
+
+        if patch_size == -1:
+            # Don't extract patches
+            self.dim = 1
+
+        self.L = None   # Luminance
+        self.C = None   # Contrast
+        self.S = None   # Structure
 
     def __int_shape(self, x):
         return tf.keras.backend.int_shape(x) if tf.keras.backend.backend() == 'tensorflow' else tf.keras.backend.shape(x)
 
     def ssim(self, y_true, y_pred):
-
-        patches_true = tf.extract_volume_patches(y_true, self.__kernel_shape, self.__stride, 'VALID',
-                                                 'patches_true')
-        patches_pred = tf.extract_volume_patches(y_pred, self.__kernel_shape, self.__stride, 'VALID',
-                                                 'patches_pred')
+        if self.dim > 1:
+            # Don't use for training. The gradient doesn't backpropagate through the patch extractors
+            # patches: [B, out_rows, out_cols, ..., krows*kcols*...*channels] -> out_rows * out_cols * ... = nb patches
+            patches_true = self.patch_extractor(y_true, ksizes=self.kernel_shape, strides=self.stride, padding='VALID', name='patches_true')
+            patches_pred = self.patch_extractor(y_pred, ksizes=self.kernel_shape, strides=self.stride, padding='VALID', name='patches_pred')
+        else:
+            patches_true = y_true
+            patches_pred = y_pred
 
         #bs, w, h, d, *c = self.__int_shape(patches_pred)
         #patches_true = tf.reshape(patches_true, [-1, w, h, d, tf.reduce_prod(c)])
@@ -124,15 +358,496 @@ class StructuralSimilarity:
         v_true = tf.math.reduce_variance(patches_true, axis=-1)
         v_pred = tf.math.reduce_variance(patches_pred, axis=-1)
 
+        # Standard dev.
+        s_true = tf.sqrt(v_true)
+        s_pred = tf.sqrt(v_pred)
+
         # Covariance
         covar = tf.reduce_mean(patches_true * patches_pred, axis=-1) - u_true * u_pred
 
         # SSIM
-        numerator = (2 * u_true * u_pred + self.__c1) * (2 * covar + self.__c2)
-        denominator = ((tf.square(u_true) + tf.square(u_pred) + self.__c1) * (v_pred + v_true + self.__c2))
-        ssim = numerator / denominator
+        self.L = (2 * u_true * u_pred + self.c1) / (tf.square(u_true) + tf.square(u_pred) + self.c1)
+        self.C = (2 * s_true * s_pred + self.c2) / (v_true + v_pred + self.c2)
+        self.S = (covar + self.c3) / (s_true * s_pred + self.c3)
+        self.L = tf.reduce_mean(self.L, axis=self.reduce_axis)
+        self.C = tf.reduce_mean(self.C, axis=self.reduce_axis)
+        self.S = tf.reduce_mean(self.S, axis=self.reduce_axis)
 
-        return tf.reduce_mean(ssim)
+        return tf.pow(self.L, self.alpha) * tf.pow(self.C, self.beta) * tf.pow(self.S, self.gamma)
+
+    @function_decorator('SSIM__loss')
+    def loss(self, y_true, y_pred):
+        return tf.reduce_mean((1. - self.ssim(y_true, y_pred)) / 2.0)
+
+    @function_decorator('SSIM__metric')
+    def metric(self, y_true, y_pred):
+        return tf.reduce_mean(self.ssim(y_true, y_pred))
+
+
+class StructuralSimilarity_simplified:
+    # Based on https://github.com/keras-team/keras-contrib/blob/master/keras_contrib/losses/dssim.py
+    def __init__(self, k1=0.01, k2=0.03,
+                 patch_size=32, dynamic_range=1., overlap=0.0, dim=3,
+                 alpha=1., beta=1., gamma=1.,
+                 **kwargs):
+        """
+        Structural (Di)Similarity Index Measure:
+
+        :param k1: Internal parameter. Defaults to 0.01
+        :param k2: Internal parameter. Defaults to 0.02
+        :param patch_size: Size of the extracted patches. Defaults to 32. Recommendation: half the image size.
+        :param dynamic_range: Maximum numerical intensity value (typ. 2^bits_per_pixel - 1). Defaults to 1.
+        :param overlap: Patch overlap ratio. Must be in the range [0., 1.). Defaults to 0.
+        :param dim: Data dimensionality. Must be {1, 2, 3}. Defaults to 3.
+        :param alpha, beta, gamma: Exponential parameters to balance the contribution of the luminance, contrast and
+                                    structure measures. Default to 1.
+        """
+        assert (dim > 0) and (dim < 4), 'Invalid dimension. It must be 1, 2, or 3'
+        assert overlap < 1., 'Invalid overlap. It must be in the range [0., 1.)'
+        self.c1 = (k1 * dynamic_range) ** 2
+        self.c2 = (k2 * dynamic_range) ** 2
+        self.c3 = self.c2 / 2
+        self.alpha = tf.cast(alpha, tf.float32)
+        self.beta = tf.cast(beta, tf.float32)
+        self.gamma = tf.cast(gamma, tf.float32)
+
+        self.kernel_shape = [1] + [patch_size] * dim + [1]
+        stride = int(patch_size * (1 - overlap))
+        self.stride = [1] + [stride if stride else 1] * dim + [1]
+        self.dim = dim
+        self.patch_extractor = None
+        if dim == 2:
+            self.patch_extractor = tf.extract_image_patches
+        elif dim == 3:
+            self.patch_extractor = tf.extract_volume_patches
+
+        if patch_size == -1:
+            # Don't extract patches
+            self.dim = 1
+
+        self.L = None   # Luminance
+        self.C = None   # Contrast
+        self.S = None   # Structure
 
-    def dssim(self, y_true, y_pred):
+    def __int_shape(self, x):
+        return tf.keras.backend.int_shape(x) if tf.keras.backend.backend() == 'tensorflow' else tf.keras.backend.shape(x)
+
+    def ssim(self, y_true, y_pred):
+        if self.dim > 1:
+            # Don't use for training. The gradient doesn't backpropagate through the patch extractors
+            # patches: [B, out_rows, out_cols, ..., krows*kcols*...*channels] -> out_rows * out_cols * ... = nb patches
+            patches_true = self.patch_extractor(y_true, ksizes=self.kernel_shape, strides=self.stride, padding='VALID', name='patches_true')
+            patches_pred = self.patch_extractor(y_pred, ksizes=self.kernel_shape, strides=self.stride, padding='VALID', name='patches_pred')
+        else:
+            patches_true = y_true
+            patches_pred = y_pred
+
+        #bs, w, h, d, *c = self.__int_shape(patches_pred)
+        #patches_true = tf.reshape(patches_true, [-1, w, h, d, tf.reduce_prod(c)])
+        #patches_pred = tf.reshape(patches_pred, [-1, w, h, d, tf.reduce_prod(c)])
+
+        # Mean
+        u_true = tf.reduce_mean(patches_true, axis=-1)
+        u_pred = tf.reduce_mean(patches_pred, axis=-1)
+
+        # Variance
+        v_true = tf.math.reduce_variance(patches_true, axis=-1)
+        v_pred = tf.math.reduce_variance(patches_pred, axis=-1)
+
+        # Covariance
+        covar = tf.reduce_mean(patches_true * patches_pred, axis=-1) - u_true * u_pred
+
+        # return tf.pow(self.L, self.alpha) * tf.pow(self.C, self.beta) * tf.pow(self.S, self.gamma)
+        num = (2 * u_true * u_pred + self.c1) * (2 * covar + self.c2)
+        den = ((tf.square(u_true) + tf.square(u_pred) + self.c1) * (v_pred + v_true + self.c2))
+        return num / den
+
+    @function_decorator('SSIM_simple__loss')
+    def loss(self, y_true, y_pred):
         return tf.reduce_mean((1. - self.ssim(y_true, y_pred)) / 2.0)
+
+    @function_decorator('SSIM_simple__metric')
+    def metric(self, y_true, y_pred):
+        return tf.reduce_mean(self.ssim(y_true, y_pred))
+
+
+class MultiScaleStructuralSimilarity(StructuralSimilarity):
+    def __init__(self, k1=0.01, k2=0.03, patch_size=3, dynamic_range=1., overlap=0.0, dim=3, nscales=3, alpha=1., beta=1., gamma=1.):
+        """
+        Multi Scale Structural (Di)Similarity Index Measure:
+        Ref:    [1] https://www.cns.nyu.edu/pub/eero/wang03b.pdf
+
+        :param k1: Internal parameter. Defaults to 0.01
+        :param k2: Internal parameter. Defaults to 0.02
+        :param patch_size: Size of the extracted patches. Defaults to 32. Recommendation: half the image size.
+        :param dynamic_range: Maximum numerical intensity value (typ. 2^bits_per_pixel - 1). Defaults to 1.
+        :param overlap: Patch overlap ratio. Must be in the range [0., 1.). Defaults to 0.
+        :param dim: Data dimensionality. Must be {2, 3}. Defaults to 3.
+        :param nscales: Number of scales to analyze. Defaults to 3.
+        :param alpha, beta, gamma: Exponential parameters to balance the contribution of the luminance, contrast and
+                                    structure measures. Default to 1.
+        """
+        assert dim > 1, 'Cannot be used with 1-D data'
+        super(MultiScaleStructuralSimilarity, self).__init__(k1=k1, k2=k2, patch_size=patch_size,
+                                                             dynamic_range=dynamic_range, overlap=overlap, dim=dim,
+                                                             alpha=alpha, beta=beta, gamma=gamma)
+        self.num_scales = nscales
+        self.avg_pool = getattr(tf.nn, 'avg_pool%dd' % dim)
+        self.ds_stride = self.ds_kernel = [1] + [2]*dim + [1]
+
+        # In [1] these are set to the same value at the same scales and normalized across scales
+        self.alpha = self.beta = self.gamma = 1 / nscales
+
+    def _cond(self, cs_prod, scale_level, y_true, y_pred):
+        return tf.less_equal(scale_level, self.num_scales)
+
+    def _iteration(self, cs_prod, scale_level, y_true, y_pred):
+        super(MultiScaleStructuralSimilarity, self).ssim(y_true, y_pred)
+        cs_prod *= tf.reduce_mean(tf.pow(self.C, self.beta) * tf.pow(self.S, self.gamma))
+        y_true = self.avg_pool(y_true, ksize=self.ds_kernel, strides=self.ds_stride, padding='VALID')
+        y_pred = self.avg_pool(y_pred, ksize=self.ds_kernel, strides=self.ds_stride, padding='VALID')
+        scale_level += 1
+        return cs_prod, scale_level, y_true, y_pred,
+
+    def ssim(self, y_true, y_pred):
+        return self.ms_ssim(y_true, y_pred)
+
+    def ms_ssim(self, y_true, y_pred):
+        cs_prod = tf.constant(1.)
+        scale_level = tf.constant(1.)
+        cs_prod, *_ = tf.while_loop(self._cond,
+                                    self._iteration,
+                                    (cs_prod, scale_level, y_true, y_pred),
+                                    (cs_prod.get_shape(), scale_level.get_shape(),
+                                     tf.TensorShape(([1] + [None] * self.dim + [1])),
+                                     tf.TensorShape(([1] + [None] * self.dim + [1]))))
+
+        ms_ssim = tf.reduce_mean(tf.pow(self.L, self.alpha)) * cs_prod
+
+        return tf.reduce_mean(ms_ssim)
+
+    @function_decorator('MS_SSIM__loss')
+    def loss(self, y_true, y_pred):
+        return tf.reduce_mean((1. - self.ms_ssim(y_true, y_pred)) / 2.0)
+
+
+class MultiScaleStructuralSimilarity_v2(StructuralSimilarity):
+    def __init__(self, k1=0.01, k2=0.03, patch_size=3, dynamic_range=1., overlap=0.0, dim=3, nscales=3, alpha=1., beta=1., gamma=1.):
+        """
+        Multi Scale Structural (Di)Similarity Index Measure:
+        Ref:    [1] https://www.cns.nyu.edu/pub/eero/wang03b.pdf
+
+        :param k1: Internal parameter. Defaults to 0.01
+        :param k2: Internal parameter. Defaults to 0.02
+        :param patch_size: Size of the extracted patches. Defaults to 32. Recommendation: half the image size.
+        :param dynamic_range: Maximum numerical intensity value (typ. 2^bits_per_pixel - 1). Defaults to 1.
+        :param overlap: Patch overlap ratio. Must be in the range [0., 1.). Defaults to 0.
+        :param dim: Data dimensionality. Must be {2, 3}. Defaults to 3.
+        :param nscales: Number of scales to analyze. Defaults to 3.
+        :param alpha, beta, gamma: Exponential parameters to balance the contribution of the luminance, contrast and
+                                    structure measures. Default to 1.
+        """
+        assert dim > 1, 'Cannot be used with 1-D data'
+        super(MultiScaleStructuralSimilarity_v2, self).__init__(k1=k1, k2=k2, patch_size=patch_size,
+                                                                dynamic_range=dynamic_range, overlap=overlap, dim=dim,
+                                                                alpha=alpha, beta=beta, gamma=gamma)
+        self.num_scales = nscales
+        self.avg_pool = getattr(tf.nn, 'avg_pool%dd' % dim)
+        self.ds_stride = self.ds_kernel = [1] + [2]*dim + [1]
+
+        # In [1] these are set to the same value at the same scales and normalized across scales
+        self.alpha = self.beta = self.gamma = 1 / nscales
+
+    def _cond(self, cs_prod, scale_level, y_true, y_pred):
+        return tf.less_equal(scale_level, self.num_scales)
+
+    def _iteration(self, cs_prod, scale_level, y_true, y_pred):
+        super(MultiScaleStructuralSimilarity_v2, self).ssim(y_true, y_pred)
+        cs_prod *= tf.reduce_mean(tf.pow(self.C, self.beta) * tf.pow(self.S, self.gamma))
+        y_true = self.avg_pool(y_true, ksize=self.ds_kernel, strides=self.ds_stride, padding='VALID')
+        y_pred = self.avg_pool(y_pred, ksize=self.ds_kernel, strides=self.ds_stride, padding='VALID')
+        scale_level += 1
+        return cs_prod, scale_level, y_true, y_pred,
+
+    def ssim(self, y_true, y_pred):
+        return self.ms_ssim(y_true, y_pred)
+
+    def ms_ssim(self, y_true, y_pred):
+        cs_prod = tf.constant(1.)
+        scale_level = tf.constant(1.)
+        cs_prod, *_ = tf.while_loop(self._cond,
+                                    self._iteration,
+                                    (cs_prod, scale_level, y_true, y_pred),
+                                    (cs_prod.get_shape(), scale_level.get_shape(),
+                                     tf.TensorShape(([1] + [None] * self.dim + [1])),
+                                     tf.TensorShape(([1] + [None] * self.dim + [1]))))
+
+        ms_ssim = tf.reduce_mean(tf.pow(self.L, self.alpha)) * cs_prod
+
+        return tf.reduce_mean(ms_ssim)
+
+    @function_decorator('MS_SSIM_v2__loss')
+    def loss(self, y_true, y_pred):
+        return tf.reduce_mean((1. - self.ms_ssim(y_true, y_pred)) / 2.0)
+
+
+class StructuralSimilarityGaussian:
+    # This is equivalent to StructuralSimilarity(patch_size=img_size)
+    def __init__(self, k1=0.01, k2=0.03, dynamic_range=1., gauss_sigma=5., dim=3, alpha=1., beta=1., gamma=1.):
+        """
+        SSIM using Gaussian filter to approximate the statistics of the images
+        Ref:    https://www.cns.nyu.edu/pub/eero/wang03b.pdf
+                https://arxiv.org/pdf/1511.08861.pdf
+                https://github.com/NVlabs/PL4NN/blob/master/src/loss.py
+
+        :param k1: Internal parameter. Defaults to 0.01
+        :param k2: Internal parameter. Defaults to 0.02
+        :param dynamic_range: Maximum numerical intensity value (typ. 2^bits_per_pixel - 1). Defaults to 1.
+        :param gauss_sigma: Sigma of the Gaussian filter. Defaults to 1.5.
+        :param dim: Data dimensionality. Must be {2, 3}. Defaults to 3.
+        :param alpha, beta, gamma: Exponential parameters to balance the contribution of the luminance, contrast and
+                                    structure measures. Default to 1.
+        """
+        self.c1 = (k1 * dynamic_range) ** 2
+        self.c2 = (k2 * dynamic_range) ** 2
+        self.c3 = self.c2 / 2
+        self.alpha = tf.cast(alpha, tf.float32)
+        self.beta = tf.cast(beta, tf.float32)
+        self.gamma = tf.cast(gamma, tf.float32)
+        self.dim = dim
+        self.convDN = getattr(tf.nn, 'conv%dd' % dim)
+        self.sigma = gauss_sigma
+
+    def build_gaussian_filter(self, size, sigma, num_channels=1):
+        range_1d = tf.range(-(size/2) + 1, size//2 + 1)
+        g_1d = tf.math.exp(-1.0 * tf.pow(range_1d, 2) / (2. * tf.pow(sigma, 2)))
+        g_1d_expanded = tf.expand_dims(g_1d, -1)
+        iterator = tf.constant(1)
+        self.__GF = tf.while_loop(lambda iterator, g_1d: tf.less(iterator, self.dim),
+                                  lambda iterator, g_1d: (iterator + 1, tf.expand_dims(g_1d, -1) * tf.transpose(g_1d_expanded)),
+                                  [iterator, g_1d],
+                                  [iterator.get_shape(), tf.TensorShape([None]*self.dim)]  # Shape invariants
+                                  )[-1]
+
+        self.__GF = tf.divide(self.__GF, tf.reduce_sum(self.__GF))  # Normalization
+        self.__GF = tf.reshape(self.__GF, (*[size]*self.dim, 1, 1))  # Add Ch_in and Ch_out for convolution
+        self.__GF = tf.tile(self.__GF, (*[1] * self.dim, num_channels, num_channels,))
+
+    def format_data(self, in_data):
+        ret_val = in_data
+        if self.dim == 3:
+            ret_val = tf.transpose(ret_val, [0, 3, 1, 2, 4])
+        return ret_val
+
+    def ssim(self, y_true, y_pred):
+        self.build_gaussian_filter(y_pred.shape[1], self.sigma)
+        y_true_tr = self.format_data(y_true)
+        y_pred_tr = self.format_data(y_pred)
+
+        u_true = self.convDN(y_true_tr, self.__GF, [1] * (self.dim + 2), 'SAME')
+        u_pred = self.convDN(y_pred_tr, self.__GF, [1] * (self.dim + 2), 'SAME')
+
+        v_true = self.convDN(tf.pow(y_true_tr, 2), self.__GF, [1] * (self.dim + 2), 'SAME') - tf.pow(u_true, 2)
+        v_pred = self.convDN(tf.pow(y_pred_tr, 2), self.__GF, [1] * (self.dim + 2), 'SAME') - tf.pow(u_pred, 2)
+        covar = self.convDN(tf.multiply(y_true_tr, y_pred_tr), self.__GF, [1] * (self.dim + 2), 'SAME') - u_true * u_pred
+
+        self.L = (2 * u_true * u_pred + self.c1) / (tf.square(u_true) + tf.square(u_pred) + self.c1)
+        self.C = (2 * tf.sqrt(v_true) * tf.sqrt(v_pred) + self.c2) / (v_true + v_pred + self.c2)
+        self.S = (covar + self.c3) / (tf.sqrt(v_true) * tf.sqrt(v_pred) + self.c3)
+        ssim = tf.pow(self.L, self.alpha) * tf.pow(self.C, self.beta) * tf.pow(self.S, self.gamma)
+
+        return tf.reduce_mean(ssim)
+
+    @function_decorator('SSIM_Gaus__loss')
+    def loss(self, y_true, y_pred):
+        return tf.reduce_mean((1. - self.ssim(y_true, y_pred))/2.)
+
+    @function_decorator('SSIM_Gaus__metric')
+    def metric(self, y_true, y_pred):
+        return tf.reduce_mean(self.ssim(y_true, y_pred))
+
+
+class MultiScaleStructuralSimilarityGaussian(StructuralSimilarityGaussian):
+    def __init__(self, k1=0.01, k2=0.03, dynamic_range=1., gauss_sigma=5., dim=3, nscales=3, alpha=1., beta=1., gamma=1.):
+        """
+        Multi Scale SSIM inheriting from StructuralSimilarityGaussian classed
+        Ref:    https://www.cns.nyu.edu/pub/eero/wang03b.pdf
+                https://arxiv.org/pdf/1511.08861.pdf
+                https://github.com/NVlabs/PL4NN/blob/master/src/loss.py
+
+        :param k1: Internal parameter. Defaults to 0.01
+        :param k2: Internal parameter. Defaults to 0.02
+        :param dynamic_range: Maximum numerical intensity value (typ. 2^bits_per_pixel - 1). Defaults to 1.
+        :param gauss_sigma: Sigma of the Gaussian filter. Defaults to 1.5.
+        :param dim: Data dimensionality. Must be {2, 3}. Defaults to 3.
+        :param nscales: Number of scales to analyze. Defaults to 3.
+        :param alpha, beta, gamma: Exponential parameters to balance the contribution of the luminance, contrast and
+                                    structure measures. Default to 1.
+        """
+        super(MultiScaleStructuralSimilarityGaussian, self).__init__(k1=k1, k2=k2, dynamic_range=dynamic_range,
+                                                                     gauss_sigma=gauss_sigma, dim=dim,
+                                                                     alpha=alpha, beta=beta, gamma=gamma)
+        self.__num_scales = nscales
+
+    # # If using the Gaussian approximation of the pyramid MS approach described in https://arxiv.org/pdf/1511.08861.pdf
+    # def build_sigma_scales(self):
+    #     iterator = tf.constant(0)
+    #     scales = tf.expand_dims(self.sigma, -1)
+    #     last_sigma = scales
+    #     self.sigma_scales = tf.while_loop(lambda iterator, last_sigma, scales: tf.less_equal(iterator, self.__num_scales),
+    #                                        lambda iterator, last_sigma, scales: (iterator + 1, tf.concat([scales, last_sigma/2], 0), last_sigma/2),
+    #                                        [iterator, last_sigma, scales])[-1]
+    #
+    # def build_gaussian_filters_scales(self, size):
+    #    self.__GFS = tf.map_fn(lambda sigma: self.build_gaussian_filter(size, sigma), self.sigma, tf.float32)
+
+    def _iteration(self, cs_prod, scale_level, y_true, y_pred):
+        # Compute the SSIM, so CS and L have the correct value
+        self.ssim(y_true, y_pred)
+
+        cs_prod *= tf.reduce_mean(tf.pow(self.C, self.beta) * tf.pow(self.S, self.gamma))
+        scale_level += 1
+
+        # Downsample the images to half the resolution for the next iteration
+        y_true = tf.nn.avg_pool(y_true, [1] + [2]*self.dim + [1], [1] + [2]*self.dim + [1], 'SAME')
+        y_pred = tf.nn.avg_pool(y_true, [1] + [2]*self.dim + [1], [1] + [2]*self.dim + [1], 'SAME')
+        return cs_prod, scale_level, y_true, y_pred
+
+    def ms_ssim(self, y_true, y_pred):
+        scale_level = tf.constant(0.)
+        cs_prod = tf.constant(1.)
+        cs_prod, *_ = tf.while_loop(tf.less(scale_level, self.__num_scales),
+                                    self._iteration,
+                                    (cs_prod, scale_level, y_true, y_pred),
+                                    (cs_prod.get_shape(), scale_level.get_shape(),
+                                     tf.TensorShape(([1] + [None]*self.dim + [1])),
+                                     tf.TensorShape(([1] + [None]*self.dim + [1]))))
+        # L is taken from the last scale
+        return tf.reduce_mean(tf.pow(self.L, self.alfa)) * cs_prod
+
+    @function_decorator('MS_SSIM_Gaus__metric')
+    def loss(self, y_true, y_pred):
+        return tf.reduce_mean((1. - self.ms_ssim(y_true, y_pred))/2.)
+
+
+class DICEScore:
+    def __init__(self, input_shape: list):
+        """
+        DICE Score.
+        :param input_shape: Shape of the input image, without the batch dimension, e.g., 2D: [H, W, C], 3D: [H, W, D, C]
+        """
+        self.axes = list(range(1, len(input_shape)))  # The list will not include the channel axis [1, ..., num_dims)
+
+    def dice(self, y_true, y_pred):
+        numerator = 2 * tf.reduce_sum(y_true * y_pred, self.axes)
+        denominator = tf.reduce_sum(y_true + y_pred, self.axes)
+        return tf.reduce_mean(tf.div_no_nan(numerator, denominator))
+
+    @function_decorator('DICE__loss')
+    def loss(self, y_true, y_pred):
+        return 1 - 2 * tf.reduce_mean(self.dice(y_true, y_pred))
+
+    @function_decorator('DICE__metric')
+    def metric(self, y_true, y_pred):
+        return tf.reduce_mean(self.dice(y_true, y_pred))
+
+
+class GeneralizedDICEScore:
+    def __init__(self, input_shape: list, num_labels: int=None):
+        """
+        Generalized DICE Score. Implementation based on Carole H. Sudre, et al., "Generalised DIce Overlap as a Deep
+        Learning Los Function for Highly Unbalanced Segmentations" https://arxiv.org/abs/1707.03237
+        :param input_shape: Shape of the input image, without the batch dimension, e.g., 2D: [H, W, C], 3D: [H, W, D, C]
+        """
+        if input_shape[-1] > 1:
+            self.flat_shape = [-1, np.prod(np.asarray(input_shape[:-1])), input_shape[-1]]
+            self.hot_encode = False
+        elif num_labels is not None:
+            self.flat_shape = [-1, np.prod(np.asarray(input_shape[:-1])), num_labels]
+            self.one_hot_enc_shape = [-1, *input_shape[:-1]]
+            self.hot_encode = True
+            warnings.warn('Differentiable one-hot encoding not yet implemented')
+        else:
+            raise ValueError('If input_shape is not one hot encoded, then num_labels must be provided')
+
+    def one_hot_encoding(self, in_img, name=''):
+        # TODO: Test if differentiable!
+        labels, indices = tf.unique(tf.reshape(in_img, [-1]), tf.int32, name=name+'_unique')
+        one_hot = tf.one_hot(indices, tf.size(labels), name=name + '_one_hot')
+        one_hot = tf.reshape(one_hot, self.one_hot_enc_shape + [tf.size(labels)], name=name + '_reshape')
+        one_hot = tf.slice(one_hot, [0]*len(self.one_hot_enc_shape) + [1], [-1]*(len(self.one_hot_enc_shape) + 1),
+                           name=name + '_remove_bg')
+        return one_hot
+
+    def weigthed_dice(self, y_true, y_pred):
+        # y_true = [B, -1, L]
+        # y_pred = [B, -1, L]
+        if self.hot_encode:
+            y_true = self.one_hot_encoding(y_true, name='GDICE_one_hot_encoding_y_true')
+            y_pred = self.one_hot_encoding(y_pred, name='GDICE_one_hot_encoding_y_pred')
+        y_true = tf.reshape(y_true, self.flat_shape, name='GDICE_reshape_y_true')    # Flatten along the volume dimensions
+        y_pred = tf.reshape(y_pred, self.flat_shape, name='GDICE_reshape_y_pred')    # Flatten along the volume dimensions
+
+        size_y_true = tf.reduce_sum(y_true, axis=1, name='GDICE_size_y_true')
+        size_y_pred = tf.reduce_sum(y_pred, axis=1, name='GDICE_size_y_pred')
+        w = tf.div_no_nan(1., tf.pow(size_y_true, 2), name='GDICE_weight')
+        numerator = w * tf.reduce_sum(y_true * y_pred, axis=1)
+        denominator = w * (size_y_true + size_y_pred)
+        return tf.div_no_nan(2 * tf.reduce_sum(numerator, axis=-1), tf.reduce_sum(denominator, axis=-1))
+
+    def macro_dice(self, y_true, y_pred):
+        # y_true = [B, -1, L]
+        # y_pred = [B, -1, L]
+        if self.hot_encode:
+            y_true = self.one_hot_encoding(y_true, name='GDICE_one_hot_encoding_y_true')
+            y_pred = self.one_hot_encoding(y_pred, name='GDICE_one_hot_encoding_y_pred')
+        y_true = tf.reshape(y_true, self.flat_shape, name='GDICE_reshape_y_true')    # Flatten along the volume dimensions
+        y_pred = tf.reshape(y_pred, self.flat_shape, name='GDICE_reshape_y_pred')    # Flatten along the volume dimensions
+
+        size_y_true = tf.reduce_sum(y_true, axis=1, name='GDICE_size_y_true')
+        size_y_pred = tf.reduce_sum(y_pred, axis=1, name='GDICE_size_y_pred')
+        numerator = tf.reduce_sum(y_true * y_pred, axis=1)
+        denominator = (size_y_true + size_y_pred)
+        return tf.div_no_nan(2 * numerator, denominator)
+
+    @function_decorator('GeneralizeDICE__loss')
+    def loss(self, y_true, y_pred):
+        return 1 - tf.reduce_mean(self.weigthed_dice(y_true, y_pred))
+
+    @function_decorator('GeneralizeDICE__metric')
+    def metric(self, y_true, y_pred):
+        return tf.reduce_mean(self.weigthed_dice(y_true, y_pred))
+
+    @function_decorator('GeneralizeDICE__loss_macro')
+    def loss_macro(self, y_true, y_pred):
+        return 1 - tf.reduce_mean(self.macro_dice(y_true, y_pred))
+
+    @function_decorator('GeneralizeDICE__metric_macro')
+    def metric_macro(self, y_true, y_pred):
+        return tf.reduce_mean(self.macro_dice(y_true, y_pred))
+
+
+def target_registration_error(y_true, y_pred, average=True):
+    '''
+    Target Registration Error measured as the average distance between y_true and y_pred
+    :param y_true: [N, D] target points
+    :param y_pred: [N, D] predicted points
+    :param average: return the average TRE or an [N,] array
+    :return: averate TRE or [N,] array of TRE for each point
+    '''
+    assert y_true.shape == y_pred.shape, "y_true and y_pred must have the same shape"
+    if average:
+        return tf.reduce_mean(tf.linalg.norm(y_pred - y_true, axis=1))
+    else:
+        return tf.linalg.norm(y_pred - y_true, axis=1)
+
+# TODO: tensorflow-graphic has an implementation of Hausdorff ditance.
+#  However, this is not where it should and I can't find it
+# def HausdorffDistance_exact(y_true, y_pred, ohe=False, name='hd_exact'):
+#     if ohe:
+#         y_true = tf.transpose(y_true, [0, 4, 1, 2, 3])
+#         y_pred = tf.transpose(y_pred, [0, 4, 1, 2, 3])
+#     y_true_coords = tf.where(y_true)
+#     y_pred_coords = tf.where(y_pred)
+#
+#     return tfg_nn.loss.hausdorff_distance.evaluate(y_true_coords, y_pred_coords, name=name)

DeepDeformationMapRegistration/ms_ssim_tf.py

@@ -0,0 +1,642 @@
+# SRC: https://github.com/tensorflow/tensorflow/blob/a4dfb8d1a71385bd6d122e4f27f86dcebb96712d/tensorflow/python/ops/image_ops_impl.py
+from tensorflow.python import nn_ops
+from tensorflow.python import math_ops
+from tensorflow.python import array_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.util.tf_export import tf_export
+from tensorflow.python.util import dispatch
+from DeepDeformationMapRegistration.utils.misc import function_decorator
+
+
+@tf_export('image.convert_image_dtype')
+@dispatch.add_dispatch_support
+def convert_image_dtype(image, dtype, saturate=False, name=None):
+    """Convert `image` to `dtype`, scaling its values if needed.
+    The operation supports data types (for `image` and `dtype`) of
+    `uint8`, `uint16`, `uint32`, `uint64`, `int8`, `int16`, `int32`, `int64`,
+    `float16`, `float32`, `float64`, `bfloat16`.
+    Images that are represented using floating point values are expected to have
+    values in the range [0,1). Image data stored in integer data types are
+    expected to have values in the range `[0,MAX]`, where `MAX` is the largest
+    positive representable number for the data type.
+    This op converts between data types, scaling the values appropriately before
+    casting.
+    Usage Example:
+    >>> x = [[[1, 2, 3], [4, 5, 6]],
+    ...      [[7, 8, 9], [10, 11, 12]]]
+    >>> x_int8 = tf.convert_to_tensor(x, dtype=tf.int8)
+    >>> tf.image.convert_image_dtype(x_int8, dtype=tf.float16, saturate=False)
+    <tf.Tensor: shape=(2, 2, 3), dtype=float16, numpy=
+    array([[[0.00787, 0.01575, 0.02362],
+            [0.0315 , 0.03937, 0.04724]],
+           [[0.0551 , 0.063  , 0.07086],
+            [0.07874, 0.0866 , 0.0945 ]]], dtype=float16)>
+    Converting integer types to floating point types returns normalized floating
+    point values in the range [0, 1); the values are normalized by the `MAX` value
+    of the input dtype. Consider the following two examples:
+    >>> a = [[[1], [2]], [[3], [4]]]
+    >>> a_int8 = tf.convert_to_tensor(a, dtype=tf.int8)
+    >>> tf.image.convert_image_dtype(a_int8, dtype=tf.float32)
+    <tf.Tensor: shape=(2, 2, 1), dtype=float32, numpy=
+    array([[[0.00787402],
+            [0.01574803]],
+           [[0.02362205],
+            [0.03149606]]], dtype=float32)>
+    >>> a_int32 = tf.convert_to_tensor(a, dtype=tf.int32)
+    >>> tf.image.convert_image_dtype(a_int32, dtype=tf.float32)
+    <tf.Tensor: shape=(2, 2, 1), dtype=float32, numpy=
+    array([[[4.6566129e-10],
+            [9.3132257e-10]],
+           [[1.3969839e-09],
+            [1.8626451e-09]]], dtype=float32)>
+    Despite having identical values of `a` and output dtype of `float32`, the
+    outputs differ due to the different input dtypes (`int8` vs. `int32`). This
+    is, again, because the values are normalized by the `MAX` value of the input
+    dtype.
+    Note that converting floating point values to integer type may lose precision.
+    In the example below, an image tensor `b` of dtype `float32` is converted to
+    `int8` and back to `float32`. The final output, however, is different from
+    the original input `b` due to precision loss.
+    >>> b = [[[0.12], [0.34]], [[0.56], [0.78]]]
+    >>> b_float32 = tf.convert_to_tensor(b, dtype=tf.float32)
+    >>> b_int8 = tf.image.convert_image_dtype(b_float32, dtype=tf.int8)
+    >>> tf.image.convert_image_dtype(b_int8, dtype=tf.float32)
+    <tf.Tensor: shape=(2, 2, 1), dtype=float32, numpy=
+    array([[[0.11811024],
+            [0.33858266]],
+           [[0.5590551 ],
+            [0.77952754]]], dtype=float32)>
+    Scaling up from an integer type (input dtype) to another integer type (output
+    dtype) will not map input dtype's `MAX` to output dtype's `MAX` but converting
+    back and forth should result in no change. For example, as shown below, the
+    `MAX` value of int8 (=127) is not mapped to the `MAX` value of int16 (=32,767)
+    but, when scaled back, we get the same, original values of `c`.
+    >>> c = [[[1], [2]], [[127], [127]]]
+    >>> c_int8 = tf.convert_to_tensor(c, dtype=tf.int8)
+    >>> c_int16 = tf.image.convert_image_dtype(c_int8, dtype=tf.int16)
+    >>> print(c_int16)
+    tf.Tensor(
+    [[[  256]
+      [  512]]
+     [[32512]
+      [32512]]], shape=(2, 2, 1), dtype=int16)
+    >>> c_int8_back = tf.image.convert_image_dtype(c_int16, dtype=tf.int8)
+    >>> print(c_int8_back)
+    tf.Tensor(
+    [[[  1]
+      [  2]]
+     [[127]
+      [127]]], shape=(2, 2, 1), dtype=int8)
+    Scaling down from an integer type to another integer type can be a lossy
+    conversion. Notice in the example below that converting `int16` to `uint8` and
+    back to `int16` has lost precision.
+    >>> d = [[[1000], [2000]], [[3000], [4000]]]
+    >>> d_int16 = tf.convert_to_tensor(d, dtype=tf.int16)
+    >>> d_uint8 = tf.image.convert_image_dtype(d_int16, dtype=tf.uint8)
+    >>> d_int16_back = tf.image.convert_image_dtype(d_uint8, dtype=tf.int16)
+    >>> print(d_int16_back)
+    tf.Tensor(
+    [[[ 896]
+      [1920]]
+     [[2944]
+      [3968]]], shape=(2, 2, 1), dtype=int16)
+    Note that converting from floating point inputs to integer types may lead to
+    over/underflow problems. Set saturate to `True` to avoid such problem in
+    problematic conversions. If enabled, saturation will clip the output into the
+    allowed range before performing a potentially dangerous cast (and only before
+    performing such a cast, i.e., when casting from a floating point to an integer
+    type, and when casting from a signed to an unsigned type; `saturate` has no
+    effect on casts between floats, or on casts that increase the type's range).
+    Args:
+      image: An image.
+      dtype: A `DType` to convert `image` to.
+      saturate: If `True`, clip the input before casting (if necessary).
+      name: A name for this operation (optional).
+    Returns:
+      `image`, converted to `dtype`.
+    Raises:
+      AttributeError: Raises an attribute error when dtype is neither
+      float nor integer
+    """
+    image = ops.convert_to_tensor(image, name='image')
+    dtype = dtypes.as_dtype(dtype)
+    if not dtype.is_floating and not dtype.is_integer:
+        raise AttributeError('dtype must be either floating point or integer')
+    if dtype == image.dtype:
+        return array_ops.identity(image, name=name)
+
+    with ops.name_scope(name, 'convert_image', [image]) as name:
+        # Both integer: use integer multiplication in the larger range
+        if image.dtype.is_integer and dtype.is_integer:
+            scale_in = image.dtype.max
+            scale_out = dtype.max
+            if scale_in > scale_out:
+                # Scaling down, scale first, then cast. The scaling factor will
+                # cause in.max to be mapped to above out.max but below out.max+1,
+                # so that the output is safely in the supported range.
+                scale = (scale_in + 1) // (scale_out + 1)
+                scaled = math_ops.floordiv(image, scale)
+
+                if saturate:
+                    return math_ops.saturate_cast(scaled, dtype, name=name)
+                else:
+                    return math_ops.cast(scaled, dtype, name=name)
+            else:
+                # Scaling up, cast first, then scale. The scale will not map in.max to
+                # out.max, but converting back and forth should result in no change.
+                if saturate:
+                    cast = math_ops.saturate_cast(image, dtype)
+                else:
+                    cast = math_ops.cast(image, dtype)
+                scale = (scale_out + 1) // (scale_in + 1)
+                return math_ops.multiply(cast, scale, name=name)
+        elif image.dtype.is_floating and dtype.is_floating:
+            # Both float: Just cast, no possible overflows in the allowed ranges.
+            # Note: We're ignoring float overflows. If your image dynamic range
+            # exceeds float range, you're on your own.
+            return math_ops.cast(image, dtype, name=name)
+        else:
+            if image.dtype.is_integer:
+                # Converting to float: first cast, then scale. No saturation possible.
+                cast = math_ops.cast(image, dtype)
+                scale = 1. / image.dtype.max
+                return math_ops.multiply(cast, scale, name=name)
+            else:
+                # Converting from float: first scale, then cast
+                scale = dtype.max + 0.5  # avoid rounding problems in the cast
+                scaled = math_ops.multiply(image, scale)
+                if saturate:
+                    return math_ops.saturate_cast(scaled, dtype, name=name)
+                else:
+                    return math_ops.cast(scaled, dtype, name=name)
+
+
+def _verify_compatible_image_shapes(img1, img2):
+    """Checks if two image tensors are compatible for applying SSIM or PSNR.
+    This function checks if two sets of images have ranks at least 3, and if the
+    last three dimensions match.
+    Args:
+      img1: Tensor containing the first image batch.
+      img2: Tensor containing the second image batch.
+    Returns:
+      A tuple containing: the first tensor shape, the second tensor shape, and a
+      list of control_flow_ops.Assert() ops implementing the checks.
+    Raises:
+      ValueError: When static shape check fails.
+    """
+    shape1 = img1.get_shape().with_rank_at_least(4)     # at least [H, W, D, C]
+    shape2 = img2.get_shape().with_rank_at_least(4)     # at least [H, W, D, C]
+    shape1[-4:].assert_is_compatible_with(shape2[-4:])
+
+    if shape1.ndims is not None and shape2.ndims is not None:
+        for dim1, dim2 in zip(
+                reversed(shape1.dims[:-4]), reversed(shape2.dims[:-4])):
+            if not (dim1 == 1 or dim2 == 1 or dim1.is_compatible_with(dim2)):
+                raise ValueError('Two images are not compatible: %s and %s' %
+                                 (shape1, shape2))
+
+    # Now assign shape tensors.
+    shape1, shape2 = array_ops.shape_n([img1, img2])
+
+    # TODO(sjhwang): Check if shape1[:-4] and shape2[:-4] are broadcastable.
+    checks = []
+    checks.append(
+        control_flow_ops.Assert(
+            math_ops.greater_equal(array_ops.size(shape1), 4), [shape1, shape2],
+            summarize=10))
+    checks.append(
+        control_flow_ops.Assert(
+            math_ops.reduce_all(math_ops.equal(shape1[-4:], shape2[-4:])),
+            [shape1, shape2],
+            summarize=10))
+    return shape1, shape2, checks
+
+
+def _ssim_helper(x, y, reducer, max_val, compensation=1.0, k1=0.01, k2=0.03):
+    r"""Helper function for computing SSIM.
+    SSIM estimates covariances with weighted sums.  The default parameters
+    use a biased estimate of the covariance:
+    Suppose `reducer` is a weighted sum, then the mean estimators are
+      \mu_x = \sum_i w_i x_i,
+      \mu_y = \sum_i w_i y_i,
+    where w_i's are the weighted-sum weights, and covariance estimator is
+      cov_{xy} = \sum_i w_i (x_i - \mu_x) (y_i - \mu_y)
+    with assumption \sum_i w_i = 1. This covariance estimator is biased, since
+      E[cov_{xy}] = (1 - \sum_i w_i ^ 2) Cov(X, Y).
+    For SSIM measure with unbiased covariance estimators, pass as `compensation`
+    argument (1 - \sum_i w_i ^ 2).
+    Args:
+      x: First set of images.
+      y: Second set of images.
+      reducer: Function that computes 'local' averages from the set of images. For
+        non-convolutional version, this is usually tf.reduce_mean(x, [1, 2]), and
+        for convolutional version, this is usually tf.nn.avg_pool2d or
+        tf.nn.conv3d with weighted-sum kernel.
+      max_val: The dynamic range (i.e., the difference between the maximum
+        possible allowed value and the minimum allowed value).
+      compensation: Compensation factor. See above.
+      k1: Default value 0.01
+      k2: Default value 0.03 (SSIM is less sensitivity to K2 for lower values, so
+        it would be better if we took the values in the range of 0 < K2 < 0.4).
+    Returns:
+      A pair containing the luminance measure, and the contrast-structure measure.
+    """
+
+    c1 = (k1 * max_val)**2
+    c2 = (k2 * max_val)**2
+
+    # SSIM luminance measure is
+    # (2 * mu_x * mu_y + c1) / (mu_x ** 2 + mu_y ** 2 + c1).
+    mean0 = reducer(x)
+    mean1 = reducer(y)
+    num0 = mean0 * mean1 * 2.0
+    den0 = math_ops.square(mean0) + math_ops.square(mean1)
+    luminance = (num0 + c1) / (den0 + c1)
+
+    # SSIM contrast-structure measure is
+    #   (2 * cov_{xy} + c2) / (cov_{xx} + cov_{yy} + c2).
+    # Note that `reducer` is a weighted sum with weight w_k, \sum_i w_i = 1, then
+    #   cov_{xy} = \sum_i w_i (x_i - \mu_x) (y_i - \mu_y)
+    #          = \sum_i w_i x_i y_i - (\sum_i w_i x_i) (\sum_j w_j y_j).
+    num1 = reducer(x * y) * 2.0
+    den1 = reducer(math_ops.square(x) + math_ops.square(y))
+    c2 *= compensation
+    cs = (num1 - num0 + c2) / (den1 - den0 + c2)
+
+    # SSIM score is the product of the luminance and contrast-structure measures.
+    return luminance, cs
+
+
+def _fspecial_gauss(size, sigma):
+    """Function to mimic the 'fspecial' gaussian MATLAB function."""
+    size = ops.convert_to_tensor(size, dtypes.int32)
+    sigma = ops.convert_to_tensor(sigma, dtypes.float32)
+
+    coords = math_ops.cast(math_ops.range(size), sigma.dtype)
+    coords -= math_ops.cast(size - 1, sigma.dtype) / 2.0
+
+    g = math_ops.square(coords)
+    g *= -0.5 / math_ops.square(sigma)
+
+    g = array_ops.reshape(g, shape=[1, -1]) + array_ops.reshape(g, shape=[-1, 1])
+    g = array_ops.reshape(g, shape=[size, size, 1]) + array_ops.reshape(g, shape=[1, size, size])
+    g = array_ops.reshape(g, shape=[1, -1])  # For tf.nn.softmax().
+    g = nn_ops.softmax(g)
+    return array_ops.reshape(g, shape=[size, size, size, 1, 1])
+
+
+def _ssim_per_channel(img1,
+                      img2,
+                      max_val=1.0,
+                      filter_size=11,
+                      filter_sigma=1.5,
+                      k1=0.01,
+                      k2=0.03):
+    """Computes SSIM index between img1 and img2 per color channel.
+    This function matches the standard SSIM implementation from:
+    Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). Image
+    quality assessment: from error visibility to structural similarity. IEEE
+    transactions on image processing.
+    Details:
+      - 11x11 Gaussian filter of width 1.5 is used.
+      - k1 = 0.01, k2 = 0.03 as in the original paper.
+    Args:
+      img1: First image batch.
+      img2: Second image batch.
+      max_val: The dynamic range of the images (i.e., the difference between the
+        maximum the and minimum allowed values).
+      filter_size: Default value 11 (size of gaussian filter).
+      filter_sigma: Default value 1.5 (width of gaussian filter).
+      k1: Default value 0.01
+      k2: Default value 0.03 (SSIM is less sensitivity to K2 for lower values, so
+        it would be better if we took the values in the range of 0 < K2 < 0.4).
+    Returns:
+      A pair of tensors containing and channel-wise SSIM and contrast-structure
+      values. The shape is [..., channels].
+    """
+    filter_size = constant_op.constant(filter_size, dtype=dtypes.int32)
+    filter_sigma = constant_op.constant(filter_sigma, dtype=img1.dtype)
+
+    shape1, shape2 = array_ops.shape_n([img1, img2])
+    checks = [
+        control_flow_ops.Assert(
+            math_ops.reduce_all(
+                math_ops.greater_equal(shape1[-4:-1], filter_size)),
+            [shape1, filter_size],
+            summarize=8),
+        control_flow_ops.Assert(
+            math_ops.reduce_all(
+                math_ops.greater_equal(shape2[-4:-1], filter_size)),
+            [shape2, filter_size],
+            summarize=8)
+    ]
+
+    # Enforce the check to run before computation.
+    with ops.control_dependencies(checks):
+        img1 = array_ops.identity(img1)
+
+    # TODO(sjhwang): Try to cache kernels and compensation factor.
+    kernel = _fspecial_gauss(filter_size, filter_sigma)
+    kernel = array_ops.tile(kernel, multiples=[1, 1, 1, shape1[-1], 1])
+
+    # The correct compensation factor is `1.0 - tf.reduce_sum(tf.square(kernel))`,
+    # but to match MATLAB implementation of MS-SSIM, we use 1.0 instead.
+    compensation = 1.0
+
+    # TODO(sjhwang): Try FFT.
+    # TODO(sjhwang): Gaussian kernel is separable in space. Consider applying
+    #   1-by-n and n-by-1 Gaussian filters instead of an n-by-n filter.
+    def reducer(x):
+        shape = array_ops.shape(x)
+        x = array_ops.reshape(x, shape=array_ops.concat([[-1], shape[-4:]], 0))
+        y = nn.conv3d(x, kernel, strides=[1, 1, 1, 1, 1], padding='VALID')
+        return array_ops.reshape(y, array_ops.concat([shape[:-4], array_ops.shape(y)[1:]], 0))
+
+    luminance, cs = _ssim_helper(img1, img2, reducer, max_val, compensation, k1,
+                                 k2)
+
+    # Average over the second, third and the fourth from the last: height, width, depth.
+    axes = constant_op.constant([-4, -3, -2], dtype=dtypes.int32)
+    ssim_val = math_ops.reduce_mean(luminance * cs, axes)
+    cs = math_ops.reduce_mean(cs, axes)
+    return ssim_val, cs
+
+
+@tf_export('image.ssim')
+@dispatch.add_dispatch_support
+def ssim(img1,
+         img2,
+         max_val,
+         filter_size=11,
+         filter_sigma=1.5,
+         k1=0.01,
+         k2=0.03):
+    """Computes SSIM index between img1 and img2.
+    This function is based on the standard SSIM implementation from:
+    Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). Image
+    quality assessment: from error visibility to structural similarity. IEEE
+    transactions on image processing.
+    Note: The true SSIM is only defined on grayscale.  This function does not
+    perform any colorspace transform.  (If the input is already YUV, then it will
+    compute YUV SSIM average.)
+    Details:
+      - 11x11 Gaussian filter of width 1.5 is used.
+      - k1 = 0.01, k2 = 0.03 as in the original paper.
+    The image sizes must be at least 11x11 because of the filter size.
+    Example:
+    ```python
+        # Read images (of size 255 x 255) from file.
+        im1 = tf.image.decode_image(tf.io.read_file('path/to/im1.png'))
+        im2 = tf.image.decode_image(tf.io.read_file('path/to/im2.png'))
+        tf.shape(im1)  # `img1.png` has 3 channels; shape is `(255, 255, 3)`
+        tf.shape(im2)  # `img2.png` has 3 channels; shape is `(255, 255, 3)`
+        # Add an outer batch for each image.
+        im1 = tf.expand_dims(im1, axis=0)
+        im2 = tf.expand_dims(im2, axis=0)
+        # Compute SSIM over tf.uint8 Tensors.
+        ssim1 = tf.image.ssim(im1, im2, max_val=255, filter_size=11,
+                              filter_sigma=1.5, k1=0.01, k2=0.03)
+        # Compute SSIM over tf.float32 Tensors.
+        im1 = tf.image.convert_image_dtype(im1, tf.float32)
+        im2 = tf.image.convert_image_dtype(im2, tf.float32)
+        ssim2 = tf.image.ssim(im1, im2, max_val=1.0, filter_size=11,
+                              filter_sigma=1.5, k1=0.01, k2=0.03)
+        # ssim1 and ssim2 both have type tf.float32 and are almost equal.
+    ```
+    Args:
+      img1: First image batch. 4-D Tensor of shape `[batch, height, width,
+        channels]` with only Positive Pixel Values.
+      img2: Second image batch. 4-D Tensor of shape `[batch, height, width,
+        channels]` with only Positive Pixel Values.
+      max_val: The dynamic range of the images (i.e., the difference between the
+        maximum the and minimum allowed values).
+      filter_size: Default value 11 (size of gaussian filter).
+      filter_sigma: Default value 1.5 (width of gaussian filter).
+      k1: Default value 0.01
+      k2: Default value 0.03 (SSIM is less sensitivity to K2 for lower values, so
+        it would be better if we took the values in the range of 0 < K2 < 0.4).
+    Returns:
+      A tensor containing an SSIM value for each image in batch.  Returned SSIM
+      values are in range (-1, 1], when pixel values are non-negative. Returns
+      a tensor with shape: broadcast(img1.shape[:-3], img2.shape[:-3]).
+    """
+    with ops.name_scope(None, 'SSIM', [img1, img2]):
+        # Convert to tensor if needed.
+        img1 = ops.convert_to_tensor(img1, name='img1')
+        img2 = ops.convert_to_tensor(img2, name='img2')
+        # Shape checking.
+        _, _, checks = _verify_compatible_image_shapes(img1, img2)
+        with ops.control_dependencies(checks):
+            img1 = array_ops.identity(img1)
+
+        # Need to convert the images to float32.  Scale max_val accordingly so that
+        # SSIM is computed correctly.
+        max_val = math_ops.cast(max_val, img1.dtype)
+        max_val = convert_image_dtype(max_val, dtypes.float32)
+        img1 = convert_image_dtype(img1, dtypes.float32)
+        img2 = convert_image_dtype(img2, dtypes.float32)
+        ssim_per_channel, _ = _ssim_per_channel(img1, img2, max_val, filter_size,
+                                                filter_sigma, k1, k2)
+        # Compute average over color channels.
+        return math_ops.reduce_mean(ssim_per_channel, [-1])
+
+
+# Default values obtained by Wang et al.
+_MSSSIM_WEIGHTS = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
+
+
+@tf_export('image.ssim_multiscale')
+@dispatch.add_dispatch_support
+def ssim_multiscale(img1,
+                    img2,
+                    max_val,
+                    power_factors=_MSSSIM_WEIGHTS,
+                    filter_size=11,
+                    filter_sigma=1.5,
+                    k1=0.01,
+                    k2=0.03):
+    """Computes the MS-SSIM between img1 and img2.
+    This function assumes that `img1` and `img2` are image batches, i.e. the last
+    three dimensions are [height, width, channels].
+    Note: The true SSIM is only defined on grayscale.  This function does not
+    perform any colorspace transform.  (If the input is already YUV, then it will
+    compute YUV SSIM average.)
+    Original paper: Wang, Zhou, Eero P. Simoncelli, and Alan C. Bovik. "Multiscale
+    structural similarity for image quality assessment." Signals, Systems and
+    Computers, 2004.
+    Args:
+      img1: First image batch with only Positive Pixel Values.
+      img2: Second image batch with only Positive Pixel Values. Must have the
+      same rank as img1.
+      max_val: The dynamic range of the images (i.e., the difference between the
+        maximum the and minimum allowed values).
+      power_factors: Iterable of weights for each of the scales. The number of
+        scales used is the length of the list. Index 0 is the unscaled
+        resolution's weight and each increasing scale corresponds to the image
+        being downsampled by 2.  Defaults to (0.0448, 0.2856, 0.3001, 0.2363,
+        0.1333), which are the values obtained in the original paper.
+      filter_size: Default value 11 (size of gaussian filter).
+      filter_sigma: Default value 1.5 (width of gaussian filter).
+      k1: Default value 0.01
+      k2: Default value 0.03 (SSIM is less sensitivity to K2 for lower values, so
+        it would be better if we took the values in the range of 0 < K2 < 0.4).
+    Returns:
+      A tensor containing an MS-SSIM value for each image in batch.  The values
+      are in range [0, 1].  Returns a tensor with shape:
+      broadcast(img1.shape[:-3], img2.shape[:-3]).
+    """
+    with ops.name_scope(None, 'MS-SSIM', [img1, img2]):
+        # Convert to tensor if needed.
+        img1 = ops.convert_to_tensor(img1, name='img1')
+        img2 = ops.convert_to_tensor(img2, name='img2')
+        # Shape checking.
+        shape1, shape2, checks = _verify_compatible_image_shapes(img1, img2)
+        with ops.control_dependencies(checks):
+            img1 = array_ops.identity(img1)
+
+        # Need to convert the images to float32.  Scale max_val accordingly so that
+        # SSIM is computed correctly.
+        max_val = math_ops.cast(max_val, img1.dtype)
+        max_val = convert_image_dtype(max_val, dtypes.float32)
+        img1 = convert_image_dtype(img1, dtypes.float32)
+        img2 = convert_image_dtype(img2, dtypes.float32)
+
+        imgs = [img1, img2]
+        shapes = [shape1, shape2]
+
+        # img1 and img2 are assumed to be a (multi-dimensional) batch of
+        # 4-dimensional images (height, width, depth, channels). `heads` contain the batch
+        # dimensions, and `tails` contain the image dimensions.
+        heads = [s[:-4] for s in shapes]
+        tails = [s[-4:] for s in shapes]
+
+        divisor = [1, 2, 2, 2, 1]
+        divisor_tensor = constant_op.constant(divisor[1:], dtype=dtypes.int32)
+
+        def do_pad(images, remainder):
+            padding = array_ops.expand_dims(remainder, -1)
+            padding = array_ops.pad(padding, [[1, 0], [1, 0]])
+            return [array_ops.pad(x, padding, mode='SYMMETRIC') for x in images]
+
+        mcs = []
+        for k in range(len(power_factors)):
+            with ops.name_scope(None, 'Scale%d' % k, imgs):
+                if k > 0:
+                    # Avg pool takes rank 4 tensors. Flatten leading dimensions.
+                    flat_imgs = [
+                        array_ops.reshape(x, array_ops.concat([[-1], t], 0))
+                        for x, t in zip(imgs, tails)
+                    ]
+
+                    remainder = tails[0] % divisor_tensor
+                    need_padding = math_ops.reduce_any(math_ops.not_equal(remainder, 0))
+                    # pylint: disable=cell-var-from-loop
+                    padded = control_flow_ops.cond(need_padding,
+                                                   lambda: do_pad(flat_imgs, remainder),
+                                                   lambda: flat_imgs)
+                    # pylint: enable=cell-var-from-loop
+
+                    downscaled = [
+                        nn_ops.avg_pool3d(
+                            x, ksize=divisor, strides=divisor, padding='VALID', data_format='NDHWC',)
+                        for x in padded
+                    ]
+                    tails = [x[1:] for x in array_ops.shape_n(downscaled)]
+                    imgs = [
+                        array_ops.reshape(x, array_ops.concat([h, t], 0))
+                        for x, h, t in zip(downscaled, heads, tails)
+                    ]
+
+                # Overwrite previous ssim value since we only need the last one.
+                ssim_per_channel, cs = _ssim_per_channel(
+                    *imgs,
+                    max_val=max_val,
+                    filter_size=filter_size,
+                    filter_sigma=filter_sigma,
+                    k1=k1,
+                    k2=k2)
+                mcs.append(nn_ops.relu(cs))
+
+        # Remove the cs score for the last scale. In the MS-SSIM calculation,
+        # we use the l(p) at the highest scale. l(p) * cs(p) is ssim(p).
+        mcs.pop()  # Remove the cs score for the last scale.
+        mcs_and_ssim = array_ops.stack(
+            mcs + [nn_ops.relu(ssim_per_channel)], axis=-1)
+        # Take weighted geometric mean across the scale axis.
+        ms_ssim = math_ops.reduce_prod(
+            math_ops.pow(mcs_and_ssim, power_factors), [-1])
+
+        return math_ops.reduce_mean(ms_ssim, [-1])  # Avg over color channels.
+
+
+class MultiScaleStructuralSimilarity:
+    def __init__(self,
+                 max_val,
+                 power_factors=_MSSSIM_WEIGHTS,
+                 filter_size=11,
+                 filter_sigma=1.5,
+                 k1=0.01,
+                 k2=0.03):
+        self.max_val = max_val
+        self.power_factors = power_factors
+        self.filter_size = int(filter_size)
+        self.filter_sigma = filter_sigma
+        self.k1 = k1
+        self.k2 = k2
+
+    @function_decorator('MS_SSIM__loss')
+    def loss(self, y_true, y_pred):
+        return math_ops.reduce_mean((1 - ssim_multiscale(y_true, y_pred, self.max_val, self.power_factors,
+                                                         self.filter_size, self.filter_sigma, self.k1, self.k2))/2)
+
+    @function_decorator('MS_SSIM__metric')
+    def metric(self, y_true, y_pred):
+        return ssim_multiscale(y_true, y_pred, self.max_val, self.power_factors, self.filter_size, self.filter_sigma,
+                               self.k1, self.k2)
+
+
+if __name__ == '__main__':
+    import os
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    import tensorflow as tf
+    tf.enable_eager_execution()
+    import nibabel as nib
+    import numpy as np
+    from DeepDeformationMapRegistration.utils.operators import min_max_norm
+    from skimage.metrics import structural_similarity
+
+    img1 = nib.load('test_images/ixi_image.nii.gz')
+    img1 = np.asarray(img1.dataobj)
+    img1 = img1[np.newaxis, ..., np.newaxis]    # Add Batch and Channel dimensions
+
+    img2 = nib.load('test_images/ixi_image2.nii.gz')
+    img2 = np.asarray(img2.dataobj)
+    img2 = img2[np.newaxis, ..., np.newaxis]
+
+    img1 = min_max_norm(img1)
+    img2 = min_max_norm(img2)
+
+    ssim_tf_1_2 = ssim(img1, img2, 1., filter_size=5)
+    assert ssim(img1, img1, 1., filter_size=5).numpy()[0] == 1., 'TF SSIM returned an unexpected value'
+    ssim_sklearn = structural_similarity(img1[0, ..., 0], img2[0, ..., 0], win_size=5)
+
+    ms_ssim_tf_1_2 = ssim_multiscale(img1, img2, 1., filter_size=5)
+    assert ssim_multiscale(img1, img1, 1., filter_size=5).numpy()[0] == 1., 'TF MS-SSIM returned an unexpected value'
+
+    print('SSIM TF: {}\nSSIM SKLEARN: {}\nMS SSIM TF: {}\n'.format(ssim_tf_1_2, ssim_sklearn, ms_ssim_tf_1_2))
+
+    batch_img1 = np.stack([img1, img2], axis=0)
+    batch_img2 = np.stack([img2, img2], axis=0)
+    batch_ssim_tf = ssim(batch_img1, batch_img2, 1., filter_size=5)
+    batch_ms_ssim_tf = ssim_multiscale(batch_img1, batch_img2, 1., filter_size=5)
+
+    print('Batch SSIM TF: {}\nBatch MS SSIM TF: {}\n'.format(batch_ssim_tf, batch_ms_ssim_tf))
+
+    img1 = img1[:, :127, :127, :127, :]
+    img2 = img2[:, :127, :127, :127, :]
+    MS_SSIM = MultiScaleStructuralSimilarity(1., filter_size=5)
+    print('MS SSIM Loss{}'.format(MS_SSIM.loss(img1, img2)))

DeepDeformationMapRegistration/networks.py

@@ -8,12 +8,14 @@ PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
 import tensorflow as tf
 import voxelmorph as vxm
 from voxelmorph.tf.modelio import LoadableModel, store_config_args
+from tensorflow.keras.layers import UpSampling3D
 
 
 class WeaklySupervised(LoadableModel):
 
     @store_config_args
-    def __init__(self, inshape, all_labels: [list, tuple], nb_unet_features=None, int_steps=5, bidir=False, **kwargs):
+    def __init__(self, inshape, all_labels: [list, tuple], nb_unet_features=None, int_steps=5, bidir=False,
+                 int_downsize=1, outshape=None, **kwargs):
         """
         Parameters:
             inshape: Input shape. e.g. (192, 192, 192)
@@ -21,34 +23,53 @@ class WeaklySupervised(LoadableModel):
             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.
+            int_downsize: Dowsampling of the displacement map. Integer
             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')
 
+        fix_img = tf.keras.Input((*inshape, 1), name='fix_image_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])
+        input_model = tf.keras.Model(inputs=[mov_img, fix_img], outputs=[mov_img, fix_img])
 
         vxm_model = vxm.networks.VxmDense(inshape=inshape,
                                           nb_unet_features=nb_unet_features,
-                                          input_model=unet_input_model,
+                                          input_model=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
+                                          bidir=bidir,
+                                          int_downsize=int_downsize,
+                                          **kwargs)
+
+        pred_segm = vxm.layers.SpatialTransformer(interp_method='linear', indexing='ij', name='interp_segm')(
+            [mov_segm, vxm_model.references.pos_flow])
+
+        inputs = [mov_img, fix_img, mov_segm] # mov_img, mov_segm, fix_segm
+        model_outputs = vxm_model.outputs
+        if outshape is not None:
+            scale_factors = [o//i for i, o in zip(inshape, outshape)]
+            upsampling_layer = UpSampling3D(scale_factors)  # Doesn't perform trilinear, only nearest
+            # Image
+            model_outputs[0] = upsampling_layer(model_outputs[0])
+            # Segmentation
+            pred_segm = upsampling_layer(pred_segm)
+            # Displacement map
+            model_outputs[1] = upsampling_layer(scale_factors)(model_outputs[1])
+            model_outputs[1] = tf.multiply(model_outputs[1], tf.cast(scale_factors, model_outputs[1].dtype))
+
+        # Just renaming
+        pred_fix_image = tf.identity(model_outputs[0], name='pred_fix_image')
+        pred_dm = tf.identity(model_outputs[1], name='pred_dm')
+        pred_segm = tf.identity(pred_segm, name='pred_fix_segm')
+        outputs = [pred_fix_image, pred_segm, pred_dm]
 
         self.references = LoadableModel.ReferenceContainer()
-        self.references.pred_segm = vxm_model.outputs[0]
-        self.references.pred_img = pred_img
+        self.references.pred_segm = pred_segm
+        self.references.pred_img = vxm_model.outputs[0]
         self.references.pos_flow = vxm_model.references.pos_flow
 
-        super().__init__(inputs=inputs, outputs=outputs)
+        super(WeaklySupervised, self).__init__(inputs=inputs, outputs=outputs)
 
     def get_registration_model(self):
         return tf.keras.Model(self.inputs, self.references.pos_flow)

DeepDeformationMapRegistration/utils/acummulated_optimizer.py

@@ -1,5 +1,9 @@
-from tensorflow.keras.optimizers import Optimizer
-from tensorflow.keras import backend as K
+from tensorflow.python.keras.optimizers import Optimizer
+from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2
+from tensorflow.python import ops, math_ops, state_ops, control_flow_ops
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras import backend_config
+import tensorflow as tf
 
 
 class AccumOptimizer(Optimizer):
@@ -14,7 +18,7 @@ class AccumOptimizer(Optimizer):
         a new keras optimizer.
     """
     def __init__(self, optimizer, steps_per_update=1, **kwargs):
-        super(AccumOptimizer, self).__init__(**kwargs)
+        super(AccumOptimizer, self).__init__(name='AccumOptimizer', **kwargs)
         self.optimizer = optimizer
         with K.name_scope(self.__class__.__name__):
             self.steps_per_update = steps_per_update
@@ -55,3 +59,134 @@ class AccumOptimizer(Optimizer):
         config = self.optimizer.get_config()
         K.set_value(self.iterations, iterations)
         return config
+
+
+__all__ = ['AdamAccumulated']
+
+
+# SRC: https://github.com/CyberZHG/keras-gradient-accumulation/blob/master/keras_gradient_accumulation/optimizer_v2.py
+class AdamAccumulated(OptimizerV2):
+    """Optimizer that implements the Adam algorithm with gradient accumulation."""
+
+    def __init__(self,
+                 accumulation_steps,
+                 learning_rate=0.001,
+                 beta_1=0.9,
+                 beta_2=0.999,
+                 epsilon=1e-7,
+                 amsgrad=False,
+                 name='Adam',
+                 **kwargs):
+        r"""Construct a new Adam optimizer.
+        Args:
+            accumulation_steps: An integer. Update gradient in every accumulation steps.
+            learning_rate: A Tensor or a floating point value.    The learning rate.
+            beta_1: A float value or a constant float tensor. The exponential decay
+                rate for the 1st moment estimates.
+            beta_2: A float value or a constant float tensor. The exponential decay
+                rate for the 2nd moment estimates.
+            epsilon: A small constant for numerical stability. This epsilon is
+                "epsilon hat" in the Kingma and Ba paper (in the formula just before
+                Section 2.1), not the epsilon in Algorithm 1 of the paper.
+            amsgrad: boolean. Whether to apply AMSGrad variant of this algorithm from
+                the paper "On the Convergence of Adam and beyond".
+            name: Optional name for the operations created when applying gradients.
+                Defaults to "Adam".    @compatibility(eager) When eager execution is
+                enabled, `learning_rate`, `beta_1`, `beta_2`, and `epsilon` can each be
+                a callable that takes no arguments and returns the actual value to use.
+                This can be useful for changing these values across different
+                invocations of optimizer functions. @end_compatibility
+            **kwargs: keyword arguments. Allowed to be {`clipnorm`, `clipvalue`, `lr`,
+                `decay`}. `clipnorm` is clip gradients by norm; `clipvalue` is clip
+                gradients by value, `decay` is included for backward compatibility to
+                allow time inverse decay of learning rate. `lr` is included for backward
+                compatibility, recommended to use `learning_rate` instead.
+        """
+
+        super(AdamAccumulated, self).__init__(name, **kwargs)
+        self._set_hyper('accumulation_steps', accumulation_steps)
+        self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
+        self._set_hyper('decay', self._initial_decay)
+        self._set_hyper('beta_1', beta_1)
+        self._set_hyper('beta_2', beta_2)
+        self.epsilon = epsilon or backend_config.epsilon()
+        self.amsgrad = amsgrad
+
+    def _create_slots(self, var_list):
+        for var in var_list:
+            self.add_slot(var, 'g')
+        for var in var_list:
+            self.add_slot(var, 'm')
+        for var in var_list:
+            self.add_slot(var, 'v')
+        if self.amsgrad:
+            for var in var_list:
+                self.add_slot(var, 'vhat')
+
+    def set_weights(self, weights):
+        params = self.weights
+        num_vars = int((len(params) - 1) / 2)
+        if len(weights) == 3 * num_vars + 1:
+            weights = weights[:len(params)]
+        super(AdamAccumulated, self).set_weights(weights)
+
+    def _resource_apply_dense(self, grad, var):
+        var_dtype = var.dtype.base_dtype
+        lr_t = self._decayed_lr(var_dtype)
+        beta_1_t = self._get_hyper('beta_1', var_dtype)
+        beta_2_t = self._get_hyper('beta_2', var_dtype)
+        accumulation_steps = self._get_hyper('accumulation_steps', 'int64')
+        update_cond = tf.equal((self.iterations + 1) % accumulation_steps, 0)
+        sub_step = self.iterations % accumulation_steps + 1
+        local_step = math_ops.cast(self.iterations // accumulation_steps + 1, var_dtype)
+        beta_1_power = math_ops.pow(beta_1_t, local_step)
+        beta_2_power = math_ops.pow(beta_2_t, local_step)
+        epsilon_t = ops.convert_to_tensor(self.epsilon, var_dtype)
+        lr = (lr_t * math_ops.sqrt(1 - beta_2_power) / (1 - beta_1_power))
+        lr = tf.where(update_cond, lr, 0.0)
+
+        g = self.get_slot(var, 'g')
+        g_a = grad / math_ops.cast(accumulation_steps, var_dtype)
+        g_t = tf.where(tf.equal(sub_step, 1),
+                       g_a,
+                       g + (g_a - g) / math_ops.cast(sub_step, var_dtype))
+        g_t = state_ops.assign(g, g_t, use_locking=self._use_locking)
+
+        m = self.get_slot(var, 'm')
+        m_t = tf.where(update_cond, m * beta_1_t + g_t * (1 - beta_1_t), m)
+        m_t = state_ops.assign(m, m_t, use_locking=self._use_locking)
+
+        v = self.get_slot(var, 'v')
+        v_t = tf.where(update_cond, v * beta_2_t + (g_t * g_t) * (1 - beta_2_t), v)
+        v_t = state_ops.assign(v, v_t, use_locking=self._use_locking)
+
+        if not self.amsgrad:
+            v_sqrt = math_ops.sqrt(v_t)
+            var_update = state_ops.assign_sub(
+                    var, lr * m_t / (v_sqrt + epsilon_t), use_locking=self._use_locking)
+            return control_flow_ops.group(*[var_update, m_t, v_t])
+        else:
+            v_hat = self.get_slot(var, 'vhat')
+            v_hat_t = tf.where(update_cond, math_ops.maximum(v_hat, v_t), v_hat)
+            with ops.control_dependencies([v_hat_t]):
+                v_hat_t = state_ops.assign(
+                        v_hat, v_hat_t, use_locking=self._use_locking)
+            v_hat_sqrt = math_ops.sqrt(v_hat_t)
+            var_update = state_ops.assign_sub(
+                    var,
+                    lr * m_t / (v_hat_sqrt + epsilon_t),
+                    use_locking=self._use_locking)
+            return control_flow_ops.group(*[var_update, m_t, v_t, v_hat_t])
+
+    def get_config(self):
+        config = super(AdamAccumulated, self).get_config()
+        config.update({
+            'accumulation_steps': self._serialize_hyperparameter('accumulation_steps'),
+            'learning_rate': self._serialize_hyperparameter('learning_rate'),
+            'decay': self._serialize_hyperparameter('decay'),
+            'beta_1': self._serialize_hyperparameter('beta_1'),
+            'beta_2': self._serialize_hyperparameter('beta_2'),
+            'epsilon': self.epsilon,
+            'amsgrad': self.amsgrad,
+        })
+        return config

DeepDeformationMapRegistration/utils/constants.py

@@ -52,6 +52,8 @@ 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_FIX_CENTROID = 'input/fix_centroid'
+H5_MOV_CENTROID = 'input/mov_centroid'
 
 H5_GT_DISP = 'output/{}'.format(DISP_MAP_GT)
 H5_GT_IMG = 'output/{}'.format(PRED_IMG_GT)
@@ -64,6 +66,7 @@ MAX_ANGLE = 45.0  # degrees
 MAX_FLIPS = 2  # Axes to flip over
 NUM_ROTATIONS = 5
 MAX_WORKERS = 10
+DEG_TO_RAD = np.pi/180.
 
 # Labels to pass to the input_labels and output_labels parameter of DataGeneratorManager
 DG_LBL_FIX_IMG = H5_FIX_IMG
@@ -75,6 +78,8 @@ DG_LBL_MOV_VESSELS = H5_MOV_VESSELS_MASK
 DG_LBL_MOV_PARENCHYMA = H5_MOV_PARENCHYMA_MASK
 DG_LBL_MOV_TUMOR = H5_MOV_TUMORS_MASK
 DG_LBL_ZERO_GRADS = 'zero_gradients'
+DG_LBL_FIX_CENTROID = H5_FIX_CENTROID
+DG_LBL_MOV_CENTROID = H5_MOV_CENTROID
 
 # Training constants
 MODEL = 'unet'
@@ -91,6 +96,8 @@ DATA_FORMAT = 'channels_last'  # or 'channels_fist'
 DATA_DIR = './data'
 MODEL_CHECKPOINT = './model_checkpoint'
 BATCH_SIZE = 8
+ACCUM_GRADIENT_STEP = 1
+EARLY_STOP_PATIENCE = 30  # Weights are updated every ACCUM_GRADIENT_STEPth step
 EPOCHS = 100
 SAVE_EPOCH = EPOCHS // 10  # Epoch when to save the model
 VERBOSE_EPOCH = EPOCHS // 10
@@ -99,7 +106,6 @@ VALIDATION_ERR_LIMIT_COUNTER = 10  # Number of successive times the validation e
 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',
@@ -362,10 +368,13 @@ 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
+        self.__spacing = 15
 
     def set_spacing(self, img_shape: tf.TensorShape):
-        self.__spacing = int(5 * np.log(img_shape[W]))
+        if isinstance(img_shape, tf.TensorShape):
+            self.__spacing = int(5 * np.log(img_shape[W]))
+        else:
+            self.__spacing = img_shape
 
     @property
     def spacing(self):
@@ -495,3 +504,19 @@ MANUAL_W = [1.] * len(PRIOR_W)
 
 REG_PRIOR_W = [1e-3]
 REG_MANUAL_W = [1.] * len(REG_PRIOR_W)
+
+# Constants for augmentation layer
+# .../T1/training/zoom_factors.csv contain the scale factors of all the training samples from isotropic to 128x128x128
+#   The augmentation values will be scaled using the average+std
+IXI_DATASET_iso_to_cubic_scales = np.asarray([0.655491 + 0.039223, 0.496783 + 0.029349, 0.499691 + 0.028155])
+MAX_AUG_DISP_ISOT = 30
+MAX_AUG_DEF_ISOT = 6
+MAX_AUG_DISP = np.max(MAX_AUG_DISP_ISOT * IXI_DATASET_iso_to_cubic_scales)  # Scaled displacements
+MAX_AUG_DEF = np.max(MAX_AUG_DEF_ISOT * IXI_DATASET_iso_to_cubic_scales)  # Scaled deformations
+MAX_AUG_ANGLE = np.max([np.arctan(np.tan(10*np.pi/180) * IXI_DATASET_iso_to_cubic_scales[1] / IXI_DATASET_iso_to_cubic_scales[0]) * 180 / np.pi,
+                        np.arctan(np.tan(10*np.pi/180) * IXI_DATASET_iso_to_cubic_scales[2] / IXI_DATASET_iso_to_cubic_scales[1]) * 180 / np.pi,
+                        np.arctan(np.tan(10*np.pi/180) * IXI_DATASET_iso_to_cubic_scales[2] / IXI_DATASET_iso_to_cubic_scales[0]) * 180 / np.pi])  # Scaled angles
+GAMMA_AUGMENTATION = False
+BRIGHTNESS_AUGMENTATION = False
+NUM_CONTROL_PTS_AUG = 10
+NUM_AUGMENTATIONS = 1

DeepDeformationMapRegistration/utils/misc.py

@@ -1,11 +1,19 @@
 import os
 import errno
+import shutil
+import numpy as np
+from scipy.interpolate import griddata, Rbf, LinearNDInterpolator, NearestNDInterpolator
+from skimage.measure import regionprops
+from DeepDeformationMapRegistration.layers.b_splines import interpolate_spline
+from DeepDeformationMapRegistration.utils.thin_plate_splines import ThinPlateSplines
+from tensorflow import squeeze
 
-def try_mkdir(dir):
+
+def try_mkdir(dir, verbose=True):
     try:
         os.makedirs(dir)
     except OSError as err:
-        if err.errno == errno.EEXIST:
+        if err.errno == errno.EEXIST and verbose:
             print("Directory " + dir + " already exists")
         else:
             raise ValueError("Can't create dir " + dir)
@@ -23,3 +31,120 @@ def function_decorator(new_name):
         func.__name__ = new_name
         return func
     return decorator
+
+
+class DatasetCopy:
+    def __init__(self, dataset_location, copy_location=None, verbose=True):
+        self.__copy_loc = os.path.join(os.getcwd(), 'temp_dataset') if copy_location is None else copy_location
+        self.__dst_loc = dataset_location
+        self.__verbose = verbose
+
+    def copy_dataset(self):
+        shutil.copytree(self.__dst_loc, self.__copy_loc)
+        if self.__verbose:
+            print('{} copied to {}'.format(self.__dst_loc, self.__copy_loc))
+        return self.__copy_loc
+
+    def delete_temp(self):
+        shutil.rmtree(self.__copy_loc)
+        if self.__verbose:
+            print('Deleted: ', self.__copy_loc)
+
+
+class DisplacementMapInterpolator:
+    def __init__(self,
+                 image_shape=[64, 64, 64],
+                 method='rbf'):
+        assert method in ['rbf', 'griddata', 'tf', 'tps'], "Method must be 'rbf' or 'griddata'"
+        self.method = method
+        self.image_shape = image_shape
+
+        self.grid = self.__regular_grid()
+
+    def __regular_grid(self):
+        xx = np.linspace(0, self.image_shape[0], self.image_shape[0], endpoint=False, dtype=np.uint16)
+        yy = np.linspace(0, self.image_shape[0], self.image_shape[0], endpoint=False, dtype=np.uint16)
+        zz = np.linspace(0, self.image_shape[0], self.image_shape[0], endpoint=False, dtype=np.uint16)
+
+        xx, yy, zz = np.meshgrid(xx, yy, zz)
+
+        return np.stack([xx.flatten(), yy.flatten(), zz.flatten()], axis=0).T
+
+    def __call__(self, disp_map, interp_points, backwards=False):
+        disp_map = disp_map.reshape([-1, 3])
+        grid_pts = self.grid.copy()
+        if backwards:
+            grid_pts = np.add(grid_pts, disp_map).astype(np.float32)
+            disp_map *= -1
+
+        if self.method == 'rbf':
+            interpolator = Rbf(grid_pts[:, 0], grid_pts[:, 1], grid_pts[:, 2], disp_map[:, :],
+                               method='thin_plate', mode='N-D')
+            disp = interpolator(interp_points)
+        elif self.method == 'griddata':
+            linear_interp = LinearNDInterpolator(grid_pts, disp_map)
+            disp = linear_interp(interp_points).copy()
+            del linear_interp
+
+            if np.any(np.isnan(disp)):
+                # It might happen (though it shouldn't) that the interpolation point is outside the convex hull of grid points.
+                #   in this situation, linear interpolation fails and will put NaN. Nearest can give a value, so we are going to
+                #   substitute those unexpected NaNs with the nearest value. Unexpected == not in interp_points
+                nan_disp_idx = set(np.unique(np.argwhere(np.isnan(disp))[:, 0]))
+                nan_interp_pts_idx = set(np.unique(np.argwhere(np.isnan(interp_points))[:, 0]))
+                idx = nan_disp_idx - nan_interp_pts_idx if len(nan_disp_idx) > len(nan_interp_pts_idx) else nan_interp_pts_idx - nan_disp_idx
+                idx = list(idx)
+                if len(idx):
+                    # We have unexpected NaNs
+                    near_interp = NearestNDInterpolator(grid_pts, disp_map)
+                    near_disp = near_interp(interp_points[idx, ...]).copy()
+                    del near_interp
+                    for n, i in enumerate(idx):
+                        disp[i, ...] = near_disp[n, ...]
+        elif self.method == 'tf':
+            # Order: 1 -> linear, 2 -> thin plate, 3 -> cubic
+            disp = squeeze(interpolate_spline(grid_pts[np.newaxis, ...][::4, :],    # Batch axis
+                                              disp_map[np.newaxis, ...][::4, :],
+                                              interp_points[np.newaxis, ...], order=2), axis=0)
+        else:
+            tps_interp = ThinPlateSplines(grid_pts[::8, :], self.grid.copy().astype(np.float32)[::8, :])
+            disp = tps_interp.interpolate(interp_points).eval()
+            del tps_interp
+
+        return disp
+
+
+def get_segmentations_centroids(segmentations, ohe=True, expected_lbls=range(0, 28), missing_centroid=[np.nan]*3, brain_study=True):
+    segmentations = np.squeeze(segmentations)
+    if ohe:
+        segmentations = np.sum(segmentations, axis=-1).astype(np.uint8)
+        missing_lbls = set(expected_lbls) - set(np.unique(segmentations))
+        if brain_study:
+            segmentations += np.ones_like(segmentations)  # Regionsprops neglect the label 0. But we need it, so offset all labels by 1
+    else:
+        missing_lbls = set(expected_lbls) - set(np.unique(segmentations))
+
+    seg_props = regionprops(segmentations)
+    centroids = np.asarray([c.centroid for c in seg_props]).astype(np.float32)
+
+    for lbl in missing_lbls:
+        idx = expected_lbls.index(lbl)
+        centroids = np.insert(centroids, idx, missing_centroid, axis=0)
+    return centroids.copy(), missing_lbls
+
+
+def segmentation_ohe_to_cardinal(segmentation):
+    cpy = segmentation.copy()
+    for lbl in range(segmentation.shape[-1]):
+        cpy[..., lbl] *= (lbl + 1)
+    # Add the Background
+    cpy = np.concatenate([np.zeros(segmentation.shape[:-1])[..., np.newaxis], cpy], axis=-1)
+    return np.argmax(cpy, axis=-1)[..., np.newaxis]
+
+
+def segmentation_cardinal_to_ohe(segmentation):
+    # Keep in mind that we don't handle the overlap between the segmentations!
+    cpy = np.tile(np.zeros_like(segmentation), (1, 1, 1, len(np.unique(segmentation)[1:])))
+    for ch, lbl in enumerate(np.unique(segmentation)[1:]):
+        cpy[segmentation == lbl, ch] = 1
+    return cpy

DeepDeformationMapRegistration/utils/nifti_utils.py

@@ -0,0 +1,43 @@
+import os
+import nibabel as nb
+import numpy as np
+import zipfile
+
+
+TEMP_UNZIP_PATH = '/mnt/EncryptedData1/Users/javier/ext_datasets/LITS17/temp'
+NII_EXTENSION = '.nii'
+
+
+def save_nifti(data, save_path, header=None, verbose=True):
+    if header is None:
+        data_nifti = nb.Nifti1Image(data, affine=np.eye(4))
+    else:
+        data_nifti = nb.Nifti1Image(data, affine=None, header=header)
+
+    data_nifti.header.get_xyzt_units()
+    try:
+        nb.save(data_nifti, save_path)  # Save as NiBabel file
+        if verbose:
+            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.py

@@ -11,14 +11,23 @@ def min_max_norm(img: np.ndarray, out_max_val=1.):
     return out_img * out_max_val
 
 
-def soft_threshold(x, threshold, name=None):
+def soft_threshold(x, threshold, name='soft_threshold'):
     # https://www.tensorflow.org/probability/api_docs/python/tfp/math/soft_threshold
-    with tf.name_scope(name or 'soft_threshold'):
+    # Foucart S., Rauhut H. (2013) Basic Algorithms. In: A Mathematical Introduction to Compressive Sensing.
+    #   Applied and Numerical Harmonic Analysis. Birkhäuser, New York, NY. https://doi.org/10.1007/978-0-8176-4948-7_3
+    #   Chapter 3, page 72
+    with tf.name_scope(name):
         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 hard_threshold(x, threshold, name='hard_threshold'):
+    with tf.name_scope(name):
+        threshold = tf.convert_to_tensor(threshold, dtype=x.dtype, name='threshold')
+        return tf.sign(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)))
@@ -26,3 +35,31 @@ def binary_activation(x):
 
     return out
 
+
+def gaussian_kernel(kernel_size, sigma, in_ch, out_ch, dim, dtype=tf.float32):
+    # SRC: https://stackoverflow.com/questions/59286171/gaussian-blur-image-in-dataset-pipeline-in-tensorflow
+    x = tf.range(-kernel_size // 2 + 1, kernel_size // 2 + 1, dtype=dtype)
+    g = tf.math.exp(-(tf.pow(x, 2) / (2 * tf.pow(tf.cast(sigma, dtype), 2))))
+
+    g_kernel = tf.identity(g)
+    g_kernel = tf.tensordot(g_kernel, g, 0)
+    g_kernel = tf.tensordot(g_kernel, g, 0)
+
+    # i = tf.constant(0)
+    # cond = lambda i, g_kern: tf.less(i, dim - 1)
+    # mult_kern = lambda i, g_kern: [tf.add(i, 1), tf.tensordot(g_kern, g, 0)]
+    # _, g_kernel = tf.while_loop(cond, mult_kern,
+    #                             loop_vars=[i, g_kernel],
+    #                             shape_invariants=[i.get_shape(), tf.TensorShape([kernel_size, None, None])])
+
+    g_kernel = g_kernel / tf.reduce_sum(g_kernel)
+    g_kernel = tf.expand_dims(tf.expand_dims(g_kernel, axis=-1), axis=-1)
+    return tf.tile(g_kernel, (*(1,)*dim, in_ch, out_ch))
+
+
+def sample_unique(population, samples, tout=tf.int32):
+    # src: https://github.com/tensorflow/tensorflow/issues/9260#issuecomment-437875125
+    z = -tf.log(-tf.log(tf.random_uniform((tf.shape(population)[0],), 0, 1)))
+    _, indices = tf.nn.top_k(z, samples)
+    ret_val = tf.gather(population, indices)
+    return tf.cast(ret_val, tout)

DeepDeformationMapRegistration/utils/thin_plate_splines.py

@@ -0,0 +1,154 @@
+import numpy as np
+import tensorflow as tf
+
+
+class ThinPlateSplines:
+    def __init__(self, ctrl_pts: tf.Tensor, target_pts: tf.Tensor, reg=0.0):
+        """
+
+        :param ctrl_pts: [N, d] tensor of control d-dimensional points
+        :param target_pts: [N, d] tensor of target d-dimensional points
+        :param reg: regularization coefficient
+        """
+        self.__ctrl_pts = ctrl_pts
+        self.__target_pts = target_pts
+        self.__reg = reg
+        self.__num_ctrl_pts = ctrl_pts.shape[0]
+        self.__dim = ctrl_pts.shape[1]
+
+        self.__compute_coeffs()
+        # self.__aff_params = self.__coeffs[self.__num_ctrl_pts:, ...]   # Affine  parameters of the TPS
+        self.__non_aff_paramms = self.__coeffs[:self.__num_ctrl_pts, ...]      # Non-affine parameters of  he TPS
+
+    def __compute_coeffs(self):
+        target_pts_aug = tf.concat([self.__target_pts,
+                                    tf.zeros([self.__dim + 1, self.__dim], dtype=self.__target_pts.dtype)],
+                                   axis=0)
+
+        # T = self.__make_T()
+        T_i = tf.cast(tf.linalg.inv(self.__make_T()), target_pts_aug.dtype)
+        self.__coeffs = tf.cast(tf.matmul(T_i, target_pts_aug), tf.float32)
+
+    def __make_T(self):
+        # cp: [K x 2] control points
+        # T: [(num_pts+dim+1) x (num_pts+dim+1)]
+        num_pts = self.__ctrl_pts.shape[0]
+
+        P = tf.concat([tf.ones([self.__num_ctrl_pts, 1], dtype=tf.float32), self.__ctrl_pts], axis=1)
+        zeros = np.zeros([self.__dim + 1, self.__dim + 1], dtype=np.float)
+        self.__K = self.__U_dist(self.__ctrl_pts)
+        alfa = tf.reduce_mean(self.__K)
+
+        self.__K = self.__K + tf.ones_like(self.__K) * tf.pow(alfa, 2) * self.__reg
+
+        # top = tf.concat([self.__K, P], axis=1)
+        # bottom = tf.concat([tf.transpose(P), zeros], axis=1)
+
+        return tf.concat([tf.concat([self.__K, P], axis=1), tf.concat([tf.transpose(P), zeros], axis=1)], axis=0)
+
+    def __U_dist(self, ctrl_pts, int_pts=None):
+        if int_pts is None:
+            dist = self.__pairwise_distance_equal(ctrl_pts)  # Already squared!
+        else:
+            dist = self.__pairwise_distance_different(ctrl_pts, int_pts)  # Already squared!
+
+
+        # U(x, y) = p_w_dist(x, y)^2 * log(p_w_dist(x, y)) (dist() > =0); 0 otw
+        if ctrl_pts.shape[-1] == 2:
+            u_dist = dist * tf.math.log(dist + 1e-6)
+        else:
+            # Src: https://github.com/vaipatel/morphops/blob/master/morphops/tps.py
+            # In particular, if k = 2, then  U(r) = r^2 * log(r^2), else U(r) = r
+            u_dist = tf.sqrt(dist)
+        # tf.matrix_set_diag(u_dist, tf.constant(0, dtype=dist_sq.dtype))
+        # reg_term = self.__reg * tf.pow(alfa, 2) * tf.eye(self.__num_ctrl_pts)
+
+        return u_dist # + reg_term
+
+    def __pairwise_distance_sq(self, pts_a, pts_b):
+        with tf.variable_scope('pairwise_distance'):
+            if np.all(pts_a == pts_b):
+                # This implementation works better when doing the pairwise distance os a single set of points
+                pts_a_ = tf.reshape(pts_a, [-1, 1, 3])
+                pts_b_ = tf.reshape(pts_b, [1, -1, 3])
+                dist = tf.reduce_sum(tf.square(pts_a_ - pts_b_), 2)   # squared pairwise distance
+            else:
+                # PwD^2= A_norm^2 - 2*A*B' + B_norm^2
+                pts_a_ = tf.reduce_sum(tf.square(pts_a), 1)
+                pts_b_ = tf.reduce_sum(tf.square(pts_b), 1)
+
+                pts_a_ = tf.expand_dims(pts_a_, 1)
+                pts_b_ = tf.expand_dims(pts_b_, 0)
+
+                pts_a_pts_b_ = tf.matmul(pts_a, pts_b, adjoint_b=True)
+
+                dist = pts_a_ - 2 * pts_a_pts_b_ + pts_b_
+
+            return tf.cast(dist, tf.float32)
+
+    @staticmethod
+    def __pairwise_distance_equal(pts):
+        # This implementation works better when doing the pairwise distance os a single set of points
+        dist = tf.reduce_sum(tf.square(tf.reshape(pts, [-1, 1, 3]) - tf.reshape(pts, [1, -1, 3])), 2)  # squared pairwise distance
+        return tf.cast(dist, tf.float32)
+
+    @staticmethod
+    def __pairwise_distance_different(pts_a, pts_b):
+        pts_a_ = tf.reduce_sum(tf.square(pts_a), 1)
+        pts_b_ = tf.reduce_sum(tf.square(pts_b), 1)
+
+        pts_a_ = tf.expand_dims(pts_a_, 1)
+        pts_b_ = tf.expand_dims(pts_b_, 0)
+
+        pts_a_pts_b_ = tf.matmul(pts_a, pts_b, adjoint_b=True)
+
+        dist = pts_a_ - 2 * pts_a_pts_b_ + pts_b_
+        return tf.cast(dist, tf.float32)
+
+    def __lift_pts(self, int_pts: tf.Tensor, num_pts):
+        # int_pts: [N x 2], input points
+        # cp: [K x 2], control points
+        # pLift: [N x (3+K)], lifted input points
+
+        # u_dist = self.__U_dist(int_pts, self.__ctrl_pts)
+
+        int_pts_lift = tf.concat([self.__U_dist(int_pts, self.__ctrl_pts),
+                                 tf.ones([num_pts, 1], dtype=tf.float32),
+                                 int_pts], axis=1)
+        return int_pts_lift
+
+    @property
+    def bending_energy(self):
+        aux = tf.matmul(self.__non_aff_paramms, self.__K, transpose_a=True)
+        return tf.matmul(aux, self.__non_aff_paramms)
+
+    def interpolate(self, int_points): #, num_pts):
+        """
+
+        :param int_points: [K, d] flattened d-points of a mesh
+        :return:
+        """
+        num_pts = tf.shape(int_points)[0]
+        int_points_lift = self.__lift_pts(int_points, num_pts)
+
+        return tf.matmul(int_points_lift, self.__coeffs)
+
+    def __call__(self, int_points, num_pts, **kwargs):
+        return self.interpolate(int_points) # , num_pts)
+
+
+def thin_plate_splines_batch(ctrl_pts: tf.Tensor, target_pts: tf.Tensor, int_pts: tf.Tensor, reg=0.0):
+    _batches = ctrl_pts.shape[0]
+
+    if tf.get_default_session() is not None:
+        print('DEBUG TIME')
+
+    def tps_sample(in_data):
+        cp, tp, ip = in_data
+        # _num_pts = ip.shape[0]
+        tps = ThinPlateSplines(cp, tp, reg)
+        interp = tps.interpolate(ip) # , _num_pts)
+        return interp
+
+    return tf.map_fn(tps_sample, elems=(ctrl_pts, target_pts, int_pts), dtype=tf.float32)
+

DeepDeformationMapRegistration/utils/visualization.py

@@ -1,5 +1,5 @@
-# import matplotlib
-# matplotlib.use('TkAgg')
+import matplotlib
+#matplotlib.use('TkAgg')
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
 import matplotlib.colors as mcolors
@@ -8,7 +8,7 @@ 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
+import DeepDeformationMapRegistration.utils.constants as C
 from skimage.exposure import rescale_intensity
 import scipy.misc as scpmisc
 import os
@@ -175,11 +175,11 @@ def plot_training(list_imgs: [np.ndarray], affine_transf=True, filename='img', f
         fig.clear()
         plt.figure(fig.number)
     else:
-        fig = plt.figure(dpi=const.DPI)
+        fig = plt.figure(dpi=C.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.set_title('Fix image', fontsize=C.FONT_SIZE)
     ax_fix.tick_params(axis='both',
                        which='both',
                        bottom=False,
@@ -188,7 +188,7 @@ def plot_training(list_imgs: [np.ndarray], affine_transf=True, filename='img', f
                        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.set_title('Moving image', fontsize=C.FONT_SIZE)
     ax_mov.tick_params(axis='both',
                        which='both',
                        bottom=False,
@@ -198,7 +198,7 @@ def plot_training(list_imgs: [np.ndarray], affine_transf=True, filename='img', f
 
     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.set_title('Prediction', fontsize=C.FONT_SIZE)
     ax_pred_im.tick_params(axis='both',
                        which='both',
                        bottom=False,
@@ -228,8 +228,8 @@ def plot_training(list_imgs: [np.ndarray], affine_transf=True, filename='img', f
     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.quiver(cx, cy, dx, dy, scale=C.QUIVER_PARAMS.arrow_scale)
+        ax_pred_disp.set_title('Pred disp map', fontsize=C.FONT_SIZE)
     ax_pred_disp.tick_params(axis='both',
                                which='both',
                                bottom=False,
@@ -259,8 +259,8 @@ def plot_training(list_imgs: [np.ndarray], affine_transf=True, filename='img', f
     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.quiver(cx, cy, dx, dy, scale=C.QUIVER_PARAMS.arrow_scale)
+        ax_gt_disp.set_title('GT disp map', fontsize=C.FONT_SIZE)
     ax_gt_disp.tick_params(axis='both',
                        which='both',
                        bottom=False,
@@ -276,7 +276,7 @@ def plot_training(list_imgs: [np.ndarray], affine_transf=True, filename='img', f
 
     if filename is not None:
         plt.savefig(filename, format='png')  # Call before show
-    if not const.REMOTE:
+    if not C.REMOTE:
         plt.show()
     else:
         plt.close()
@@ -288,7 +288,7 @@ def save_centreline_img(img, title, filename, fig=None):
         fig.clear()
         plt.figure(fig.number)
     else:
-        fig = plt.figure(dpi=const.DPI)
+        fig = plt.figure(dpi=C.DPI)
 
     dim = len(img.shape[:-1])
 
@@ -321,19 +321,19 @@ def save_centreline_img(img, title, filename, fig=None):
     plt.close()
 
 
-def save_disp_map_img(disp_map, title, filename, affine_transf=False, fig=None):
+def save_disp_map_img(disp_map, title, filename, affine_transf=False, fig=None, show=False):
     if fig is not None:
         fig.clear()
         plt.figure(fig.number)
     else:
-        fig = plt.figure(dpi=const.DPI)
-
+        fig = plt.figure(dpi=C.DPI)
+    dim_h, dim_w, dim_d = disp_map.shape[1:-1]
     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])
+        im_x = ax_x.imshow(disp_map[..., C.H_DISP])
         ax_x.tick_params(axis='both',
                        which='both',
                        bottom=False,
@@ -344,7 +344,7 @@ def save_disp_map_img(disp_map, title, filename, affine_transf=False, fig=None):
 
         ax_y = fig.add_subplot(132)
         ax_y.set_title('W displacement')
-        im_y = ax_y.imshow(disp_map[..., const.W_DISP])
+        im_y = ax_y.imshow(disp_map[..., C.W_DISP])
         ax_y.tick_params(axis='both',
                          which='both',
                          bottom=False,
@@ -373,8 +373,8 @@ def save_disp_map_img(disp_map, title, filename, affine_transf=False, fig=None):
         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)
+            ax.quiver(c[C.H_DISP], c[C.W_DISP], d[C.H_DISP], d[C.W_DISP],
+                      scale=C.QUIVER_PARAMS.arrow_scale)
             cb = _set_colorbar(fig, ax, im, False)
             ax.set_title('Displacement map')
         ax.tick_params(axis='both',
@@ -387,9 +387,8 @@ def save_disp_map_img(disp_map, title, filename, affine_transf=False, fig=None):
     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)
+        ax.quiver(c[C.H_DISP], c[C.W_DISP], c[C.D_DISP], d[C.H_DISP], d[C.W_DISP], d[C.D_DISP])
+        _square_3d_plot(np.arange(0, dim_h-1), np.arange(0, dim_w-1), np.arange(0, dim_d-1), 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',
@@ -397,10 +396,14 @@ def save_disp_map_img(disp_map, title, filename, affine_transf=False, fig=None):
                        left=False,
                        labelleft=False,
                        labelbottom=False)
+        add_axes_arrows_3d(ax, xyz_label=['R', 'A', 'S'])
         fig.suptitle(title)
 
     plt.savefig(filename, format='png')
+    if show:
+        plt.show()
     plt.close()
+    return fig
 
 
 def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, filename='img', fig=None,
@@ -409,16 +412,16 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
         fig.clear()
         plt.figure(fig.number)
     else:
-        fig = plt.figure(dpi=const.DPI)
+        fig = plt.figure(dpi=C.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)
+        ax_input.set_ylabel(title_first_row, fontsize=C.FONT_SIZE)
         im_fix = ax_input.imshow(list_imgs[0][:, :, 0])
-        ax_input.set_title('Fix image', fontsize=const.FONT_SIZE)
+        ax_input.set_title('Fix image', fontsize=C.FONT_SIZE)
         ax_input.tick_params(axis='both',
                            which='both',
                            bottom=False,
@@ -427,7 +430,7 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
                            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.set_title('Moving image', fontsize=C.FONT_SIZE)
         ax_mov.tick_params(axis='both',
                            which='both',
                            bottom=False,
@@ -437,7 +440,7 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
 
         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.set_title('Predicted fix image', fontsize=C.FONT_SIZE)
         ax_pred_im.tick_params(axis='both',
                            which='both',
                            bottom=False,
@@ -467,8 +470,8 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
         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.quiver(cx, cy, dx, dy, scale=C.QUIVER_PARAMS.arrow_scale)
+            ax_pred_disp.set_title('Pred disp map', fontsize=C.FONT_SIZE)
         ax_pred_disp.tick_params(axis='both',
                                    which='both',
                                    bottom=False,
@@ -478,9 +481,9 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
 
         # VALIDATION
         axinput_val = fig.add_subplot(245)
-        axinput_val.set_ylabel(title_second_row, fontsize=const.FONT_SIZE)
+        axinput_val.set_ylabel(title_second_row, fontsize=C.FONT_SIZE)
         im_fix_val = axinput_val.imshow(list_imgs[4][:, :, 0])
-        axinput_val.set_title('Fix image', fontsize=const.FONT_SIZE)
+        axinput_val.set_title('Fix image', fontsize=C.FONT_SIZE)
         axinput_val.tick_params(axis='both',
                            which='both',
                            bottom=False,
@@ -489,7 +492,7 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
                            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.set_title('Moving image', fontsize=C.FONT_SIZE)
         ax_mov_val.tick_params(axis='both',
                            which='both',
                            bottom=False,
@@ -499,7 +502,7 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
 
         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.set_title('Predicted fix image', fontsize=C.FONT_SIZE)
         ax_pred_im_val.tick_params(axis='both',
                            which='both',
                            bottom=False,
@@ -529,8 +532,8 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
         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.quiver(c[0], c[1], d[0], d[1], scale=C.QUIVER_PARAMS.arrow_scale)
+            ax_pred_disp_val.set_title('Pred disp map', fontsize=C.FONT_SIZE)
         ax_pred_disp_val.tick_params(axis='both',
                                    which='both',
                                    bottom=False,
@@ -552,10 +555,10 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
         # 3D
         # TRAINING
         ax_input = fig.add_subplot(231, projection='3d')
-        ax_input.set_ylabel(title_first_row, fontsize=const.FONT_SIZE)
+        ax_input.set_ylabel(title_first_row, fontsize=C.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.set_title('Fix image', fontsize=C.FONT_SIZE)
         ax_input.tick_params(axis='both',
                            which='both',
                            bottom=False,
@@ -566,7 +569,7 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
         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.set_title('Predicted fix image', fontsize=C.FONT_SIZE)
         ax_pred_im.tick_params(axis='both',
                                which='both',
                                bottom=False,
@@ -578,9 +581,9 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
 
         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.quiver(c[C.H_DISP], c[C.W_DISP], c[C.D_DISP],
+                            d[C.H_DISP], d[C.W_DISP], d[C.D_DISP], scale=C.QUIVER_PARAMS.arrow_scale)
+        ax_pred_disp.set_title('Pred disp map', fontsize=C.FONT_SIZE)
         ax_pred_disp.tick_params(axis='both',
                                  which='both',
                                  bottom=False,
@@ -590,10 +593,10 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
 
         # VALIDATION
         axinput_val = fig.add_subplot(234, projection='3d')
-        axinput_val.set_ylabel(title_second_row, fontsize=const.FONT_SIZE)
+        axinput_val.set_ylabel(title_second_row, fontsize=C.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.set_title('Fix image', fontsize=C.FONT_SIZE)
         axinput_val.tick_params(axis='both',
                                which='both',
                                bottom=False,
@@ -604,7 +607,7 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
         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.set_title('Predicted fix image', fontsize=C.FONT_SIZE)
         ax_pred_im_val.tick_params(axis='both',
                                    which='both',
                                    bottom=False,
@@ -615,10 +618,10 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
         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.quiver(c[C.H_DISP], c[C.W_DISP], c[C.D_DISP],
+                                d[C.H_DISP], d[C.W_DISP], d[C.D_DISP],
+                                scale=C.QUIVER_PARAMS.arrow_scale)
+        ax_pred_disp_val.set_title('Pred disp map', fontsize=C.FONT_SIZE)
         ax_pred_disp_val.tick_params(axis='both',
                                      which='both',
                                      bottom=False,
@@ -628,7 +631,7 @@ def plot_training_and_validation(list_imgs: [np.ndarray], affine_transf=True, fi
 
     if filename is not None:
         plt.savefig(filename, format='png')  # Call before show
-    if not const.REMOTE:
+    if not C.REMOTE:
         plt.show()
     else:
         plt.close()
@@ -644,21 +647,21 @@ def _set_colorbar(fig, ax, im, drawedges=True):
     return im_cb
 
 
-def _prepare_quiver_map(disp_map: np.ndarray, dim=2, spc=const.QUIVER_PARAMS.spacing):
+def _prepare_quiver_map(disp_map: np.ndarray, dim=2, spc=C.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]
+    dx = disp_map[..., C.H_DISP]
+    dy = disp_map[..., C.W_DISP]
     if dim > 2:
-        dz = disp_map[..., const.D_DISP]
+        dz = disp_map[..., C.D_DISP]
 
-    img_size_x = disp_map.shape[const.H_DISP]
-    img_size_y = disp_map.shape[const.W_DISP]
+    img_size_x = disp_map.shape[C.H_DISP]
+    img_size_y = disp_map.shape[C.W_DISP]
     if dim > 2:
-        img_size_z = disp_map.shape[const.D_DISP]
+        img_size_z = disp_map.shape[C.D_DISP]
 
     if dim > 2:
         s = np.sqrt(np.square(dx) + np.square(dy) + np.square(dz))
@@ -728,7 +731,7 @@ def plot_input_data(fix_img, mov_img, img_size=(64, 64), title=None, filename=No
 
     if filename is not None:
         plt.savefig(filename, format='png')  # Call before show
-    if not const.REMOTE:
+    if not C.REMOTE:
         plt.show()
     else:
         plt.close()
@@ -807,29 +810,42 @@ def plot_dataset_3d(img_sets):
     return fig
 
 
-def plot_predictions(fix_img_batch, mov_img_batch, disp_map_batch, pred_img_batch, filename='predictions', fig=None):
+def plot_predictions(fix_img_batch, mov_img_batch, disp_map_batch, pred_img_batch, filename='predictions', fig=None, show=False):
     num_rows = fix_img_batch.shape[0]
-    img_size = fix_img_batch.shape[1:3]
+    img_dim = len(fix_img_batch.shape) - 2
+    img_size = fix_img_batch.shape[1:-1]
     if fig is not None:
         fig.clear()
         plt.figure(fig.number)
-        ax = fig.add_subplot(nrows=num_rows, ncols=4, dpi=const.DPI)
+        ax = fig.add_subplot(nrows=num_rows, ncols=5, figsize=(10, 3*num_rows), dpi=C.DPI)
     else:
-        fig, ax = plt.subplots(nrows=num_rows, ncols=4, dpi=const.DPI)
+        fig, ax = plt.subplots(nrows=num_rows, ncols=5, figsize=(10, 3*num_rows), dpi=C.DPI)
+    if num_rows == 1:
+        ax = ax[np.newaxis, ...]
+
+    if img_dim == 3:    # Extract slices from the images
+        selected_slice = img_size[0] // 2
+        fix_img_batch = fix_img_batch[:, selected_slice, ...]
+        mov_img_batch = mov_img_batch[:, selected_slice, ...]
+        pred_img_batch = pred_img_batch[:, selected_slice, ...]
+        disp_map_batch = disp_map_batch[:, selected_slice, ..., 1:]  # Only the sagittal and longitudinal axes
+        img_size = fix_img_batch.shape[1:-1]
+    elif img_dim != 2:
+        raise ValueError('Images have a bad shape: {}'.format(fix_img_batch.shape))
 
     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)
+        fix_img = fix_img_batch[row, :, :, 0].transpose()
+        mov_img = mov_img_batch[row, :, :, 0].transpose()
+        disp_map = disp_map_batch[row, :, :, :].transpose((1, 0, 2))
+        pred_img = pred_img_batch[row, :, :, 0].transpose()
+        ax[row, 0].imshow(fix_img, origin='lower')
         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].imshow(mov_img, origin='lower')
         ax[row, 1].tick_params(axis='both',
                                which='both',
                                bottom=False,
@@ -837,11 +853,12 @@ def plot_predictions(fix_img_batch, mov_img_batch, disp_map_batch, pred_img_batc
                                labelleft=False,
                                labelbottom=False)
 
-        cx, cy, dx, dy, s = _prepare_quiver_map(disp_map)
+        c, d, s = _prepare_quiver_map(disp_map, spc=5)
+        cx, cy = c
+        dx, dy = d
         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].imshow(disp_map_color, interpolation='none', aspect='equal', origin='lower')
+        ax[row, 2].quiver(cx, cy, dx, dy, units='dots', scale=1)
         ax[row, 2].tick_params(axis='both',
                                which='both',
                                bottom=False,
@@ -849,6 +866,8 @@ def plot_predictions(fix_img_batch, mov_img_batch, disp_map_batch, pred_img_batc
                                labelleft=False,
                                labelbottom=False)
 
+        ax[row, 3].imshow(mov_img, origin='lower')
+        ax[row, 3].quiver(cx, cy, dx, dy, units='dots', scale=1, color='w')
         ax[row, 3].tick_params(axis='both',
                                which='both',
                                bottom=False,
@@ -856,18 +875,26 @@ def plot_predictions(fix_img_batch, mov_img_batch, disp_map_batch, pred_img_batc
                                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)
+        ax[row, 4].imshow(pred_img, origin='lower')
+        ax[row, 4].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=C.FONT_SIZE)
+    ax[0, 1].set_title('Moving img ($I_m$)', fontsize=C.FONT_SIZE)
+    ax[0, 2].set_title('Backwards\ndisp .map ($\delta$)', fontsize=C.FONT_SIZE)
+    ax[0, 3].set_title('Disp. map over $I_m$', fontsize=C.FONT_SIZE)
+    ax[0, 4].set_title('Predicted $I_m$', fontsize=C.FONT_SIZE)
+    plt.tight_layout()
     if filename is not None:
         plt.savefig(filename, format='png')  # Call before show
-    if not const.REMOTE:
+    if show:
         plt.show()
-    else:
-        plt.close()
+    plt.close()
     return fig
 
 
@@ -876,7 +903,7 @@ def inspect_disp_map_generation(fix_img, mov_img, disp_map, filename=None, fig=N
         fig.clear()
         plt.figure(fig.number)
     else:
-        fig = plt.figure(dpi=const.DPI)
+        fig = plt.figure(dpi=C.DPI)
 
     ax0 = fig.add_subplot(221)
     im0 = ax0.imshow(fix_img[..., 0])
@@ -900,7 +927,7 @@ def inspect_disp_map_generation(fix_img, mov_img, disp_map, filename=None, fig=N
     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.quiver(cx, cy, dx, dy, scale=C.QUIVER_PARAMS.arrow_scale)
     # ax2.figure.set_size_inches(img_size)
     ax2.tick_params(axis='both',
                       which='both',
@@ -912,7 +939,7 @@ def inspect_disp_map_generation(fix_img, mov_img, disp_map, filename=None, fig=N
     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.quiver(cx, cy, dx, dy, scale=C.QUIVER_PARAMS.arrow_scale)
     ax3.tick_params(axis='both',
                     which='both',
                     bottom=False,
@@ -921,10 +948,10 @@ def inspect_disp_map_generation(fix_img, mov_img, disp_map, filename=None, fig=N
                     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)
+    ax0.set_title('Fixed img ($I_f$)', fontsize=C.FONT_SIZE)
+    ax1.set_title('Moving img ($I_m$)', fontsize=C.FONT_SIZE)
+    ax2.set_title('Displacement map', fontsize=C.FONT_SIZE)
+    ax3.set_title('Fix - Mov', fontsize=C.FONT_SIZE)
 
     im0_cb = _set_colorbar(fig, ax0, im0, False)
     im1_cb = _set_colorbar(fig, ax1, im1, False)
@@ -933,7 +960,7 @@ def inspect_disp_map_generation(fix_img, mov_img, disp_map, filename=None, fig=N
 
     if filename is not None:
         plt.savefig(filename, format='png')  # Call before show
-    if not const.REMOTE:
+    if not C.REMOTE:
         plt.show()
     else:
         plt.close()
@@ -950,7 +977,7 @@ def inspect_displacement_grid(ctrl_coords, dense_coords, target_coords, disp_coo
         fig = plt.figure()
 
     ax_grid = fig.add_subplot(231)
-    ax_grid.set_title('Grids', fontsize=const.FONT_SIZE)
+    ax_grid.set_title('Grids', fontsize=C.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',
@@ -966,7 +993,7 @@ def inspect_displacement_grid(ctrl_coords, dense_coords, target_coords, disp_coo
     ax_grid.set_aspect('equal')
 
     ax_disp = fig.add_subplot(232)
-    ax_disp.set_title('Displacement map', fontsize=const.FONT_SIZE)
+    ax_disp.set_title('Displacement map', fontsize=C.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',
@@ -977,7 +1004,7 @@ def inspect_displacement_grid(ctrl_coords, dense_coords, target_coords, disp_coo
                         labelbottom=False)
 
     ax_mask = fig.add_subplot(233)
-    ax_mask.set_title('Mask', fontsize=const.FONT_SIZE)
+    ax_mask.set_title('Mask', fontsize=C.FONT_SIZE)
     ax_mask.imshow(mask)
     ax_mask.tick_params(axis='both',
                         which='both',
@@ -987,7 +1014,7 @@ def inspect_displacement_grid(ctrl_coords, dense_coords, target_coords, disp_coo
                         labelbottom=False)
 
     ax_fix = fig.add_subplot(234)
-    ax_fix.set_title('Fix image', fontsize=const.FONT_SIZE)
+    ax_fix.set_title('Fix image', fontsize=C.FONT_SIZE)
     ax_fix.imshow(fix_img[..., 0])
     ax_fix.tick_params(axis='both',
                         which='both',
@@ -997,7 +1024,7 @@ def inspect_displacement_grid(ctrl_coords, dense_coords, target_coords, disp_coo
                         labelbottom=False)
 
     ax_mov = fig.add_subplot(235)
-    ax_mov.set_title('Moving image', fontsize=const.FONT_SIZE)
+    ax_mov.set_title('Moving image', fontsize=C.FONT_SIZE)
     ax_mov.imshow(mov_img[..., 0])
     ax_mov.tick_params(axis='both',
                         which='both',
@@ -1007,7 +1034,7 @@ def inspect_displacement_grid(ctrl_coords, dense_coords, target_coords, disp_coo
                         labelbottom=False)
 
     ax_dif = fig.add_subplot(236)
-    ax_dif.set_title('Fix - Moving image', fontsize=const.FONT_SIZE)
+    ax_dif.set_title('Fix - Moving image', fontsize=C.FONT_SIZE)
     ax_dif.imshow(fix_img[..., 0] - mov_img[..., 0], cmap=cmap_bin)
     ax_dif.tick_params(axis='both',
                         which='both',
@@ -1023,7 +1050,7 @@ def inspect_displacement_grid(ctrl_coords, dense_coords, target_coords, disp_coo
 
     if filename is not None:
         plt.savefig(filename, format='png')  # Call before show
-    if not const.REMOTE:
+    if not C.REMOTE:
         plt.show()
 
     return fig
@@ -1037,10 +1064,10 @@ def compare_disp_maps(disp_m_f, disp_f_m, fix_img, mov_img, filename=None, fig=N
         fig = plt.figure()
 
     ax_d_m_f = fig.add_subplot(131)
-    ax_d_m_f.set_title('Disp M->F', fontsize=const.FONT_SIZE)
+    ax_d_m_f.set_title('Disp M->F', fontsize=C.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.quiver(cx, cy, dx, dy, scale=C.QUIVER_PARAMS.arrow_scale)
     ax_d_m_f.tick_params(axis='both',
                         which='both',
                         bottom=False,
@@ -1049,9 +1076,9 @@ def compare_disp_maps(disp_m_f, disp_f_m, fix_img, mov_img, filename=None, fig=N
                         labelbottom=False)
 
     ax_d_f_m = fig.add_subplot(132)
-    ax_d_f_m.set_title('Disp F->M', fontsize=const.FONT_SIZE)
+    ax_d_f_m.set_title('Disp F->M', fontsize=C.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.quiver(cx, cy, dx, dy, scale=C.QUIVER_PARAMS.arrow_scale)
     ax_d_f_m.imshow(s, interpolation='none', aspect='equal')
     ax_d_f_m.tick_params(axis='both',
                          which='both',
@@ -1061,7 +1088,7 @@ def compare_disp_maps(disp_m_f, disp_f_m, fix_img, mov_img, filename=None, fig=N
                          labelbottom=False)
 
     ax_dif = fig.add_subplot(133)
-    ax_dif.set_title('Fix - Moving image', fontsize=const.FONT_SIZE)
+    ax_dif.set_title('Fix - Moving image', fontsize=C.FONT_SIZE)
     ax_dif.imshow(fix_img[..., 0] - mov_img[..., 0], cmap=cmap_bin)
     ax_dif.tick_params(axis='both',
                        which='both',
@@ -1078,7 +1105,7 @@ def compare_disp_maps(disp_m_f, disp_f_m, fix_img, mov_img, filename=None, fig=N
 
     if filename is not None:
         plt.savefig(filename, format='png')  # Call before show
-    if not const.REMOTE:
+    if not C.REMOTE:
         plt.show()
     else:
         plt.close()
@@ -1104,7 +1131,7 @@ def plot_train_step(list_imgs: [np.ndarray], fig_title='TRAINING', dest_folder='
     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)
+    ax.set_title('Fix image', fontsize=C.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):
@@ -1116,23 +1143,23 @@ def plot_train_step(list_imgs: [np.ndarray], fig_title='TRAINING', dest_folder='
 
     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)
+    ax.set_title('Fix image', fontsize=C.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],
+        ax.quiver(c[C.H_DISP], c[C.W_DISP], c[C.D_DISP],
+                  d[C.H_DISP], d[C.W_DISP], d[C.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)
+    fig.suptitle(fig_title, fontsize=C.FONT_SIZE)
 
     if save_file:
         plt.savefig(os.path.join(dest_folder, fig_title+'.png'), format='png')  # Call before show
-    if not const.REMOTE:
+    if not C.REMOTE:
         plt.show()
     else:
         plt.close()
@@ -1149,3 +1176,66 @@ def _square_3d_plot(X, Y, Z, ax):
     ax.set_ylim(mid_y - max_range, mid_y + max_range)
     ax.set_zlim(mid_z - max_range, mid_z + max_range)
 
+
+def remove_tick_labels(ax, project_3d=False):
+    ax.set_xticklabels([])
+    ax.set_yticklabels([])
+    if project_3d:
+        ax.set_zticklabels([])
+    return ax
+
+
+def add_axes_arrows_3d(ax, arrow_length=10, xyz_colours=['r', 'g', 'b'], xyz_label=['X', 'Y', 'Z'], dist_arrow_text=3):
+    x_limits = ax.get_xlim3d()
+    y_limits = ax.get_ylim3d()
+    z_limits = ax.get_zlim3d()
+
+    x_len = x_limits[1] - x_limits[0]
+    y_len = y_limits[1] - y_limits[0]
+    z_len = z_limits[1] - z_limits[0]
+
+    ax.quiver(x_limits[0], y_limits[0], z_limits[0], x_len, 0, 0, color=xyz_colours[0], arrow_length_ratio=0) # (init_loc, end_loc, params)
+    ax.quiver(x_limits[0], y_limits[0], z_limits[0], 0, y_len, 0, color=xyz_colours[1], arrow_length_ratio=0)
+    ax.quiver(x_limits[0], y_limits[0], z_limits[0], 0, 0, z_len, color=xyz_colours[2], arrow_length_ratio=0)
+
+    # X axis
+    ax.quiver(x_limits[1], y_limits[0], z_limits[0], arrow_length, 0, 0, color=xyz_colours[0])
+    ax.text(x_limits[1] + arrow_length + dist_arrow_text, y_limits[0], z_limits[0], xyz_label[0], fontsize=20, ha='right', va='top')
+
+    # Y axis
+    ax.quiver(x_limits[0], y_limits[1], z_limits[0], 0, arrow_length, 0, color=xyz_colours[1])
+    ax.text(x_limits[0], y_limits[1] + arrow_length + dist_arrow_text, z_limits[0], xyz_label[0], fontsize=20, ha='left', va='top')
+
+    # Z axis
+    ax.quiver(x_limits[0], y_limits[0], z_limits[1], 0, 0, arrow_length, color=xyz_colours[2])
+    ax.text(x_limits[0], y_limits[0], z_limits[1] + arrow_length + dist_arrow_text, xyz_label[0], fontsize=20, ha='center', va='bottom')
+
+    return ax
+
+
+def add_axes_arrows_2d(ax, arrow_length=10, xy_colour=['r', 'g'], xy_label=['X', 'Y']):
+    x_limits = list(ax.get_xlim())
+    y_limits = list(ax.get_ylim())
+    origin = (x_limits[0], y_limits[1])
+
+    ax.annotate('', xy=(origin[0] + arrow_length, origin[1]), xytext=origin,
+                arrowprops=dict(headlength=8., headwidth=10., width=5., color=xy_colour[0]))
+    ax.annotate('', xy=(origin[0], origin[1] + arrow_length), xytext=origin,
+                arrowprops=dict(headlength=8., headwidth=10., width=5., color=xy_colour[0]))
+
+    ax.text(origin[0] + arrow_length, origin[1], xy_label[0], fontsize=25, ha='left', va='bottom')
+    ax.text(origin[0] - 1, origin[1] + arrow_length, xy_label[1], fontsize=25, ha='right', va='top')
+
+    return ax
+
+
+def set_axes_size(w,h, ax=None):
+    """ w, h: width, height in inches """
+    if not ax: ax=plt.gca()
+    l = ax.figure.subplotpars.left
+    r = ax.figure.subplotpars.right
+    t = ax.figure.subplotpars.top
+    b = ax.figure.subplotpars.bottom
+    figw = float(w)/(r-l)
+    figh = float(h)/(t-b)
+    ax.figure.set_size_inches(figw, figh)