text: from keras import backend as K
  from keras.layers import Layer
  import tensorflow as tf
  """
   Code comes from:
   Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration
      Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu
      MICCAI 2018.
  """
  
  # local
  from .utils import transform, resize, integrate_vec, affine_to_shift
  
  
  class SpatialTransformer(Layer):
      """
      N-D Spatial Transformer Tensorflow / Keras Layer
  
      The Layer can handle both affine and dense transforms. 
      Both transforms are meant to give a 'shift' from the current position.
      Therefore, a dense transform gives displacements (not absolute locations) at each voxel,
      and an affine transform gives the *difference* of the affine matrix from 
      the identity matrix.
  
      If you find this function useful, please cite:
        Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration
        Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu
        MICCAI 2018.
  
      Originally, this code was based on voxelmorph code, which 
      was in turn transformed to be dense with the help of (affine) STN code 
      via https://github.com/kevinzakka/spatial-transformer-network
  
      Since then, we've re-written the code to be generalized to any 
      dimensions, and along the way wrote grid and interpolation functions
      """
  
      def __init__(self,
                   interp_method='linear',
                   indexing='ij',
                   single_transform=False,
                   **kwargs):
          """
          Parameters: 
              interp_method: 'linear' or 'nearest'
              single_transform: whether a single transform supplied for the whole batch
              indexing (default: 'ij'): 'ij' (matrix) or 'xy' (cartesian)
                  'xy' indexing will have the first two entries of the flow 
                  (along last axis) flipped compared to 'ij' indexing
          """
          self.interp_method = interp_method
          self.ndims = None
          self.inshape = None
          self.single_transform = single_transform
  
          assert indexing in ['ij', 'xy'], "indexing has to be 'ij' (matrix) or 'xy' (cartesian)"
          self.indexing = indexing
  
          super(self.__class__, self).__init__(**kwargs)
  
  
      def build(self, input_shape):
          """
          input_shape should be a list for two inputs:
          input1: image.
          input2: transform Tensor
              if affine:
                  should be a N x N+1 matrix
                  *or* a N*N+1 tensor (which will be reshape to N x (N+1) and an identity row added)
              if not affine:
                  should be a *vol_shape x N
          """
  
          if len(input_shape) > 2:
              raise Exception('Spatial Transformer must be called on a list of length 2.'
                              'First argument is the image, second is the transform.')
          
          # set up number of dimensions
          self.ndims = len(input_shape[0]) - 2
          self.inshape = input_shape
          vol_shape = input_shape[0][1:-1]
          trf_shape = input_shape[1][1:]
  
          # the transform is an affine iff:
          # it's a 1D Tensor [dense transforms need to be at least ndims + 1]
          # it's a 2D Tensor and shape == [N+1, N+1]. 
          #   [dense with N=1, which is the only one that could have a transform shape of 2, would be of size Mx1]
          self.is_affine = len(trf_shape) == 1 or \
                           (len(trf_shape) == 2 and all([f == (self.ndims+1) for f in trf_shape]))
  
          # check sizes
          if self.is_affine and len(trf_shape) == 1:
              ex = self.ndims * (self.ndims + 1)
              if trf_shape[0] != ex:
                  raise Exception('Expected flattened affine of len %d but got %d'
                                  % (ex, trf_shape[0]))
  
          if not self.is_affine:
              if trf_shape[-1] != self.ndims:
                  raise Exception('Offset flow field size expected: %d, found: %d' 
                                  % (self.ndims, trf_shape[-1]))
  
          # confirm built
          self.built = True
  
      def call(self, inputs):
          """
          Parameters
              inputs: list with two entries
          """
  
          # check shapes
          assert len(inputs) == 2, "inputs has to be len 2, found: %d" % len(inputs)
          vol = inputs[0]
          trf = inputs[1]
  
          # necessary for multi_gpu models...
          vol = K.reshape(vol, [-1, *self.inshape[0][1:]])
          trf = K.reshape(trf, [-1, *self.inshape[1][1:]])
  
          # go from affine
          if self.is_affine:
              trf = tf.map_fn(lambda x: self._single_aff_to_shift(x, vol.shape[1:-1]), trf, dtype=tf.float32)
  
          # prepare location shift
          if self.indexing == 'xy':  # shift the first two dimensions
              trf_split = tf.split(trf, trf.shape[-1], axis=-1)
              trf_lst = [trf_split[1], trf_split[0], *trf_split[2:]]
              trf = tf.concat(trf_lst, -1)
  
          # map transform across batch
          if self.single_transform:
              fn = lambda x: self._single_transform([x, trf[0,:]])
              return tf.map_fn(fn, vol, dtype=tf.float32)
          else:
              return tf.map_fn(self._single_transform, [vol, trf], dtype=tf.float32)
  
      def _single_aff_to_shift(self, trf, volshape):
          if len(trf.shape) == 1:  # go from vector to matrix
              trf = tf.reshape(trf, [self.ndims, self.ndims + 1])
  
          # note this is unnecessarily extra graph since at every batch entry we have a tf.eye graph
          trf += tf.eye(self.ndims+1)[:self.ndims,:]  # add identity, hence affine is a shift from identitiy
          return affine_to_shift(trf, volshape, shift_center=True)
  
      def _single_transform(self, inputs):
          return transform(inputs[0], inputs[1], interp_method=self.interp_method)
  
  
  class Resize(Layer):
      """
      N-D Resize Tensorflow / Keras Layer
      Note: this is not re-shaping an existing volume, but resizing, like scipy's "Zoom"
  
      If you find this function useful, please cite:
          Anatomical Priors in Convolutional Networks for Unsupervised Biomedical Segmentation,Dalca AV, Guttag J, Sabuncu MR
          CVPR 2018  
  
      Since then, we've re-written the code to be generalized to any 
      dimensions, and along the way wrote grid and interpolation functions
      """
  
      def __init__(self,
                   zoom_factor,
                   interp_method='linear',
                   **kwargs):
          """
          Parameters: 
              interp_method: 'linear' or 'nearest'
                  'xy' indexing will have the first two entries of the flow 
                  (along last axis) flipped compared to 'ij' indexing
          """
          self.zoom_factor = zoom_factor
          self.interp_method = interp_method
          self.ndims = None
          self.inshape = None
          super(Resize, self).__init__(**kwargs)
  
      def build(self, input_shape):
          """
          input_shape should be an element of list of one inputs:
          input1: volume
                  should be a *vol_shape x N
          """
  
          if isinstance(input_shape[0], (list, tuple)) and len(input_shape) > 1:
              raise Exception('Resize must be called on a list of length 1.'
                              'First argument is the image, second is the transform.')
  
          if isinstance(input_shape[0], (list, tuple)):
              input_shape = input_shape[0]
  
          # set up number of dimensions
          self.ndims = len(input_shape) - 2
          self.inshape = input_shape
  
          # confirm built
          self.built = True
  
      def call(self, inputs):
          """
          Parameters
              inputs: volume of list with one volume
          """
  
          # check shapes
          if isinstance(inputs, (list, tuple)):
              assert len(inputs) == 1, "inputs has to be len 1. found: %d" % len(inputs)
              vol = inputs[0]
          else:
              vol = inputs
  
          # necessary for multi_gpu models...
          vol = K.reshape(vol, [-1, *self.inshape[1:]])
  
          # map transform across batch
          return tf.map_fn(self._single_resize, vol, dtype=tf.float32)
  
      def compute_output_shape(self, input_shape):
          output_shape = [input_shape[0]]
          output_shape += [int(f * self.zoom_factor) for f in input_shape[1:-1]]
          output_shape += [input_shape[-1]]
          return tuple(output_shape)
  
      def _single_resize(self, inputs):
          return resize(inputs, self.zoom_factor, interp_method=self.interp_method)
  
  
  # Zoom naming of resize, to match scipy's naming
  Zoom = Resize
  
  
  class VecInt(Layer):
      """
      Vector Integration Layer
  
      Enables vector integration via several methods 
      (ode or quadrature for time-dependent vector fields, 
      scaling and squaring for stationary fields)
  
      If you find this function useful, please cite:
        Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration
        Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu
        MICCAI 2018.
      """
  
      def __init__(self, indexing='ij', method='ss', int_steps=7, **kwargs):
          """        
          Parameters:
              method can be any of the methods in neuron.utils.integrate_vec
          """
  
          assert indexing in ['ij', 'xy'], "indexing has to be 'ij' (matrix) or 'xy' (cartesian)"
          self.indexing = indexing
          self.method = method
          self.int_steps = int_steps
          self.inshape = None
          super(self.__class__, self).__init__(**kwargs)
  
      def build(self, input_shape):
          # confirm built
          self.built = True
          self.inshape = input_shape
  
          if input_shape[-1] != len(input_shape) - 2:
              raise Exception('transform ndims %d does not match expected ndims %d' \
                  % (input_shape[-1], len(input_shape) - 2))
  
      def call(self, inputs):
          loc_shift = inputs
  
          # necessary for multi_gpu models...
          loc_shift = K.reshape(loc_shift, [-1, *self.inshape[1:]])
          # loc_shift._keras_shape = inputs._keras_shape
          
          # prepare location shift
          if self.indexing == 'xy':  # shift the first two dimensions
              loc_shift_split = tf.split(loc_shift, loc_shift.shape[-1], axis=-1)
              loc_shift_lst = [loc_shift_split[1], loc_shift_split[0], *loc_shift_split[2:]]
              loc_shift = tf.concat(loc_shift_lst, -1)
  
          # map transform across batch
          out = tf.map_fn(self._single_int, loc_shift, dtype=tf.float32)
          # out._keras_shape = inputs._keras_shape
          return out
  
      def _single_int(self, inputs):
  
          vel = inputs
          return integrate_vec(vel, method=self.method,
                        nb_steps=self.int_steps,
                        ode_args={'rtol':1e-6, 'atol':1e-12},
                        time_pt=1)
  