--- a
+++ b/ext/lab2im/layers.py
@@ -0,0 +1,2060 @@
+"""
+This file regroups several custom keras layers used in the generation model:
+    - RandomSpatialDeformation,
+    - RandomCrop,
+    - RandomFlip,
+    - SampleConditionalGMM,
+    - SampleResolution,
+    - GaussianBlur,
+    - DynamicGaussianBlur,
+    - MimicAcquisition,
+    - BiasFieldCorruption,
+    - IntensityAugmentation,
+    - DiceLoss,
+    - WeightedL2Loss,
+    - ResetValuesToZero,
+    - ConvertLabels,
+    - PadAroundCentre,
+    - MaskEdges
+    - ImageGradients
+    - RandomDilationErosion
+
+
+If you use this code, please cite the first SynthSeg paper:
+https://github.com/BBillot/lab2im/blob/master/bibtex.bib
+
+Copyright 2020 Benjamin Billot
+
+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
+https://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.
+"""
+
+
+# python imports
+import keras
+import numpy as np
+import tensorflow as tf
+import keras.backend as K
+from keras.layers import Layer
+
+# project imports
+from ext.lab2im import utils
+from ext.lab2im import edit_tensors as l2i_et
+
+# third-party imports
+from ext.neuron import utils as nrn_utils
+import ext.neuron.layers as nrn_layers
+
+
+class RandomSpatialDeformation(Layer):
+    """This layer spatially deforms one or several tensors with a combination of affine and elastic transformations.
+    The input tensors are expected to have the same shape [batchsize, shape_dim1, ..., shape_dimn, channel].
+    The non-linear deformation is obtained by:
+    1) a small-size SVF is sampled from a centred normal distribution of random standard deviation.
+    2) it is resized with trilinear interpolation to half the shape of the input tensor
+    3) it is integrated to obtain a diffeomorphic transformation
+    4) finally, it is resized (again with trilinear interpolation) to full image size
+    :param scaling_bounds: (optional) range of the random scaling to apply. The scaling factor for each dimension is
+    sampled from a uniform distribution of predefined bounds. Can either be:
+    1) a number, in which case the scaling factor is independently sampled from the uniform distribution of bounds
+    [1-scaling_bounds, 1+scaling_bounds] for each dimension.
+    2) a sequence, in which case the scaling factor is sampled from the uniform distribution of bounds
+    (1-scaling_bounds[i], 1+scaling_bounds[i]) for the i-th dimension.
+    3) a numpy array of shape (2, n_dims), in which case the scaling factor is sampled from the uniform distribution
+     of bounds (scaling_bounds[0, i], scaling_bounds[1, i]) for the i-th dimension.
+    4) False, in which case scaling is completely turned off.
+    Default is scaling_bounds = 0.15 (case 1)
+    :param rotation_bounds: (optional) same as scaling bounds but for the rotation angle, except that for cases 1
+    and 2, the bounds are centred on 0 rather than 1, i.e. [0+rotation_bounds[i], 0-rotation_bounds[i]].
+    Default is rotation_bounds = 15.
+    :param shearing_bounds: (optional) same as scaling bounds. Default is shearing_bounds = 0.012.
+    :param translation_bounds: (optional) same as scaling bounds. Default is translation_bounds = False, but we
+    encourage using it when cropping is deactivated (i.e. when output_shape=None in BrainGenerator).
+    :param enable_90_rotations: (optional) whether to rotate the input by a random angle chosen in {0, 90, 180, 270}.
+    This is done regardless of the value of rotation_bounds. If true, a different value is sampled for each dimension.
+    :param nonlin_std: (optional) maximum value of the standard deviation of the normal distribution from which we
+    sample the small-size SVF. Set to 0 if you wish to completely turn the elastic deformation off.
+    :param nonlin_scale: (optional) if nonlin_std is not False, factor between the shapes of the input tensor
+    and the shape of the input non-linear tensor.
+    :param inter_method: (optional) interpolation method when deforming the input tensor. Can be 'linear', or 'nearest'
+    :param prob_deform: (optional) probability to apply spatial deformation
+    """
+
+    def __init__(self,
+                 scaling_bounds=0.15,
+                 rotation_bounds=10,
+                 shearing_bounds=0.02,
+                 translation_bounds=False,
+                 enable_90_rotations=False,
+                 nonlin_std=4.,
+                 nonlin_scale=.0625,
+                 inter_method='linear',
+                 prob_deform=1,
+                 **kwargs):
+
+        # shape attributes
+        self.n_inputs = 1
+        self.inshape = None
+        self.n_dims = None
+        self.small_shape = None
+
+        # deformation attributes
+        self.scaling_bounds = scaling_bounds
+        self.rotation_bounds = rotation_bounds
+        self.shearing_bounds = shearing_bounds
+        self.translation_bounds = translation_bounds
+        self.enable_90_rotations = enable_90_rotations
+        self.nonlin_std = nonlin_std
+        self.nonlin_scale = nonlin_scale
+
+        # boolean attributes
+        self.apply_affine_trans = (self.scaling_bounds is not False) | (self.rotation_bounds is not False) | \
+                                  (self.shearing_bounds is not False) | (self.translation_bounds is not False) | \
+                                  self.enable_90_rotations
+        self.apply_elastic_trans = self.nonlin_std > 0
+        self.prob_deform = prob_deform
+
+        # interpolation methods
+        self.inter_method = inter_method
+
+        super(RandomSpatialDeformation, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["scaling_bounds"] = self.scaling_bounds
+        config["rotation_bounds"] = self.rotation_bounds
+        config["shearing_bounds"] = self.shearing_bounds
+        config["translation_bounds"] = self.translation_bounds
+        config["enable_90_rotations"] = self.enable_90_rotations
+        config["nonlin_std"] = self.nonlin_std
+        config["nonlin_scale"] = self.nonlin_scale
+        config["inter_method"] = self.inter_method
+        config["prob_deform"] = self.prob_deform
+        return config
+
+    def build(self, input_shape):
+
+        if not isinstance(input_shape, list):
+            inputshape = [input_shape]
+        else:
+            self.n_inputs = len(input_shape)
+            inputshape = input_shape
+        self.inshape = inputshape[0][1:]
+        self.n_dims = len(self.inshape) - 1
+
+        if self.apply_elastic_trans:
+            self.small_shape = utils.get_resample_shape(self.inshape[:self.n_dims],
+                                                        self.nonlin_scale, self.n_dims)
+        else:
+            self.small_shape = None
+
+        self.inter_method = utils.reformat_to_list(self.inter_method, length=self.n_inputs, dtype='str')
+
+        self.built = True
+        super(RandomSpatialDeformation, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # reformat inputs and get its shape
+        if self.n_inputs < 2:
+            inputs = [inputs]
+        types = [v.dtype for v in inputs]
+        inputs = [tf.cast(v, dtype='float32') for v in inputs]
+        batchsize = tf.split(tf.shape(inputs[0]), [1, self.n_dims + 1])[0]
+
+        # initialise list of transforms to operate
+        list_trans = list()
+
+        # add affine deformation to inputs list
+        if self.apply_affine_trans:
+            affine_trans = utils.sample_affine_transform(batchsize,
+                                                         self.n_dims,
+                                                         self.rotation_bounds,
+                                                         self.scaling_bounds,
+                                                         self.shearing_bounds,
+                                                         self.translation_bounds,
+                                                         self.enable_90_rotations)
+            list_trans.append(affine_trans)
+
+        # prepare non-linear deformation field and add it to inputs list
+        if self.apply_elastic_trans:
+
+            # sample small field from normal distribution of specified std dev
+            trans_shape = tf.concat([batchsize, tf.convert_to_tensor(self.small_shape, dtype='int32')], axis=0)
+            trans_std = tf.random.uniform((1, 1), maxval=self.nonlin_std)
+            elastic_trans = tf.random.normal(trans_shape, stddev=trans_std)
+
+            # reshape this field to half size (for smoother SVF), integrate it, and reshape to full image size
+            resize_shape = [max(int(self.inshape[i] / 2), self.small_shape[i]) for i in range(self.n_dims)]
+            elastic_trans = nrn_layers.Resize(size=resize_shape, interp_method='linear')(elastic_trans)
+            elastic_trans = nrn_layers.VecInt()(elastic_trans)
+            elastic_trans = nrn_layers.Resize(size=self.inshape[:self.n_dims], interp_method='linear')(elastic_trans)
+            list_trans.append(elastic_trans)
+
+        # apply deformations and return tensors with correct dtype
+        if self.apply_affine_trans | self.apply_elastic_trans:
+            if self.prob_deform == 1:
+                inputs = [nrn_layers.SpatialTransformer(m)([v] + list_trans) for (m, v) in
+                          zip(self.inter_method, inputs)]
+            else:
+                rand_trans = tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_deform))
+                inputs = [K.switch(rand_trans, nrn_layers.SpatialTransformer(m)([v] + list_trans), v)
+                          for (m, v) in zip(self.inter_method, inputs)]
+        if self.n_inputs < 2:
+            return tf.cast(inputs[0], types[0])
+        else:
+            return [tf.cast(v, t) for (t, v) in zip(types, inputs)]
+
+
+class RandomCrop(Layer):
+    """Randomly crop all input tensors to a given shape. This cropping is applied to all channels.
+    The input tensors are expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
+    :param crop_shape: list with cropping shape in each dimension (excluding batch and channel dimension)
+
+    example:
+    if input is a tensor of shape [batchsize, 160, 160, 160, 3],
+    output = RandomCrop(crop_shape=[96, 128, 96])(input)
+    will yield an output of shape [batchsize, 96, 128, 96, 3] that is obtained by cropping with randomly selected
+    cropping indices.
+    """
+
+    def __init__(self, crop_shape, **kwargs):
+
+        self.several_inputs = True
+        self.crop_max_val = None
+        self.crop_shape = crop_shape
+        self.n_dims = len(crop_shape)
+        self.list_n_channels = None
+        super(RandomCrop, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["crop_shape"] = self.crop_shape
+        return config
+
+    def build(self, input_shape):
+
+        if not isinstance(input_shape, list):
+            self.several_inputs = False
+            inputshape = [input_shape]
+        else:
+            inputshape = input_shape
+        self.crop_max_val = np.array(np.array(inputshape[0][1:self.n_dims + 1])) - np.array(self.crop_shape)
+        self.list_n_channels = [i[-1] for i in inputshape]
+        self.built = True
+        super(RandomCrop, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # if one input only is provided, performs the cropping directly
+        if not self.several_inputs:
+            return tf.map_fn(self._single_slice, inputs, dtype=inputs.dtype)
+
+        # otherwise we concatenate all inputs before cropping, so that they are all cropped at the same location
+        else:
+            types = [v.dtype for v in inputs]
+            inputs = tf.concat([tf.cast(v, 'float32') for v in inputs], axis=-1)
+            inputs = tf.map_fn(self._single_slice, inputs, dtype=tf.float32)
+            inputs = tf.split(inputs, self.list_n_channels, axis=-1)
+            return [tf.cast(v, t) for (t, v) in zip(types, inputs)]
+
+    def _single_slice(self, vol):
+        crop_idx = tf.cast(tf.random.uniform([self.n_dims], 0, np.array(self.crop_max_val), 'float32'), dtype='int32')
+        crop_idx = tf.concat([crop_idx, tf.zeros([1], dtype='int32')], axis=0)
+        crop_size = tf.convert_to_tensor(self.crop_shape + [-1], dtype='int32')
+        return tf.slice(vol, begin=crop_idx, size=crop_size)
+
+    def compute_output_shape(self, input_shape):
+        output_shape = [tuple([None] + self.crop_shape + [v]) for v in self.list_n_channels]
+        return output_shape if self.several_inputs else output_shape[0]
+
+
+class RandomFlip(Layer):
+    """This layer randomly flips the input tensor along the specified axes with a specified probability.
+    It can also take multiple tensors as inputs (if they have the same shape). The same flips will be applied to all
+    input tensors. These are expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
+    If specified, this layer can also swap corresponding values. This is especially useful when flipping label maps
+    with different labels for right/left structures, such that the flipped label maps keep a consistent labelling.
+    :param axis: integer, or list of integers specifying the dimensions along which to flip.
+    If a list, the input tensors can be flipped simultaneously in several directions. The values in flip_axis exclude
+    the batch dimension (e.g. 0 will flip the tensor along the first axis after the batch dimension).
+    Default is None, where the tensors can be flipped along all axes (except batch and channel axes).
+    :param swap_labels: boolean to specify whether to swap the values of each input. Values are only swapped if an odd
+    number of flips is applied.
+    Can also be a list if several tensors are given as input.
+    All the inputs for which the values need to be swapped must be int32 or int64.
+    :param label_list: if swap_labels is True, list of all labels contained in labels. Must be ordered as follows, first
+     the neutral labels (i.e. non-sided), then left labels and right labels.
+    :param n_neutral_labels: if swap_labels is True, number of non-sided labels
+    :param prob: probability to flip along each specified axis
+
+    example 1:
+    if input is a tensor of shape (batchsize, 10, 100, 200, 3)
+    output = RandomFlip()(input) will randomly flip input along one of the 1st, 2nd, or 3rd axis (i.e. those with shape
+    10, 100, 200).
+
+    example 2:
+    if input is a tensor of shape (batchsize, 10, 100, 200, 3)
+    output = RandomFlip(flip_axis=1)(input) will randomly flip input along the 3rd axis (with shape 100), i.e. the axis
+    with index 1 if we don't count the batch axis.
+
+    example 3:
+    input = tf.convert_to_tensor(np.array([[1, 0, 0, 0, 0, 0, 0],
+                                           [1, 0, 0, 0, 2, 2, 0],
+                                           [1, 0, 0, 0, 2, 2, 0],
+                                           [1, 0, 0, 0, 2, 2, 0],
+                                           [1, 0, 0, 0, 0, 0, 0]]))
+    label_list = np.array([0, 1, 2])
+    n_neutral_labels = 1
+    output = RandomFlip(flip_axis=1, swap_labels=True, label_list=label_list, n_neutral_labels=n_neutral_labels)(input)
+    where output will either be equal to input (bear in mind the flipping occurs with a 0.5 probability), or:
+    output = [[0, 0, 0, 0, 0, 0, 2],
+              [0, 1, 1, 0, 0, 0, 2],
+              [0, 1, 1, 0, 0, 0, 2],
+              [0, 1, 1, 0, 0, 0, 2],
+              [0, 0, 0, 0, 0, 0, 2]]
+    Note that the input must have a dtype int32 or int64 for its values to be swapped, otherwise an error will be raised
+
+    example 4:
+    if labels is the same as in the input of example 3, and image is a float32 image, then we can swap consistently both
+    the labels and the image with:
+    labels, image = RandomFlip(flip_axis=1, swap_labels=[True, False], label_list=label_list,
+                               n_neutral_labels=n_neutral_labels)([labels, image]])
+    Note that the labels must have a dtype int32 or int64 to be swapped, otherwise an error will be raised.
+    This doesn't concern the image input, as its values are not swapped.
+    """
+
+    def __init__(self, axis=None, swap_labels=False, label_list=None, n_neutral_labels=None, prob=0.5, **kwargs):
+
+        # shape attributes
+        self.several_inputs = True
+        self.n_dims = None
+        self.list_n_channels = None
+
+        # axis along which to flip
+        self.axis = utils.reformat_to_list(axis)
+        self.flip_axes = None
+
+        # whether to swap labels, and corresponding label list
+        self.swap_labels = utils.reformat_to_list(swap_labels)
+        self.label_list = label_list
+        self.n_neutral_labels = n_neutral_labels
+        self.swap_lut = None
+
+        self.prob = prob
+
+        super(RandomFlip, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["axis"] = self.axis
+        config["swap_labels"] = self.swap_labels
+        config["label_list"] = self.label_list
+        config["n_neutral_labels"] = self.n_neutral_labels
+        config["prob"] = self.prob
+        return config
+
+    def build(self, input_shape):
+
+        if not isinstance(input_shape, list):
+            self.several_inputs = False
+            inputshape = [input_shape]
+        else:
+            inputshape = input_shape
+        self.n_dims = len(inputshape[0][1:-1])
+        self.list_n_channels = [i[-1] for i in inputshape]
+        self.swap_labels = utils.reformat_to_list(self.swap_labels, length=len(inputshape))
+        self.flip_axes = np.arange(self.n_dims).tolist() if self.axis is None else self.axis
+
+        # create label list with swapped labels
+        if any(self.swap_labels):
+            assert (self.label_list is not None) & (self.n_neutral_labels is not None), \
+                'please provide a label_list, and n_neutral_labels when swapping the values of at least one input'
+            n_labels = len(self.label_list)
+            if self.n_neutral_labels == n_labels:
+                self.swap_labels = [False] * len(self.swap_labels)
+            else:
+                rl_split = np.split(self.label_list, [self.n_neutral_labels,
+                                                      self.n_neutral_labels + int((n_labels-self.n_neutral_labels)/2)])
+                label_list_swap = np.concatenate((rl_split[0], rl_split[2], rl_split[1]))
+                swap_lut = utils.get_mapping_lut(self.label_list, label_list_swap)
+                self.swap_lut = tf.convert_to_tensor(swap_lut, dtype='int32')
+
+        self.built = True
+        super(RandomFlip, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # convert inputs to list, and get each input type
+        inputs = [inputs] if not self.several_inputs else inputs
+        types = [v.dtype for v in inputs]
+
+        # store whether to flip along each specified dimension
+        batchsize = tf.split(tf.shape(inputs[0]), [1, self.n_dims + 1])[0]
+        size = tf.concat([batchsize, len(self.flip_axes) * tf.ones(1, dtype='int32')], axis=0)
+        rand_flip = K.less(tf.random.uniform(size, 0, 1), self.prob)
+
+        # swap right/left labels if we apply an odd number of flips
+        odd = tf.math.floormod(tf.reduce_sum(tf.cast(rand_flip, 'int32'), -1, keepdims=True), 2) != 0
+        swapped_inputs = list()
+        for i in range(len(inputs)):
+            if self.swap_labels[i]:
+                swapped_inputs.append(tf.map_fn(self._single_swap, [inputs[i], odd], dtype=types[i]))
+            else:
+                swapped_inputs.append(inputs[i])
+
+        # flip inputs and convert them back to their original type
+        inputs = tf.concat([tf.cast(v, 'float32') for v in swapped_inputs], axis=-1)
+        inputs = tf.map_fn(self._single_flip, [inputs, rand_flip], dtype=tf.float32)
+        inputs = tf.split(inputs, self.list_n_channels, axis=-1)
+
+        if self.several_inputs:
+            return [tf.cast(v, t) for (t, v) in zip(types, inputs)]
+        else:
+            return tf.cast(inputs[0], types[0])
+
+    def _single_swap(self, inputs):
+        return K.switch(inputs[1], tf.gather(self.swap_lut, inputs[0]), inputs[0])
+
+    @staticmethod
+    def _single_flip(inputs):
+        flip_axis = tf.where(inputs[1])
+        return K.switch(tf.equal(tf.size(flip_axis), 0), inputs[0], tf.reverse(inputs[0], axis=flip_axis[..., 0]))
+
+
+class SampleConditionalGMM(Layer):
+    """This layer generates an image by sampling a Gaussian Mixture Model conditioned on a label map given as input.
+    The parameters of the GMM are given as two additional inputs to the layer (means and standard deviations):
+    image = SampleConditionalGMM(generation_labels)([label_map, means, stds])
+
+    :param generation_labels: list of all possible label values contained in the input label maps.
+    Must be a list or a 1D numpy array of size N, where N is the total number of possible label values.
+
+    Layer inputs:
+    label_map: input label map of shape [batchsize, shape_dim1, ..., shape_dimn, n_channel].
+    All the values of label_map must be contained in generation_labels, but the input label_map doesn't necessarily have
+    to contain all the values in generation_labels.
+    means: tensor containing the mean values of all Gaussian distributions of the GMM.
+           It must be of shape [batchsize, N, n_channel], and in the same order as generation label,
+           i.e. the ith value of generation_labels will be associated to the ith value of means.
+    stds: same as means but for the standard deviations of the GMM.
+    """
+
+    def __init__(self, generation_labels, **kwargs):
+        self.generation_labels = generation_labels
+        self.n_labels = None
+        self.n_channels = None
+        self.max_label = None
+        self.indices = None
+        self.shape = None
+        super(SampleConditionalGMM, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["generation_labels"] = self.generation_labels
+        return config
+
+    def build(self, input_shape):
+
+        # check n_labels and n_channels
+        assert len(input_shape) == 3, 'should have three inputs: labels, means, std devs (in that order).'
+        self.n_channels = input_shape[1][-1]
+        self.n_labels = len(self.generation_labels)
+        assert self.n_labels == input_shape[1][1], 'means should have the same number of values as generation_labels'
+        assert self.n_labels == input_shape[2][1], 'stds should have the same number of values as generation_labels'
+
+        # scatter parameters (to build mean/std lut)
+        self.max_label = np.max(self.generation_labels) + 1
+        indices = np.concatenate([self.generation_labels + self.max_label * i for i in range(self.n_channels)], axis=-1)
+        self.shape = tf.convert_to_tensor([np.max(indices) + 1], dtype='int32')
+        self.indices = tf.convert_to_tensor(utils.add_axis(indices, axis=[0, -1]), dtype='int32')
+
+        self.built = True
+        super(SampleConditionalGMM, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # reformat labels and scatter indices
+        batch = tf.split(tf.shape(inputs[0]), [1, -1])[0]
+        tmp_indices = tf.tile(self.indices, tf.concat([batch, tf.convert_to_tensor([1, 1], dtype='int32')], axis=0))
+        labels = tf.concat([tf.cast(inputs[0], dtype='int32') + self.max_label * i for i in range(self.n_channels)], -1)
+
+        # build mean map
+        means = tf.concat([inputs[1][..., i] for i in range(self.n_channels)], 1)
+        tile_shape = tf.concat([batch, tf.convert_to_tensor([1, ], dtype='int32')], axis=0)
+        means = tf.tile(tf.expand_dims(tf.scatter_nd(tmp_indices, means, self.shape), 0), tile_shape)
+        means_map = tf.map_fn(lambda x: tf.gather(x[0], x[1]), [means, labels], dtype=tf.float32)
+
+        # same for stds
+        stds = tf.concat([inputs[2][..., i] for i in range(self.n_channels)], 1)
+        stds = tf.tile(tf.expand_dims(tf.scatter_nd(tmp_indices, stds, self.shape), 0), tile_shape)
+        stds_map = tf.map_fn(lambda x: tf.gather(x[0], x[1]), [stds, labels], dtype=tf.float32)
+
+        return stds_map * tf.random.normal(tf.shape(labels)) + means_map
+
+    def compute_output_shape(self, input_shape):
+        return input_shape[0] if (self.n_channels == 1) else tuple(list(input_shape[0][:-1]) + [self.n_channels])
+
+
+class SampleResolution(Layer):
+    """Build a random resolution tensor by sampling a uniform distribution of provided range.
+
+    You can use this layer in the following ways:
+        resolution = SampleConditionalGMM(min_resolution)() in this case resolution will be a tensor of shape (n_dims,),
+        where n_dims is the length of the min_resolution parameter (provided as a list, see below).
+        resolution = SampleConditionalGMM(min_resolution)(input), where input is a tensor for which the first dimension
+        represents the batch_size. In this case resolution will be a tensor of shape (batchsize, n_dims,).
+
+    :param min_resolution: list of length n_dims specifying the inferior bounds of the uniform distributions to
+    sample from for each value.
+    :param max_res_iso: If not None, all the values of resolution will be equal to the same value, which is randomly
+    sampled at each minibatch in U(min_resolution, max_res_iso).
+    :param max_res_aniso: If not None, we first randomly select a direction i in the range [0, n_dims-1], and we sample
+    a value in the corresponding uniform distribution U(min_resolution[i], max_res_aniso[i]).
+    The other values of resolution will be set to min_resolution.
+    :param prob_iso: if both max_res_iso and max_res_aniso are specified, this allows to specify the probability of
+    sampling an isotropic resolution (therefore using max_res_iso) with respect to anisotropic resolution
+    (which would use max_res_aniso).
+    :param prob_min: if not zero, this allows to return with the specified probability an output resolution equal
+    to min_resolution.
+    :param return_thickness: if set to True, this layer will also return a thickness value of the same shape as
+    resolution, which will be sampled independently for each axis from the uniform distribution
+    U(min_resolution, resolution).
+
+    """
+
+    def __init__(self,
+                 min_resolution,
+                 max_res_iso=None,
+                 max_res_aniso=None,
+                 prob_iso=0.1,
+                 prob_min=0.05,
+                 return_thickness=True,
+                 **kwargs):
+
+        self.min_res = min_resolution
+        self.max_res_iso_input = max_res_iso
+        self.max_res_iso = None
+        self.max_res_aniso_input = max_res_aniso
+        self.max_res_aniso = None
+        self.prob_iso = prob_iso
+        self.prob_min = prob_min
+        self.return_thickness = return_thickness
+        self.n_dims = len(self.min_res)
+        self.add_batchsize = False
+        self.min_res_tens = None
+        super(SampleResolution, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["min_resolution"] = self.min_res
+        config["max_res_iso"] = self.max_res_iso
+        config["max_res_aniso"] = self.max_res_aniso
+        config["prob_iso"] = self.prob_iso
+        config["prob_min"] = self.prob_min
+        config["return_thickness"] = self.return_thickness
+        return config
+
+    def build(self, input_shape):
+
+        # check maximum resolutions
+        assert ((self.max_res_iso_input is not None) | (self.max_res_aniso_input is not None)), \
+            'at least one of maximum isotropic or anisotropic resolutions must be provided, received none'
+
+        # reformat resolutions as numpy arrays
+        self.min_res = np.array(self.min_res)
+        if self.max_res_iso_input is not None:
+            self.max_res_iso = np.array(self.max_res_iso_input)
+            assert len(self.min_res) == len(self.max_res_iso), \
+                'min and isotropic max resolution must have the same length, ' \
+                'had {0} and {1}'.format(self.min_res, self.max_res_iso)
+            if np.array_equal(self.min_res, self.max_res_iso):
+                self.max_res_iso = None
+        if self.max_res_aniso_input is not None:
+            self.max_res_aniso = np.array(self.max_res_aniso_input)
+            assert len(self.min_res) == len(self.max_res_aniso), \
+                'min and anisotropic max resolution must have the same length, ' \
+                'had {} and {}'.format(self.min_res, self.max_res_aniso)
+            if np.array_equal(self.min_res, self.max_res_aniso):
+                self.max_res_aniso = None
+
+        # check prob iso
+        if (self.max_res_iso is not None) & (self.max_res_aniso is not None) & (self.prob_iso == 0):
+            raise Exception('prob iso is 0 while sampling either isotropic and anisotropic resolutions is enabled')
+
+        if input_shape:
+            self.add_batchsize = True
+
+        self.min_res_tens = tf.convert_to_tensor(self.min_res, dtype='float32')
+
+        self.built = True
+        super(SampleResolution, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        if not self.add_batchsize:
+            shape = [self.n_dims]
+            dim = tf.random.uniform(shape=(1, 1), minval=0, maxval=self.n_dims, dtype='int32')
+            mask = tf.tensor_scatter_nd_update(tf.zeros([self.n_dims], dtype='bool'), dim,
+                                               tf.convert_to_tensor([True], dtype='bool'))
+        else:
+            batch = tf.split(tf.shape(inputs), [1, -1])[0]
+            tile_shape = tf.concat([batch, tf.convert_to_tensor([1], dtype='int32')], axis=0)
+            self.min_res_tens = tf.tile(tf.expand_dims(self.min_res_tens, 0), tile_shape)
+
+            shape = tf.concat([batch, tf.convert_to_tensor([self.n_dims], dtype='int32')], axis=0)
+            indices = tf.stack([tf.range(0, batch[0]), tf.random.uniform(batch, 0, self.n_dims, dtype='int32')], 1)
+            mask = tf.tensor_scatter_nd_update(tf.zeros(shape, dtype='bool'), indices, tf.ones(batch, dtype='bool'))
+
+        # return min resolution as tensor if min=max
+        if (self.max_res_iso is None) & (self.max_res_aniso is None):
+            new_resolution = self.min_res_tens
+
+        # sample isotropic resolution only
+        elif (self.max_res_iso is not None) & (self.max_res_aniso is None):
+            new_resolution_iso = tf.random.uniform(shape, minval=self.min_res, maxval=self.max_res_iso)
+            new_resolution = K.switch(tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_min)),
+                                      self.min_res_tens,
+                                      new_resolution_iso)
+
+        # sample anisotropic resolution only
+        elif (self.max_res_iso is None) & (self.max_res_aniso is not None):
+            new_resolution_aniso = tf.random.uniform(shape, minval=self.min_res, maxval=self.max_res_aniso)
+            new_resolution = K.switch(tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_min)),
+                                      self.min_res_tens,
+                                      tf.where(mask, new_resolution_aniso, self.min_res_tens))
+
+        # sample either anisotropic or isotropic resolution
+        else:
+            new_resolution_iso = tf.random.uniform(shape, minval=self.min_res, maxval=self.max_res_iso)
+            new_resolution_aniso = tf.random.uniform(shape, minval=self.min_res, maxval=self.max_res_aniso)
+            new_resolution = K.switch(tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_iso)),
+                                      new_resolution_iso,
+                                      tf.where(mask, new_resolution_aniso, self.min_res_tens))
+            new_resolution = K.switch(tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_min)),
+                                      self.min_res_tens,
+                                      new_resolution)
+
+        if self.return_thickness:
+            return [new_resolution, tf.random.uniform(tf.shape(self.min_res_tens), self.min_res_tens, new_resolution)]
+        else:
+            return new_resolution
+
+    def compute_output_shape(self, input_shape):
+        if self.return_thickness:
+            return [(None, self.n_dims)] * 2 if self.add_batchsize else [self.n_dims] * 2
+        else:
+            return (None, self.n_dims) if self.add_batchsize else self.n_dims
+
+
+class GaussianBlur(Layer):
+    """Applies gaussian blur to an input image.
+    The input image is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
+    :param sigma: standard deviation of the blurring kernels to apply. Can be a number, a list of length n_dims, or
+    a numpy array.
+    :param random_blur_range: (optional) if not None, this introduces a randomness in the blurring kernels, where
+    sigma is now multiplied by a coefficient dynamically sampled from a uniform distribution with bounds
+    [1/random_blur_range, random_blur_range].
+    :param use_mask: (optional) whether a mask of the input will be provided as an additional layer input. This is used
+    to mask the blurred image, and to correct for edge blurring effects.
+
+    example 1:
+    output = GaussianBlur(sigma=0.5)(input) will isotropically blur the input with a gaussian kernel of std 0.5.
+
+    example 2:
+    if input is a tensor of shape [batchsize, 10, 100, 200, 2]
+    output = GaussianBlur(sigma=[0.5, 1, 10])(input) will blur the input a different gaussian kernel in each dimension.
+
+    example 3:
+    output = GaussianBlur(sigma=0.5, random_blur_range=1.15)(input)
+    will blur the input a different gaussian kernel in each dimension, as each dimension will be associated with
+    a kernel, whose standard deviation will be uniformly sampled from [0.5/1.15; 0.5*1.15].
+
+    example 4:
+    output = GaussianBlur(sigma=0.5, use_mask=True)([input, mask])
+    will 1) blur the input a different gaussian kernel in each dimension, 2) mask the blurred image with the provided
+    mask, and 3) correct for edge blurring effects. If the provided mask is not of boolean type, it will be thresholded
+    above positive values.
+    """
+
+    def __init__(self, sigma, random_blur_range=None, use_mask=False, **kwargs):
+        self.sigma = utils.reformat_to_list(sigma)
+        assert np.all(np.array(self.sigma) >= 0), 'sigma should be superior or equal to 0'
+        self.use_mask = use_mask
+
+        self.n_dims = None
+        self.n_channels = None
+        self.blur_range = random_blur_range
+        self.stride = None
+        self.separable = None
+        self.kernels = None
+        self.convnd = None
+        super(GaussianBlur, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["sigma"] = self.sigma
+        config["random_blur_range"] = self.blur_range
+        config["use_mask"] = self.use_mask
+        return config
+
+    def build(self, input_shape):
+
+        # get shapes
+        if self.use_mask:
+            assert len(input_shape) == 2, 'please provide a mask as second layer input when use_mask=True'
+            self.n_dims = len(input_shape[0]) - 2
+            self.n_channels = input_shape[0][-1]
+        else:
+            self.n_dims = len(input_shape) - 2
+            self.n_channels = input_shape[-1]
+
+        # prepare blurring kernel
+        self.stride = [1] * (self.n_dims + 2)
+        self.sigma = utils.reformat_to_list(self.sigma, length=self.n_dims)
+        self.separable = np.linalg.norm(np.array(self.sigma)) > 5
+        if self.blur_range is None:  # fixed kernels
+            self.kernels = l2i_et.gaussian_kernel(self.sigma, separable=self.separable)
+        else:
+            self.kernels = None
+
+        # prepare convolution
+        self.convnd = getattr(tf.nn, 'conv%dd' % self.n_dims)
+
+        self.built = True
+        super(GaussianBlur, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        if self.use_mask:
+            image = inputs[0]
+            mask = tf.cast(inputs[1], 'bool')
+        else:
+            image = inputs
+            mask = None
+
+        # redefine the kernels at each new step when blur_range is activated
+        if self.blur_range is not None:
+            self.kernels = l2i_et.gaussian_kernel(self.sigma, blur_range=self.blur_range, separable=self.separable)
+
+        if self.separable:
+            for k in self.kernels:
+                if k is not None:
+                    image = tf.concat([self.convnd(tf.expand_dims(image[..., n], -1), k, self.stride, 'SAME')
+                                       for n in range(self.n_channels)], -1)
+                    if self.use_mask:
+                        maskb = tf.cast(mask, 'float32')
+                        maskb = tf.concat([self.convnd(tf.expand_dims(maskb[..., n], -1), k, self.stride, 'SAME')
+                                           for n in range(self.n_channels)], -1)
+                        image = image / (maskb + K.epsilon())
+                        image = tf.where(mask, image, tf.zeros_like(image))
+        else:
+            if any(self.sigma):
+                image = tf.concat([self.convnd(tf.expand_dims(image[..., n], -1), self.kernels, self.stride, 'SAME')
+                                   for n in range(self.n_channels)], -1)
+                if self.use_mask:
+                    maskb = tf.cast(mask, 'float32')
+                    maskb = tf.concat([self.convnd(tf.expand_dims(maskb[..., n], -1), self.kernels, self.stride, 'SAME')
+                                       for n in range(self.n_channels)], -1)
+                    image = image / (maskb + K.epsilon())
+                    image = tf.where(mask, image, tf.zeros_like(image))
+
+        return image
+
+
+class DynamicGaussianBlur(Layer):
+    """Applies gaussian blur to an input image, where the standard deviation of the blurring kernel is provided as a
+    layer input, which enables to perform dynamic blurring (i.e. the blurring kernel can vary at each minibatch).
+    :param max_sigma: maximum value of the standard deviation that will be provided as input. This is used to compute
+    the size of the blurring kernels. It must be provided as a list of length n_dims.
+    :param random_blur_range: (optional) if not None, this introduces a randomness in the blurring kernels, where
+    sigma is now multiplied by a coefficient dynamically sampled from a uniform distribution with bounds
+    [1/random_blur_range, random_blur_range].
+
+    example:
+    blurred_image = DynamicGaussianBlur(max_sigma=[5.]*3, random_blurring_range=1.15)([image, sigma])
+    will return a blurred version of image, where the standard deviation of each dimension (given as a tensor, and with
+    values lower than 5 for each axis) is multiplied by a random coefficient uniformly sampled from [1/1.15; 1.15].
+    """
+
+    def __init__(self, max_sigma, random_blur_range=None, **kwargs):
+        self.max_sigma = max_sigma
+        self.n_dims = None
+        self.n_channels = None
+        self.convnd = None
+        self.blur_range = random_blur_range
+        self.separable = None
+        super(DynamicGaussianBlur, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["max_sigma"] = self.max_sigma
+        config["random_blur_range"] = self.blur_range
+        return config
+
+    def build(self, input_shape):
+        assert len(input_shape) == 2, 'sigma should be provided as an input tensor for dynamic blurring'
+        self.n_dims = len(input_shape[0]) - 2
+        self.n_channels = input_shape[0][-1]
+        self.convnd = getattr(tf.nn, 'conv%dd' % self.n_dims)
+        self.max_sigma = utils.reformat_to_list(self.max_sigma, length=self.n_dims)
+        self.separable = np.linalg.norm(np.array(self.max_sigma)) > 5
+        self.built = True
+        super(DynamicGaussianBlur, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+        image = inputs[0]
+        sigma = inputs[-1]
+        kernels = l2i_et.gaussian_kernel(sigma, self.max_sigma, self.blur_range, self.separable)
+        if self.separable:
+            for kernel in kernels:
+                image = tf.map_fn(self._single_blur, [image, kernel], dtype=tf.float32)
+        else:
+            image = tf.map_fn(self._single_blur, [image, kernels], dtype=tf.float32)
+        return image
+
+    def _single_blur(self, inputs):
+        if self.n_channels > 1:
+            split_channels = tf.split(inputs[0], [1] * self.n_channels, axis=-1)
+            blurred_channel = list()
+            for channel in split_channels:
+                blurred = self.convnd(tf.expand_dims(channel, 0), inputs[1], [1] * (self.n_dims + 2), padding='SAME')
+                blurred_channel.append(tf.squeeze(blurred, axis=0))
+            output = tf.concat(blurred_channel, -1)
+        else:
+            output = self.convnd(tf.expand_dims(inputs[0], 0), inputs[1], [1] * (self.n_dims + 2), padding='SAME')
+            output = tf.squeeze(output, axis=0)
+        return output
+
+
+class MimicAcquisition(Layer):
+    """
+    Layer that takes an image as input, and simulates data that has been acquired at low resolution.
+    The output is obtained by resampling the input twice:
+     - first at a resolution given as an input (i.e. the "acquisition" resolution),
+     - then at the output resolution (specified output shape).
+    The input tensor is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
+
+    :param volume_res: resolution of the provided inputs. Must be a 1-D numpy array with n_dims elements.
+    :param min_subsample_res: lower bound of the acquisition resolutions to mimic (i.e. the input resolution must have
+    values higher than min-subsample_res).
+    :param resample_shape: shape of the output tensor
+    :param build_dist_map: whether to return distance maps as outputs. These indicate the distance between each voxel
+    and the nearest non-interpolated voxel (during the second resampling).
+    :param prob_noise: probability to apply noise injection
+
+    example 1:
+    im_res = [1., 1., 1.]
+    low_res = [1., 1., 3.]
+    res = tf.convert_to_tensor([1., 1., 4.5])
+    image is a tensor of shape (None, 256, 256, 256, 3)
+    resample_shape = [256, 256, 256]
+    output = MimicAcquisition(im_res, low_res, resample_shape)([image, res])
+    output will be a tensor of shape (None, 256, 256, 256, 3), obtained by downsampling image to [1., 1., 4.5].
+    and re-upsampling it at initial resolution (because resample_shape is equal to the input shape). In this example all
+    examples of the batch will be downsampled to the same resolution (because res has no batch dimension).
+    Note that the provided res must have higher values than min_low_res.
+
+    example 2:
+    im_res = [1., 1., 1.]
+    min_low_res = [1., 1., 1.]
+    res is a tensor of shape (None, 3), obtained for example by using the SampleResolution layer (see above).
+    image is a tensor of shape (None, 256, 256, 256, 1)
+    resample_shape = [128, 128, 128]
+    output = MimicAcquisition(im_res, low_res, resample_shape)([image, res])
+    output will be a tensor of shape (None, 128, 128, 128, 1), obtained by downsampling each examples of the batch to
+    the matching resolution in res, and resampling them all to half the initial resolution.
+    Note that the provided res must have higher values than min_low_res.
+    """
+
+    def __init__(self, volume_res, min_subsample_res, resample_shape, build_dist_map=False,
+                 noise_std=0, prob_noise=0.95, **kwargs):
+
+        # resolutions and dimensions
+        self.volume_res = volume_res
+        self.min_subsample_res = min_subsample_res
+        self.n_dims = len(self.volume_res)
+        self.n_channels = None
+        self.add_batchsize = None
+
+        # noise
+        self.noise_std = noise_std
+        self.prob_noise = prob_noise
+
+        # input and output shapes
+        self.inshape = None
+        self.resample_shape = resample_shape
+
+        # meshgrids for resampling
+        self.down_grid = None
+        self.up_grid = None
+
+        # whether to return a map indicating the distance from the interpolated voxels, to acquired ones.
+        self.build_dist_map = build_dist_map
+
+        super(MimicAcquisition, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["volume_res"] = self.volume_res
+        config["min_subsample_res"] = self.min_subsample_res
+        config["resample_shape"] = self.resample_shape
+        config["build_dist_map"] = self.build_dist_map
+        config["noise_std"] = self.noise_std
+        config["prob_noise"] = self.prob_noise
+        return config
+
+    def build(self, input_shape):
+
+        # set up input shape and acquisition shape
+        self.inshape = input_shape[0][1:]
+        self.n_channels = input_shape[0][-1]
+        self.add_batchsize = False if (input_shape[1][0] is None) else True
+        down_tensor_shape = np.int32(np.array(self.inshape[:-1]) * self.volume_res / self.min_subsample_res)
+
+        # build interpolation meshgrids
+        self.down_grid = tf.expand_dims(tf.stack(nrn_utils.volshape_to_ndgrid(down_tensor_shape), -1), axis=0)
+        self.up_grid = tf.expand_dims(tf.stack(nrn_utils.volshape_to_ndgrid(self.resample_shape), -1), axis=0)
+
+        self.built = True
+        super(MimicAcquisition, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # sort inputs
+        assert len(inputs) == 2, 'inputs must have two items, the tensor to resample, and the downsampling resolution'
+        vol = inputs[0]
+        subsample_res = tf.cast(inputs[1], dtype='float32')
+        vol = K.reshape(vol, [-1, *self.inshape])  # necessary for multi_gpu models
+        batchsize = tf.split(tf.shape(vol), [1, -1])[0]
+        tile_shape = tf.concat([batchsize, tf.ones([1], dtype='int32')], 0)
+
+        # get downsampling and upsampling factors
+        if self.add_batchsize:
+            subsample_res = tf.tile(tf.expand_dims(subsample_res, 0), tile_shape)
+        down_shape = tf.cast(tf.convert_to_tensor(np.array(self.inshape[:-1]) * self.volume_res, dtype='float32') /
+                             subsample_res, dtype='int32')
+        down_zoom_factor = tf.cast(down_shape / tf.convert_to_tensor(self.inshape[:-1]), dtype='float32')
+        up_zoom_factor = tf.cast(tf.convert_to_tensor(self.resample_shape, dtype='int32') / down_shape, dtype='float32')
+
+        # downsample
+        down_loc = tf.tile(self.down_grid, tf.concat([batchsize, tf.ones([self.n_dims + 1], dtype='int32')], 0))
+        down_loc = tf.cast(down_loc, 'float32') / l2i_et.expand_dims(down_zoom_factor, axis=[1] * self.n_dims)
+        inshape_tens = tf.tile(tf.expand_dims(tf.convert_to_tensor(self.inshape[:-1]), 0), tile_shape)
+        inshape_tens = l2i_et.expand_dims(inshape_tens, axis=[1] * self.n_dims)
+        down_loc = K.clip(down_loc, 0., tf.cast(inshape_tens, 'float32'))
+        vol = tf.map_fn(self._single_down_interpn, [vol, down_loc], tf.float32)
+
+        # add noise with predefined probability
+        if self.noise_std > 0:
+            sample_shape = tf.concat([batchsize, tf.ones([self.n_dims], dtype='int32'),
+                                      self.n_channels * tf.ones([1], dtype='int32')], 0)
+            noise = tf.random.normal(tf.shape(vol), stddev=tf.random.uniform(sample_shape, maxval=self.noise_std))
+            if self.prob_noise == 1:
+                vol += noise
+            else:
+                vol = K.switch(tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_noise)), vol + noise, vol)
+
+        # upsample
+        up_loc = tf.tile(self.up_grid, tf.concat([batchsize, tf.ones([self.n_dims + 1], dtype='int32')], axis=0))
+        up_loc = tf.cast(up_loc, 'float32') / l2i_et.expand_dims(up_zoom_factor, axis=[1] * self.n_dims)
+        vol = tf.map_fn(self._single_up_interpn, [vol, up_loc], tf.float32)
+
+        # return upsampled volume
+        if not self.build_dist_map:
+            return vol
+
+        # return upsampled volumes with distance maps
+        else:
+
+            # get grid points
+            floor = tf.math.floor(up_loc)
+            ceil = tf.math.ceil(up_loc)
+
+            # get distances of every voxel to higher and lower grid points for every dimension
+            f_dist = up_loc - floor
+            c_dist = ceil - up_loc
+
+            # keep minimum 1d distances, and compute 3d distance to nearest grid point
+            dist = tf.math.minimum(f_dist, c_dist) * l2i_et.expand_dims(subsample_res, axis=[1] * self.n_dims)
+            dist = tf.math.sqrt(tf.math.reduce_sum(tf.math.square(dist), axis=-1, keepdims=True))
+
+            return [vol, dist]
+
+    @staticmethod
+    def _single_down_interpn(inputs):
+        return nrn_utils.interpn(inputs[0], inputs[1], interp_method='nearest')
+
+    @staticmethod
+    def _single_up_interpn(inputs):
+        return nrn_utils.interpn(inputs[0], inputs[1], interp_method='linear')
+
+    def compute_output_shape(self, input_shape):
+        output_shape = tuple([None] + self.resample_shape + [input_shape[0][-1]])
+        return [output_shape] * 2 if self.build_dist_map else output_shape
+
+
+class BiasFieldCorruption(Layer):
+    """This layer applies a smooth random bias field to the input by applying the following steps:
+    1) we first sample a value for the standard deviation of a centred normal distribution
+    2) a small-size SVF is sampled from this normal distribution
+    3) the small SVF is then resized with trilinear interpolation to image size
+    4) it is rescaled to positive values by taking the voxel-wise exponential
+    5) it is multiplied to the input tensor.
+    The input tensor is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
+
+    :param bias_field_std: maximum value of the standard deviation sampled in 1 (it will be sampled from the range
+    [0, bias_field_std])
+    :param bias_scale: ratio between the shape of the input tensor and the shape of the sampled SVF.
+    :param same_bias_for_all_channels: whether to apply the same bias field to all the channels of the input tensor.
+    :param prob: probability to apply this bias field corruption.
+    """
+
+    def __init__(self, bias_field_std=.5, bias_scale=.025, same_bias_for_all_channels=False, prob=0.95, **kwargs):
+
+        # input shape
+        self.several_inputs = False
+        self.inshape = None
+        self.n_dims = None
+        self.n_channels = None
+
+        # sampling shape
+        self.std_shape = None
+        self.small_bias_shape = None
+
+        # bias field parameters
+        self.bias_field_std = bias_field_std
+        self.bias_scale = bias_scale
+        self.same_bias_for_all_channels = same_bias_for_all_channels
+        self.prob = prob
+
+        super(BiasFieldCorruption, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["bias_field_std"] = self.bias_field_std
+        config["bias_scale"] = self.bias_scale
+        config["same_bias_for_all_channels"] = self.same_bias_for_all_channels
+        config["prob"] = self.prob
+        return config
+
+    def build(self, input_shape):
+
+        # input shape
+        if isinstance(input_shape, list):
+            self.several_inputs = True
+            self.inshape = input_shape
+        else:
+            self.inshape = [input_shape]
+        self.n_dims = len(self.inshape[0]) - 2
+        self.n_channels = self.inshape[0][-1]
+
+        # sampling shapes
+        self.std_shape = [1] * (self.n_dims + 1)
+        self.small_bias_shape = utils.get_resample_shape(self.inshape[0][1:self.n_dims + 1], self.bias_scale, 1)
+        if not self.same_bias_for_all_channels:
+            self.std_shape[-1] = self.n_channels
+            self.small_bias_shape[-1] = self.n_channels
+
+        self.built = True
+        super(BiasFieldCorruption, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        if not self.several_inputs:
+            inputs = [inputs]
+
+        if self.bias_field_std > 0:
+
+            # sampling shapes
+            batchsize = tf.split(tf.shape(inputs[0]), [1, -1])[0]
+            std_shape = tf.concat([batchsize, tf.convert_to_tensor(self.std_shape, dtype='int32')], 0)
+            bias_shape = tf.concat([batchsize, tf.convert_to_tensor(self.small_bias_shape, dtype='int32')], axis=0)
+
+            # sample small bias field
+            bias_field = tf.random.normal(bias_shape, stddev=tf.random.uniform(std_shape, maxval=self.bias_field_std))
+
+            # resize bias field and take exponential
+            bias_field = nrn_layers.Resize(size=self.inshape[0][1:self.n_dims + 1], interp_method='linear')(bias_field)
+            bias_field = tf.math.exp(bias_field)
+
+            # apply bias field with predefined probability
+            if self.prob == 1:
+                return [tf.math.multiply(bias_field, v) for v in inputs]
+            else:
+                rand_trans = tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob))
+                if self.several_inputs:
+                    return [K.switch(rand_trans, tf.math.multiply(bias_field, v), v) for v in inputs]
+                else:
+                    return K.switch(rand_trans, tf.math.multiply(bias_field, inputs[0]), inputs[0])
+
+        else:
+            return inputs
+
+
+class IntensityAugmentation(Layer):
+    """This layer enables to augment the intensities of the input tensor, as well as to apply min_max normalisation.
+    The following steps are applied (all are optional):
+    1) white noise corruption, with a randomly sampled std dev.
+    2) clip the input between two values
+    3) min-max normalisation
+    4) gamma augmentation (i.e. voxel-wise exponentiation by a randomly sampled power)
+    The input tensor is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
+
+    :param noise_std: maximum value of the standard deviation of the Gaussian white noise used in 1 (it will be sampled
+    from the range [0, noise_std]). Set to 0 to skip this step.
+    :param clip: clip the input tensor between the given values. Can either be: a number (in which case we clip between
+    0 and the given value), or a list or a numpy array with two elements. Default is 0, where no clipping occurs.
+    :param normalise: whether to apply min-max normalisation, to normalise between 0 and 1. Default is True.
+    :param norm_perc: percentiles (between 0 and 1) of the sorted intensity values for robust normalisation. Can be:
+    a number (in which case the robust minimum is the provided percentile of sorted values, and the maximum is the
+    1 - norm_perc percentile), or a list/numpy array of 2 elements (percentiles for the minimum and maximum values).
+    The minimum and maximum values are computed separately for each channel if separate_channels is True.
+    Default is 0, where we simply take the minimum and maximum values.
+    :param gamma_std: standard deviation of the normal distribution from which we sample gamma (in log domain).
+    Default is 0, where no gamma augmentation occurs.
+    :param contrast_inversion: whether to perform contrast inversion (i.e. 1 - x). If True, this is performed randomly
+    for each element of the batch, as well as for each channel.
+    :param separate_channels: whether to augment all channels separately. Default is True.
+    :param prob_noise: probability to apply noise injection
+    :param prob_gamma: probability to apply gamma augmentation
+    """
+
+    def __init__(self, noise_std=0, clip=0, normalise=True, norm_perc=0, gamma_std=0, contrast_inversion=False,
+                 separate_channels=True, prob_noise=0.95, prob_gamma=1, **kwargs):
+
+        # shape attributes
+        self.n_dims = None
+        self.n_channels = None
+        self.flatten_shape = None
+        self.expand_minmax_dim = None
+        self.one = None
+
+        # inputs
+        self.noise_std = noise_std
+        self.clip = clip
+        self.clip_values = None
+        self.normalise = normalise
+        self.norm_perc = norm_perc
+        self.perc = None
+        self.gamma_std = gamma_std
+        self.separate_channels = separate_channels
+        self.contrast_inversion = contrast_inversion
+        self.prob_noise = prob_noise
+        self.prob_gamma = prob_gamma
+
+        super(IntensityAugmentation, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["noise_std"] = self.noise_std
+        config["clip"] = self.clip
+        config["normalise"] = self.normalise
+        config["norm_perc"] = self.norm_perc
+        config["gamma_std"] = self.gamma_std
+        config["separate_channels"] = self.separate_channels
+        config["prob_noise"] = self.prob_noise
+        config["prob_gamma"] = self.prob_gamma
+        return config
+
+    def build(self, input_shape):
+        self.n_dims = len(input_shape) - 2
+        self.n_channels = input_shape[-1]
+        self.flatten_shape = np.prod(np.array(input_shape[1:-1]))
+        self.flatten_shape = self.flatten_shape * self.n_channels if not self.separate_channels else self.flatten_shape
+        self.expand_minmax_dim = self.n_dims if self.separate_channels else self.n_dims + 1
+        self.one = tf.ones([1], dtype='int32')
+        if self.clip:
+            self.clip_values = utils.reformat_to_list(self.clip)
+            self.clip_values = self.clip_values if len(self.clip_values) == 2 else [0, self.clip_values[0]]
+        else:
+            self.clip_values = None
+        if self.norm_perc:
+            self.perc = utils.reformat_to_list(self.norm_perc)
+            self.perc = self.perc if len(self.perc) == 2 else [self.perc[0], 1 - self.perc[0]]
+        else:
+            self.perc = None
+
+        self.built = True
+        super(IntensityAugmentation, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # prepare shape for sampling the noise and gamma std dev (depending on whether we augment channels separately)
+        batchsize = tf.split(tf.shape(inputs), [1, -1])[0]
+        if (self.noise_std > 0) | (self.gamma_std > 0) | self.contrast_inversion:
+            sample_shape = tf.concat([batchsize, tf.ones([self.n_dims], dtype='int32')], 0)
+            if self.separate_channels:
+                sample_shape = tf.concat([sample_shape, self.n_channels * self.one], 0)
+            else:
+                sample_shape = tf.concat([sample_shape, self.one], 0)
+        else:
+            sample_shape = None
+
+        # add noise with predefined probability
+        if self.noise_std > 0:
+            noise_stddev = tf.random.uniform(sample_shape, maxval=self.noise_std)
+            if self.separate_channels:
+                noise = tf.random.normal(tf.shape(inputs), stddev=noise_stddev)
+            else:
+                noise = tf.random.normal(tf.shape(tf.split(inputs, [1, -1], -1)[0]), stddev=noise_stddev)
+                noise = tf.tile(noise, tf.convert_to_tensor([1] * (self.n_dims + 1) + [self.n_channels]))
+            if self.prob_noise == 1:
+                inputs = inputs + noise
+            else:
+                inputs = K.switch(tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_noise)),
+                                  inputs + noise, inputs)
+
+        # clip images to given values
+        if self.clip_values is not None:
+            inputs = K.clip(inputs, self.clip_values[0], self.clip_values[1])
+
+        # normalise
+        if self.normalise:
+            # define robust min and max by sorting values and taking percentile
+            if self.perc is not None:
+                if self.separate_channels:
+                    shape = tf.concat([batchsize, self.flatten_shape * self.one, self.n_channels * self.one], 0)
+                else:
+                    shape = tf.concat([batchsize, self.flatten_shape * self.one], 0)
+                intensities = tf.sort(tf.reshape(inputs, shape), axis=1)
+                m = intensities[:, max(int(self.perc[0] * self.flatten_shape), 0), ...]
+                M = intensities[:, min(int(self.perc[1] * self.flatten_shape), self.flatten_shape - 1), ...]
+            # simple min and max
+            else:
+                m = K.min(inputs, axis=list(range(1, self.expand_minmax_dim + 1)))
+                M = K.max(inputs, axis=list(range(1, self.expand_minmax_dim + 1)))
+            # normalise
+            m = l2i_et.expand_dims(m, axis=[1] * self.expand_minmax_dim)
+            M = l2i_et.expand_dims(M, axis=[1] * self.expand_minmax_dim)
+            inputs = tf.clip_by_value(inputs, m, M)
+            inputs = (inputs - m) / (M - m + K.epsilon())
+
+        # apply voxel-wise exponentiation with predefined probability
+        if self.gamma_std > 0:
+            gamma = tf.random.normal(sample_shape, stddev=self.gamma_std)
+            if self.prob_gamma == 1:
+                inputs = tf.math.pow(inputs, tf.math.exp(gamma))
+            else:
+                inputs = K.switch(tf.squeeze(K.less(tf.random.uniform([1], 0, 1), self.prob_gamma)),
+                                  tf.math.pow(inputs, tf.math.exp(gamma)), inputs)
+
+        # apply random contrast inversion
+        if self.contrast_inversion:
+            rand_invert = tf.less(tf.random.uniform(sample_shape, maxval=1), 0.5)
+            split_channels = tf.split(inputs, [1] * self.n_channels, axis=-1)
+            split_rand_invert = tf.split(rand_invert, [1] * self.n_channels, axis=-1)
+            inverted_channel = list()
+            for (channel, invert) in zip(split_channels, split_rand_invert):
+                inverted_channel.append(tf.map_fn(self._single_invert, [channel, invert], dtype=channel.dtype))
+            inputs = tf.concat(inverted_channel, -1)
+
+        return inputs
+
+    @staticmethod
+    def _single_invert(inputs):
+        return K.switch(tf.squeeze(inputs[1]), 1 - inputs[0], inputs[0])
+
+
+class DiceLoss(Layer):
+    """This layer computes the soft Dice loss between two tensors.
+    These tensors are expected to have the same shape (one-hot encoding) [batch, size_dim1, ..., size_dimN, n_labels].
+    The first input tensor is the GT and the second is the prediction: dice_loss = DiceLoss()([gt, pred])
+
+    :param class_weights: (optional) if given, the loss is obtained by a weighted average of the Dice across labels.
+    Must be a sequence or 1d numpy array of length n_labels. Can also be -1, where the weights are dynamically set to
+    the inverse of the volume of each label in the ground truth.
+    :param boundary_weights: (optional) bonus weight that we apply to the voxels close to boundaries between structures
+    when computing the loss. Default is 0 where no boundary weighting is applied.
+    :param boundary_dist: (optional) if boundary_weight is not 0, the extra boundary weighting is applied to all voxels
+    within this distance to a region boundary. Default is 3.
+    :param skip_background: (optional) whether to skip boundary weighting for the background class, as this may be
+    redundant when we have several labels. This is only used if boundary_weight is not 0.
+    :param enable_checks: (optional) whether to make sure that the 2 input tensors are probabilistic (i.e. the label
+    probabilities sum to 1 at each voxel location). Default is True.
+    """
+
+    def __init__(self,
+                 class_weights=None,
+                 boundary_weights=0,
+                 boundary_dist=3,
+                 skip_background=True,
+                 enable_checks=True,
+                 **kwargs):
+
+        self.class_weights = class_weights
+        self.dynamic_weighting = False
+        self.class_weights_tens = None
+        self.boundary_weights = boundary_weights
+        self.boundary_dist = boundary_dist
+        self.skip_background = skip_background
+        self.enable_checks = enable_checks
+        self.spatial_axes = None
+        self.avg_pooling_layer = None
+        super(DiceLoss, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["class_weights"] = self.class_weights
+        config["boundary_weights"] = self.boundary_weights
+        config["boundary_dist"] = self.boundary_dist
+        config["skip_background"] = self.skip_background
+        config["enable_checks"] = self.enable_checks
+        return config
+
+    def build(self, input_shape):
+
+        # get shape
+        assert len(input_shape) == 2, 'DiceLoss expects 2 inputs to compute the Dice loss.'
+        assert input_shape[0] == input_shape[1], 'the two inputs must have the same shape.'
+        inshape = input_shape[0][1:]
+        n_dims = len(inshape[:-1])
+        n_labels = inshape[-1]
+        self.spatial_axes = list(range(1, n_dims + 1))
+        self.avg_pooling_layer = getattr(keras.layers, 'AvgPool%dD' % n_dims)
+        self.skip_background = False if n_labels == 1 else self.skip_background
+
+        # build tensor with class weights
+        if self.class_weights is not None:
+            if self.class_weights == -1:
+                self.dynamic_weighting = True
+            else:
+                class_weights_tens = utils.reformat_to_list(self.class_weights, n_labels)
+                class_weights_tens = tf.convert_to_tensor(class_weights_tens, 'float32')
+                self.class_weights_tens = l2i_et.expand_dims(class_weights_tens, 0)
+
+        self.built = True
+        super(DiceLoss, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # make sure tensors are probabilistic
+        gt = inputs[0]
+        pred = inputs[1]
+        if self.enable_checks:  # disabling is useful to, e.g., use incomplete label maps
+            gt = K.clip(gt / (tf.math.reduce_sum(gt, axis=-1, keepdims=True) + tf.keras.backend.epsilon()), 0, 1)
+            pred = K.clip(pred / (tf.math.reduce_sum(pred, axis=-1, keepdims=True) + tf.keras.backend.epsilon()), 0, 1)
+
+        # compute dice loss for each label
+        top = 2 * gt * pred
+        bottom = tf.math.square(gt) + tf.math.square(pred)
+
+        # apply boundary weighting (ie voxels close to region boundaries will be counted several times to compute Dice)
+        if self.boundary_weights:
+            avg = self.avg_pooling_layer(pool_size=2 * self.boundary_dist + 1, strides=1, padding='same')(gt)
+            boundaries = tf.cast(avg > 0., 'float32') * tf.cast(avg < (1 / len(self.spatial_axes) - 1e-4), 'float32')
+            if self.skip_background:
+                boundaries_channels = tf.unstack(boundaries, axis=-1)
+                boundaries = tf.stack([tf.zeros_like(boundaries_channels[0])] + boundaries_channels[1:], axis=-1)
+            boundary_weights_tensor = 1 + self.boundary_weights * boundaries
+            top *= boundary_weights_tensor
+            bottom *= boundary_weights_tensor
+        else:
+            boundary_weights_tensor = None
+
+        # compute loss
+        top = tf.math.reduce_sum(top, self.spatial_axes)
+        bottom = tf.math.reduce_sum(bottom, self.spatial_axes)
+        dice = (top + tf.keras.backend.epsilon()) / (bottom + tf.keras.backend.epsilon())
+        loss = 1 - dice
+
+        # apply class weighting across labels. In this case loss will have shape (batch), otherwise (batch, n_labels).
+        if self.dynamic_weighting:  # the weight of a class is the inverse of its volume in the gt
+            if boundary_weights_tensor is not None:  # we account for the boundary weighting to compute volume
+                self.class_weights_tens = 1 / tf.reduce_sum(gt * boundary_weights_tensor, self.spatial_axes)
+            else:
+                self.class_weights_tens = 1 / tf.reduce_sum(gt, self.spatial_axes)
+        if self.class_weights_tens is not None:
+            self. class_weights_tens /= tf.reduce_sum(self.class_weights_tens, -1)
+            loss = tf.reduce_sum(loss * self.class_weights_tens, -1)
+
+        return tf.math.reduce_mean(loss)
+
+    def compute_output_shape(self, input_shape):
+        return [[]]
+
+
+class WeightedL2Loss(Layer):
+    """This layer computes a L2 loss weighted by a specified factor (target_value) between two tensors.
+    This is designed to be used on the layer before the softmax.
+    The tensors are expected to have the same shape [batchsize, size_dim1, ..., size_dimN, n_labels].
+    The first input tensor is the GT and the second is the prediction: wl2_loss = WeightedL2Loss()([gt, pred])
+
+    :param target_value: target value for the layer before softmax: target_value when gt = 1, -target_value when gt = 0.
+    """
+
+    def __init__(self, target_value=5, **kwargs):
+        self.target_value = target_value
+        self.n_labels = None
+        super(WeightedL2Loss, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["target_value"] = self.target_value
+        return config
+
+    def build(self, input_shape):
+        assert len(input_shape) == 2, 'DiceLoss expects 2 inputs to compute the Dice loss.'
+        assert input_shape[0] == input_shape[1], 'the two inputs must have the same shape.'
+        self.n_labels = input_shape[0][-1]
+        self.built = True
+        super(WeightedL2Loss, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+        gt = inputs[0]
+        pred = inputs[1]
+        weights = tf.expand_dims(1 - gt[..., 0] + 1e-8, -1)
+        return K.sum(weights * K.square(pred - self.target_value * (2 * gt - 1))) / (K.sum(weights) * self.n_labels)
+
+    def compute_output_shape(self, input_shape):
+        return [[]]
+
+
+class CrossEntropyLoss(Layer):
+    """This layer computes the cross-entropy loss between two tensors.
+    These tensors are expected to have the same shape (one-hot encoding) [batch, size_dim1, ..., size_dimN, n_labels].
+    The first input tensor is the GT and the second is the prediction: ce_loss = CrossEntropyLoss()([gt, pred])
+
+    :param class_weights: (optional) if given, the loss is obtained by a weighted average of the Dice across labels.
+    Must be a sequence or 1d numpy array of length n_labels. Can also be -1, where the weights are dynamically set to
+    the inverse of the volume of each label in the ground truth.
+    :param boundary_weights: (optional) bonus weight that we apply to the voxels close to boundaries between structures
+    when computing the loss. Default is 0 where no boundary weighting is applied.
+    :param boundary_dist: (optional) if boundary_weight is not 0, the extra boundary weighting is applied to all voxels
+    within this distance to a region boundary. Default is 3.
+    :param skip_background: (optional) whether to skip boundary weighting for the background class, as this may be
+    redundant when we have several labels. This is only used if boundary_weight is not 0.
+    :param enable_checks: (optional) whether to make sure that the 2 input tensors are probabilistic (i.e. the label
+    probabilities sum to 1 at each voxel location). Default is True.
+    """
+
+    def __init__(self,
+                 class_weights=None,
+                 boundary_weights=0,
+                 boundary_dist=3,
+                 skip_background=True,
+                 enable_checks=True,
+                 **kwargs):
+
+        self.class_weights = class_weights
+        self.dynamic_weighting = False
+        self.class_weights_tens = None
+        self.boundary_weights = boundary_weights
+        self.boundary_dist = boundary_dist
+        self.skip_background = skip_background
+        self.enable_checks = enable_checks
+        self.spatial_axes = None
+        self.avg_pooling_layer = None
+        super(CrossEntropyLoss, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["class_weights"] = self.class_weights
+        config["boundary_weights"] = self.boundary_weights
+        config["boundary_dist"] = self.boundary_dist
+        config["skip_background"] = self.skip_background
+        config["enable_checks"] = self.enable_checks
+        return config
+
+    def build(self, input_shape):
+
+        # get shape
+        assert len(input_shape) == 2, 'CrossEntropy expects 2 inputs to compute the Dice loss.'
+        assert input_shape[0] == input_shape[1], 'the two inputs must have the same shape.'
+        inshape = input_shape[0][1:]
+        n_dims = len(inshape[:-1])
+        n_labels = inshape[-1]
+        self.spatial_axes = list(range(1, n_dims + 1))
+        self.avg_pooling_layer = getattr(keras.layers, 'AvgPool%dD' % n_dims)
+        self.skip_background = False if n_labels == 1 else self.skip_background
+
+        # build tensor with class weights
+        if self.class_weights is not None:
+            if self.class_weights == -1:
+                self.dynamic_weighting = True
+            else:
+                class_weights_tens = utils.reformat_to_list(self.class_weights, n_labels)
+                class_weights_tens = tf.convert_to_tensor(class_weights_tens, 'float32')
+                self.class_weights_tens = l2i_et.expand_dims(class_weights_tens, [0] * (1 + n_dims))
+
+        self.built = True
+        super(CrossEntropyLoss, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # make sure tensors are probabilistic
+        gt = inputs[0]
+        pred = inputs[1]
+        if self.enable_checks:  # disabling is useful to, e.g., use incomplete label maps
+            gt = K.clip(gt / (tf.math.reduce_sum(gt, axis=-1, keepdims=True) + tf.keras.backend.epsilon()), 0, 1)
+            pred = pred / (tf.math.reduce_sum(pred, axis=-1, keepdims=True) + tf.keras.backend.epsilon())
+            pred = K.clip(pred, tf.keras.backend.epsilon(), 1 - tf.keras.backend.epsilon())  # to avoid log(0)
+
+        # compare prediction/target, ce has the same shape has the input tensors
+        ce = -gt * tf.math.log(pred)
+
+        # apply boundary weighting (ie voxels close to region boundaries will be counted several times to compute Dice)
+        if self.boundary_weights:
+            avg = self.avg_pooling_layer(pool_size=2 * self.boundary_dist + 1, strides=1, padding='same')(gt)
+            boundaries = tf.cast(avg > 0., 'float32') * tf.cast(avg < (1 / len(self.spatial_axes) - 1e-4), 'float32')
+            if self.skip_background:
+                boundaries_channels = tf.unstack(boundaries, axis=-1)
+                boundaries = tf.stack([tf.zeros_like(boundaries_channels[0])] + boundaries_channels[1:], axis=-1)
+            boundary_weights_tensor = 1 + self.boundary_weights * boundaries
+            ce *= boundary_weights_tensor
+        else:
+            boundary_weights_tensor = None
+
+        # apply class weighting across labels. By the end of this, ce still has the same shape has the input tensors.
+        if self.dynamic_weighting:  # the weight of a class is the inverse of its volume in the gt
+            if boundary_weights_tensor is not None:  # we account for the boundary weighting to compute volume
+                self.class_weights_tens = 1 / tf.reduce_sum(gt * boundary_weights_tensor, self.spatial_axes, True)
+            else:
+                self.class_weights_tens = 1 / tf.reduce_sum(gt, self.spatial_axes)
+        if self.class_weights_tens is not None:
+            self.class_weights_tens /= tf.reduce_sum(self.class_weights_tens, -1)
+            ce = tf.reduce_sum(ce * self.class_weights_tens, -1)
+
+        # sum along label axis, and take the mean along spatial dimensions
+        ce = tf.math.reduce_mean(tf.math.reduce_sum(ce, axis=-1))
+
+        return ce
+
+    def compute_output_shape(self, input_shape):
+        return [[]]
+
+
+class MomentLoss(Layer):
+    """This layer computes a moment loss between two tensors. Specifically, it computes the distance between the centres
+    of gravity for all the channels of the two tensors, and then returns a value averaged across all channels.
+    These tensors are expected to have the same shape [batch, size_dim1, ..., size_dimN, n_channels].
+    The first input tensor is the GT and the second is the prediction: moment_loss = MomentLoss()([gt, pred])
+
+    :param class_weights: (optional) if given, the loss is obtained by a weighted average of the Dice across labels.
+    Must be a sequence or 1d numpy array of length n_labels. Can also be -1, where the weights are dynamically set to
+    the inverse of the volume of each label in the ground truth.
+    :param enable_checks: (optional) whether to make sure that the 2 input tensors are probabilistic (i.e. the label
+    probabilities sum to 1 at each voxel location). Default is True.
+    """
+
+    def __init__(self, class_weights=None, enable_checks=False, **kwargs):
+        self.class_weights = class_weights
+        self.dynamic_weighting = False
+        self.class_weights_tens = None
+        self.enable_checks = enable_checks
+        self.spatial_axes = None
+        self.coordinates = None
+        super(MomentLoss, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["class_weights"] = self.class_weights
+        config["enable_checks"] = self.enable_checks
+        return config
+
+    def build(self, input_shape):
+
+        # get shape
+        assert len(input_shape) == 2, 'MomentLoss expects 2 inputs to compute the Dice loss.'
+        assert input_shape[0] == input_shape[1], 'the two inputs must have the same shape.'
+        inshape = input_shape[0][1:]
+        n_dims = len(inshape[:-1])
+        n_labels = inshape[-1]
+        self.spatial_axes = list(range(1, n_dims + 1))
+
+        # build coordinate meshgrid of size (1, dim1, dim2, ..., dimN, ndim, nchan)
+        self.coordinates = tf.stack(nrn_utils.volshape_to_ndgrid(inshape[:-1]), -1)
+        self.coordinates = tf.cast(l2i_et.expand_dims(tf.stack([self.coordinates] * n_labels, -1), 0), 'float32')
+
+        # build tensor with class weights
+        if self.class_weights is not None:
+            if self.class_weights == -1:
+                self.dynamic_weighting = True
+            else:
+                class_weights_tens = utils.reformat_to_list(self.class_weights, n_labels)
+                class_weights_tens = tf.convert_to_tensor(class_weights_tens, 'float32')
+                self.class_weights_tens = l2i_et.expand_dims(class_weights_tens, 0)
+
+        self.built = True
+        super(MomentLoss, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # make sure tensors are probabilistic
+        gt = inputs[0]  # (B, dim1, dim2, ..., dimN, nchan)
+        pred = inputs[1]
+        if self.enable_checks:  # disabling is useful to, e.g., use incomplete label maps
+            gt = gt / (tf.math.reduce_sum(gt, axis=-1, keepdims=True) + tf.keras.backend.epsilon())
+            pred = pred / (tf.math.reduce_sum(pred, axis=-1, keepdims=True) + tf.keras.backend.epsilon())
+
+        # compute loss
+        gt_mean_coordinates = self._mean_coordinates(gt)  # (B, ndim, nchan)
+        pred_mean_coordinates = self._mean_coordinates(pred)
+        loss = tf.math.sqrt(tf.reduce_sum(tf.square(pred_mean_coordinates - gt_mean_coordinates), axis=1))  # (B, nchan)
+
+        # apply class weighting across labels. In this case loss will have shape (batch), otherwise (batch, n_labels).
+        if self.dynamic_weighting:  # the weight of a class is the inverse of its volume in the gt
+            self.class_weights_tens = 1 / tf.reduce_sum(gt, self.spatial_axes)
+        if self.class_weights_tens is not None:
+            self.class_weights_tens /= tf.reduce_sum(self.class_weights_tens, -1)
+            loss = tf.reduce_sum(loss * self.class_weights_tens, -1)
+
+        return tf.math.reduce_mean(loss)
+
+    def _mean_coordinates(self, tensor):
+        tensor = l2i_et.expand_dims(tensor, axis=-2)  # (B, dim1, dim2, ..., dimN, 1, nchan)
+        numerator = tf.reduce_sum(tensor * self.coordinates, axis=self.spatial_axes)  # (B, ndim, nchan)
+        denominator = tf.reduce_sum(tensor, axis=self.spatial_axes) + tf.keras.backend.epsilon()
+        return numerator / denominator
+
+    def compute_output_shape(self, input_shape):
+        return [[]]
+
+
+class ResetValuesToZero(Layer):
+    """This layer enables to reset given values to 0 within the input tensors.
+
+    :param values: list of values to be reset to 0.
+
+    example:
+    input = tf.convert_to_tensor(np.array([[1, 0, 2, 2, 2, 2, 0],
+                                           [1, 3, 3, 3, 3, 3, 3],
+                                           [1, 0, 0, 0, 4, 4, 4]]))
+    values = [1, 3]
+    ResetValuesToZero(values)(input)
+    >> [[0, 0, 2, 2, 2, 2, 0],
+        [0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 4, 4, 4]]
+    """
+
+    def __init__(self, values, **kwargs):
+        assert values is not None, 'please provide correct list of values, received None'
+        self.values = utils.reformat_to_list(values)
+        self.values_tens = None
+        self.n_values = len(values)
+        super(ResetValuesToZero, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["values"] = self.values
+        return config
+
+    def build(self, input_shape):
+        self.values_tens = tf.convert_to_tensor(self.values)
+        self.built = True
+        super(ResetValuesToZero, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+        values = tf.cast(self.values_tens, dtype=inputs.dtype)
+        for i in range(self.n_values):
+            inputs = tf.where(tf.equal(inputs, values[i]), tf.zeros_like(inputs), inputs)
+        return inputs
+
+
+class ConvertLabels(Layer):
+    """Convert all labels in a tensor by the corresponding given set of values.
+    labels_converted = ConvertLabels(source_values, dest_values)(labels).
+    labels must be an int32 tensor, and labels_converted will also be int32.
+
+    :param source_values: list of all the possible values in labels. Must be a list or a 1D numpy array.
+    :param dest_values: list of all the target label values. Must be ordered the same as source values:
+    labels[labels == source_values[i]] = dest_values[i].
+    If None (default), dest_values is equal to [0, ..., N-1], where N is the total number of values in source_values,
+    which enables to remap label maps to [0, ..., N-1].
+    """
+
+    def __init__(self, source_values, dest_values=None, **kwargs):
+        self.source_values = source_values
+        self.dest_values = dest_values
+        self.lut = None
+        super(ConvertLabels, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["source_values"] = self.source_values
+        config["dest_values"] = self.dest_values
+        return config
+
+    def build(self, input_shape):
+        self.lut = tf.convert_to_tensor(utils.get_mapping_lut(self.source_values, dest=self.dest_values), dtype='int32')
+        self.built = True
+        super(ConvertLabels, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+        return tf.gather(self.lut, tf.cast(inputs, dtype='int32'))
+
+
+class PadAroundCentre(Layer):
+    """Pad the input tensor to the specified shape with the given value.
+    The input tensor is expected to have shape [batchsize, shape_dim1, ..., shape_dimn, channel].
+    :param pad_margin: margin to use for padding. The tensor will be padded by the provided margin on each side.
+    Can either be a number (all axes padded with the same margin), or a  list/numpy array of length n_dims.
+    example: if tensor is of shape [batch, x, y, z, n_channels] and margin=10, then the padded tensor will be of
+    shape [batch, x+2*10, y+2*10, z+2*10, n_channels].
+    :param pad_shape: shape to pad the tensor to. Can either be a number (all axes padded to the same shape), or a
+    list/numpy array of length n_dims.
+    :param value: value to pad the tensors with. Default is 0.
+    """
+
+    def __init__(self, pad_margin=None, pad_shape=None, value=0, **kwargs):
+        self.pad_margin = pad_margin
+        self.pad_shape = pad_shape
+        self.value = value
+        self.pad_margin_tens = None
+        self.pad_shape_tens = None
+        self.n_dims = None
+        super(PadAroundCentre, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["pad_margin"] = self.pad_margin
+        config["pad_shape"] = self.pad_shape
+        config["value"] = self.value
+        return config
+
+    def build(self, input_shape):
+        # input shape
+        self.n_dims = len(input_shape) - 2
+        shape = list(input_shape)
+        shape[0] = 0
+        shape[-1] = 0
+
+        if self.pad_margin is not None:
+            assert self.pad_shape is None, 'please do not provide a padding shape and margin at the same time.'
+
+            # reformat padding margins
+            pad = np.transpose(np.array([[0] + utils.reformat_to_list(self.pad_margin, self.n_dims) + [0]] * 2))
+            self.pad_margin_tens = tf.convert_to_tensor(pad, dtype='int32')
+
+        elif self.pad_shape is not None:
+            assert self.pad_margin is None, 'please do not provide a padding shape and margin at the same time.'
+
+            # pad shape
+            tensor_shape = tf.cast(tf.convert_to_tensor(shape), 'int32')
+            self.pad_shape_tens = np.array([0] + utils.reformat_to_list(self.pad_shape, length=self.n_dims) + [0])
+            self.pad_shape_tens = tf.convert_to_tensor(self.pad_shape_tens, dtype='int32')
+            self.pad_shape_tens = tf.math.maximum(tensor_shape, self.pad_shape_tens)
+
+            # padding margin
+            min_margins = (self.pad_shape_tens - tensor_shape) / 2
+            max_margins = self.pad_shape_tens - tensor_shape - min_margins
+            self.pad_margin_tens = tf.stack([min_margins, max_margins], axis=-1)
+
+        else:
+            raise Exception('please either provide a padding shape or a padding margin.')
+
+        self.built = True
+        super(PadAroundCentre, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+        return tf.pad(inputs, self.pad_margin_tens, mode='CONSTANT', constant_values=self.value)
+
+
+class MaskEdges(Layer):
+    """Reset the edges of a tensor to zero (i.e. with bands of zeros along the specified axes).
+    The width of the zero-band is randomly drawn from a uniform distribution, whose range is given in boundaries.
+
+    :param axes: axes along which to reset edges to zero. Can be an int (single axis), or a sequence.
+    :param boundaries: numpy array of shape (len(axes), 4). Each row contains the two bounds of the uniform
+    distributions from which we draw the width of the zero-bands on each side.
+    Those bounds must be expressed in relative side (i.e. between 0 and 1).
+    :return: a tensor of the same shape as the input, with bands of zeros along the specified axes.
+
+    example:
+    tensor=tf.constant([[[[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+                       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]]]])  # shape = [1,10,10,1]
+    axes=1
+    boundaries = np.array([[0.2, 0.45, 0.85, 0.9]])
+
+    In this case, we reset the edges along the 2nd dimension (i.e. the 1st dimension after the batch dimension),
+    the 1st zero-band will expand from the 1st row to a number drawn from [0.2*tensor.shape[1], 0.45*tensor.shape[1]],
+    and the 2nd zero-band will expand from a row drawn from [0.85*tensor.shape[1], 0.9*tensor.shape[1]], to the end of
+    the tensor. A possible output could be:
+    array([[[[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
+          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
+          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
+          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+          [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]]])  # shape = [1,10,10,1]
+    """
+
+    def __init__(self, axes, boundaries, prob_mask=1, **kwargs):
+        self.axes = utils.reformat_to_list(axes, dtype='int')
+        self.boundaries = utils.reformat_to_n_channels_array(boundaries, n_dims=4, n_channels=len(self.axes))
+        self.prob_mask = prob_mask
+        self.inputshape = None
+        super(MaskEdges, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["axes"] = self.axes
+        config["boundaries"] = self.boundaries
+        config["prob_mask"] = self.prob_mask
+        return config
+
+    def build(self, input_shape):
+        self.inputshape = input_shape
+        self.built = True
+        super(MaskEdges, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # build mask
+        mask = tf.ones_like(inputs)
+        for i, axis in enumerate(self.axes):
+
+            # select restricting indices
+            axis_boundaries = self.boundaries[i, :]
+            idx1 = tf.math.round(tf.random.uniform([1],
+                                                   minval=axis_boundaries[0] * self.inputshape[axis],
+                                                   maxval=axis_boundaries[1] * self.inputshape[axis]))
+            idx2 = tf.math.round(tf.random.uniform([1],
+                                                   minval=axis_boundaries[2] * self.inputshape[axis],
+                                                   maxval=axis_boundaries[3] * self.inputshape[axis] - 1) - idx1)
+            idx3 = self.inputshape[axis] - idx1 - idx2
+            split_idx = tf.cast(tf.concat([idx1, idx2, idx3], axis=0), dtype='int32')
+
+            # update mask
+            split_list = tf.split(inputs, split_idx, axis=axis)
+            tmp_mask = tf.concat([tf.zeros_like(split_list[0]),
+                                  tf.ones_like(split_list[1]),
+                                  tf.zeros_like(split_list[2])], axis=axis)
+            mask = mask * tmp_mask
+
+        # mask second_channel
+        tensor = K.switch(tf.squeeze(K.greater(tf.random.uniform([1], 0, 1), 1 - self.prob_mask)),
+                          inputs * mask,
+                          inputs)
+
+        return [tensor, mask]
+
+    def compute_output_shape(self, input_shape):
+        return [input_shape] * 2
+
+
+class ImageGradients(Layer):
+
+    def __init__(self, gradient_type='sobel', return_magnitude=False, **kwargs):
+
+        self.gradient_type = gradient_type
+        assert (self.gradient_type == 'sobel') | (self.gradient_type == '1-step_diff'), \
+            'gradient_type should be either sobel or 1-step_diff, had %s' % self.gradient_type
+
+        # shape
+        self.n_dims = 0
+        self.shape = None
+        self.n_channels = 0
+
+        # convolution params if sobel diff
+        self.stride = None
+        self.kernels = None
+        self.convnd = None
+
+        self.return_magnitude = return_magnitude
+
+        super(ImageGradients, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["gradient_type"] = self.gradient_type
+        config["return_magnitude"] = self.return_magnitude
+        return config
+
+    def build(self, input_shape):
+
+        # get shapes
+        self.n_dims = len(input_shape) - 2
+        self.shape = input_shape[1:]
+        self.n_channels = input_shape[-1]
+
+        # prepare kernel if sobel gradients
+        if self.gradient_type == 'sobel':
+            self.kernels = l2i_et.sobel_kernels(self.n_dims)
+            self.stride = [1] * (self.n_dims + 2)
+            self.convnd = getattr(tf.nn, 'conv%dd' % self.n_dims)
+        else:
+            self.kernels = self.convnd = self.stride = None
+
+        self.built = True
+        super(ImageGradients, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        image = inputs
+        batchsize = tf.split(tf.shape(inputs), [1, -1])[0]
+        gradients = list()
+
+        # sobel method
+        if self.gradient_type == 'sobel':
+            # get sobel gradients in each direction
+            for n in range(self.n_dims):
+                gradient = image
+                # apply 1D kernel in each direction (sobel kernels are separable), instead of applying a nD kernel
+                for k in self.kernels[n]:
+                    gradient = tf.concat([self.convnd(tf.expand_dims(gradient[..., n], -1), k, self.stride, 'SAME')
+                                          for n in range(self.n_channels)], -1)
+                gradients.append(gradient)
+
+        # 1-step method, only supports 2 and 3D
+        else:
+
+            # get 1-step diff
+            if self.n_dims == 2:
+                gradients.append(image[:, 1:, :, :] - image[:, :-1, :, :])  # dx
+                gradients.append(image[:, :, 1:, :] - image[:, :, :-1, :])  # dy
+
+            elif self.n_dims == 3:
+                gradients.append(image[:, 1:, :, :, :] - image[:, :-1, :, :, :])  # dx
+                gradients.append(image[:, :, 1:, :, :] - image[:, :, :-1, :, :])  # dy
+                gradients.append(image[:, :, :, 1:, :] - image[:, :, :, :-1, :])  # dz
+
+            else:
+                raise Exception('ImageGradients only support 2D or 3D tensors for 1-step diff, had: %dD' % self.n_dims)
+
+            # pad with zeros to return tensors of the same shape as input
+            for i in range(self.n_dims):
+                tmp_shape = list(self.shape)
+                tmp_shape[i] = 1
+                zeros = tf.zeros(tf.concat([batchsize, tf.convert_to_tensor(tmp_shape, dtype='int32')], 0), image.dtype)
+                gradients[i] = tf.concat([gradients[i], zeros], axis=i + 1)
+
+        # compute total gradient magnitude if necessary, or concatenate different gradients along the channel axis
+        if self.return_magnitude:
+            gradients = tf.sqrt(tf.reduce_sum(tf.square(tf.stack(gradients, axis=-1)), axis=-1))
+        else:
+            gradients = tf.concat(gradients, axis=-1)
+
+        return gradients
+
+    def compute_output_shape(self, input_shape):
+        if not self.return_magnitude:
+            input_shape = list(input_shape)
+            input_shape[-1] = self.n_dims
+        return tuple(input_shape)
+
+
+class RandomDilationErosion(Layer):
+    """
+    GPU implementation of binary dilation or erosion. The operation can be chosen to be always a dilation, or always an
+    erosion, or randomly choosing between them for each element of the batch.
+    The chosen operation is applied to the input with a given probability. Moreover, it is also possible to randomise
+    the factor of the operation for each element of the mini-batch.
+    :param min_factor: minimum possible value for the dilation/erosion factor. Must be an integer.
+    :param max_factor: minimum possible value for the dilation/erosion factor. Must be an integer.
+    Set it to the same value as min_factor to always perform dilation/erosion with the same factor.
+    :param prob: probability with which to apply the selected operation to the input.
+    :param operation: which operation to apply. Can be 'dilation' or 'erosion' or 'random'.
+    :param return_mask: if operation is erosion and the input of this layer is a label map, we have the
+    choice to either return the eroded label map or the mask (return_mask=True)
+    """
+
+    def __init__(self, min_factor, max_factor, max_factor_dilate=None, prob=1, operation='random', return_mask=False,
+                 **kwargs):
+
+        self.min_factor = min_factor
+        self.max_factor = max_factor
+        self.max_factor_dilate = max_factor_dilate if max_factor_dilate is not None else self.max_factor
+        self.prob = prob
+        self.operation = operation
+        self.return_mask = return_mask
+        self.n_dims = None
+        self.inshape = None
+        self.n_channels = None
+        self.convnd = None
+        super(RandomDilationErosion, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config["min_factor"] = self.min_factor
+        config["max_factor"] = self.max_factor
+        config["max_factor_dilate"] = self.max_factor_dilate
+        config["prob"] = self.prob
+        config["operation"] = self.operation
+        config["return_mask"] = self.return_mask
+        return config
+
+    def build(self, input_shape):
+
+        # input shape
+        self.inshape = input_shape
+        self.n_dims = len(self.inshape) - 2
+        self.n_channels = self.inshape[-1]
+
+        # prepare convolution
+        self.convnd = getattr(tf.nn, 'conv%dd' % self.n_dims)
+
+        self.built = True
+        super(RandomDilationErosion, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
+
+        # sample probability of applying operation. If random negative is erosion and positive is dilation
+        batchsize = tf.split(tf.shape(inputs), [1, -1])[0]
+        shape = tf.concat([batchsize, tf.convert_to_tensor([1], dtype='int32')], axis=0)
+        if self.operation == 'dilation':
+            prob = tf.random.uniform(shape, 0, 1)
+        elif self.operation == 'erosion':
+            prob = tf.random.uniform(shape, -1, 0)
+        elif self.operation == 'random':
+            prob = tf.random.uniform(shape, -1, 1)
+        else:
+            raise ValueError("operation should either be 'dilation' 'erosion' or 'random', had %s" % self.operation)
+
+        # build kernel
+        if self.min_factor == self.max_factor:
+            dist_threshold = self.min_factor * tf.ones(shape, dtype='int32')
+        else:
+            if (self.max_factor == self.max_factor_dilate) | (self.operation != 'random'):
+                dist_threshold = tf.random.uniform(shape, minval=self.min_factor, maxval=self.max_factor, dtype='int32')
+            else:
+                dist_threshold = tf.cast(tf.map_fn(self._sample_factor, [prob], dtype=tf.float32), dtype='int32')
+        kernel = l2i_et.unit_kernel(dist_threshold, self.n_dims, max_dist_threshold=self.max_factor)
+
+        # convolve input mask with kernel according to given probability
+        mask = tf.cast(tf.cast(inputs, dtype='bool'), dtype='float32')
+        mask = tf.map_fn(self._single_blur, [mask, kernel, prob], dtype=tf.float32)
+        mask = tf.cast(mask, 'bool')
+
+        if self.return_mask:
+            return mask
+        else:
+            return inputs * tf.cast(mask, dtype=inputs.dtype)
+
+    def _sample_factor(self, inputs):
+        return tf.cast(K.switch(K.less(tf.squeeze(inputs[0]), 0),
+                                tf.random.uniform((1,), self.min_factor, self.max_factor, dtype='int32'),
+                                tf.random.uniform((1,), self.min_factor, self.max_factor_dilate, dtype='int32')),
+                       dtype='float32')
+
+    def _single_blur(self, inputs):
+        # dilate...
+        new_mask = K.switch(K.greater(tf.squeeze(inputs[2]), 1 - self.prob + 0.001),
+                            tf.cast(tf.greater(tf.squeeze(self.convnd(tf.expand_dims(inputs[0], 0), inputs[1],
+                                    [1] * (self.n_dims + 2), padding='SAME'), axis=0), 0.01), dtype='float32'),
+                            inputs[0])
+        # ...or erode
+        new_mask = K.switch(K.less(tf.squeeze(inputs[2]), - (1 - self.prob + 0.001)),
+                            1 - tf.cast(tf.greater(tf.squeeze(self.convnd(tf.expand_dims(1 - new_mask, 0), inputs[1],
+                                        [1] * (self.n_dims + 2), padding='SAME'), axis=0), 0.01), dtype='float32'),
+                            new_mask)
+        return new_mask
+
+    def compute_output_shape(self, input_shape):
+        return input_shape