"""
tensorflow/keras utilities for the neuron project
If you use this code, please cite
Dalca AV, Guttag J, Sabuncu MR
Anatomical Priors in Convolutional Networks for Unsupervised Biomedical Segmentation,
CVPR 2018
or for the transformation/integration functions:
Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration
Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu
MICCAI 2018.
Contact: adalca [at] csail [dot] mit [dot] edu
License: GPLv3
"""
# third party
import tensorflow as tf
from keras import backend as K
from keras.layers import Layer
from copy import deepcopy
# local
from ext.neuron.utils import transform, resize, integrate_vec, affine_to_shift, combine_non_linear_and_aff_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
self.is_affine = list()
assert indexing in ['ij', 'xy'], "indexing has to be 'ij' (matrix) or 'xy' (cartesian)"
self.indexing = indexing
super(self.__class__, self).__init__(**kwargs)
def get_config(self):
config = super().get_config()
config["interp_method"] = self.interp_method
config["indexing"] = self.indexing
config["single_transform"] = self.single_transform
return config
def build(self, input_shape):
"""
input_shape should be a list for two inputs:
input1: image.
input2: list of transform Tensors
if affine:
should be an N+1 x N+1 matrix
*or* a N+1*N+1 tensor (which will be reshaped to N x (N+1) and an identity row added)
if not affine:
should be a *vol_shape x N
"""
if len(input_shape) > 3:
raise Exception('Spatial Transformer must be called on a list of min length 2 and max length 3.'
'First argument is the image followed by the affine and non linear transforms.')
# set up number of dimensions
self.ndims = len(input_shape[0]) - 2
self.inshape = input_shape
trf_shape = [trans_shape[1:] for trans_shape in input_shape[1:]]
for (i, shape) in enumerate(trf_shape):
# 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].
self.is_affine.append(len(shape) == 1 or
(len(shape) == 2 and all([f == (self.ndims + 1) for f in shape])))
# check sizes
if self.is_affine[i] and len(shape) == 1:
ex = self.ndims * (self.ndims + 1)
if shape[0] != ex:
raise Exception('Expected flattened affine of len %d but got %d' % (ex, shape[0]))
if not self.is_affine[i]:
if shape[-1] != self.ndims:
raise Exception('Offset flow field size expected: %d, found: %d' % (self.ndims, shape[-1]))
# confirm built
self.built = True
def call(self, inputs, **kwargs):
"""
Parameters
inputs: list with several entries: the volume followed by the transforms
"""
# check shapes
assert 1 < len(inputs) < 4, "inputs has to be len 2 or 3, found: %d" % len(inputs)
vol = inputs[0]
trf = inputs[1:]
# necessary for multi_gpu models...
vol = K.reshape(vol, [-1, *self.inshape[0][1:]])
for i in range(len(trf)):
trf[i] = K.reshape(trf[i], [-1, *self.inshape[i+1][1:]])
# reorder transforms, non-linear first and affine second
ind_nonlinear_linear = [i[0] for i in sorted(enumerate(self.is_affine), key=lambda x:x[1])]
self.is_affine = [self.is_affine[i] for i in ind_nonlinear_linear]
self.inshape = [self.inshape[i] for i in ind_nonlinear_linear]
trf = [trf[i] for i in ind_nonlinear_linear]
# go from affine to deformation field
if len(trf) == 1:
trf = trf[0]
if self.is_affine[0]:
trf = tf.map_fn(lambda x: self._single_aff_to_shift(x, vol.shape[1:-1]), trf, dtype=tf.float32)
# combine non-linear and affine to obtain a single deformation field
elif len(trf) == 2:
trf = tf.map_fn(lambda x: self._non_linear_and_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:
return tf.map_fn(self._single_transform, [vol, trf[0, :]], 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])
return affine_to_shift(trf, volshape, shift_center=True)
def _non_linear_and_aff_to_shift(self, trf, volshape):
if len(trf[1].shape) == 1: # go from vector to matrix
trf[1] = tf.reshape(trf[1], [self.ndims, self.ndims + 1])
return combine_non_linear_and_aff_to_shift(trf, volshape, shift_center=True)
def _single_transform(self, inputs):
return transform(inputs[0], inputs[1], interp_method=self.interp_method)
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, out_time_pt=1,
ode_args=None,
odeint_fn=None, **kwargs):
"""
Parameters:
method can be any of the methods in neuron.utils.integrate_vec
indexing can be 'xy' (switches first two dimensions) or 'ij'
int_steps is the number of integration steps
out_time_pt is time point at which to output if using odeint integration
"""
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
self.out_time_pt = out_time_pt
self.odeint_fn = odeint_fn # if none then will use a tensorflow function
self.ode_args = ode_args
if ode_args is None:
self.ode_args = {'rtol': 1e-6, 'atol': 1e-12}
super(self.__class__, self).__init__(**kwargs)
def get_config(self):
config = super().get_config()
config["indexing"] = self.indexing
config["method"] = self.method
config["int_steps"] = self.int_steps
config["out_time_pt"] = self.out_time_pt
config["ode_args"] = self.ode_args
config["odeint_fn"] = self.odeint_fn
return config
def build(self, input_shape):
# confirm built
self.built = True
trf_shape = input_shape
if isinstance(input_shape[0], (list, tuple)):
trf_shape = input_shape[0]
self.inshape = trf_shape
if trf_shape[-1] != len(trf_shape) - 2:
raise Exception('transform ndims %d does not match expected ndims %d' % (trf_shape[-1], len(trf_shape) - 2))
def call(self, inputs, **kwargs):
if not isinstance(inputs, (list, tuple)):
inputs = [inputs]
loc_shift = inputs[0]
# necessary for multi_gpu models...
loc_shift = K.reshape(loc_shift, [-1, *self.inshape[1:]])
# 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)
if len(inputs) > 1:
assert self.out_time_pt is None, 'out_time_pt should be None if providing batch_based out_time_pt'
# map transform across batch
out = tf.map_fn(self._single_int, [loc_shift] + inputs[1:], dtype=tf.float32)
return out
def _single_int(self, inputs):
vel = inputs[0]
out_time_pt = self.out_time_pt
if len(inputs) == 2:
out_time_pt = inputs[1]
return integrate_vec(vel, method=self.method,
nb_steps=self.int_steps,
ode_args=self.ode_args,
out_time_pt=out_time_pt,
odeint_fn=self.odeint_fn)
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=None,
size=None,
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.size = list(size)
self.zoom_factor0 = None
self.size0 = None
self.interp_method = interp_method
self.ndims = None
self.inshape = None
super(Resize, self).__init__(**kwargs)
def get_config(self):
config = super().get_config()
config["zoom_factor"] = self.zoom_factor
config["size"] = self.size
config["interp_method"] = self.interp_method
return config
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.')
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
# check zoom_factor
if isinstance(self.zoom_factor, float):
self.zoom_factor0 = [self.zoom_factor] * self.ndims
elif self.zoom_factor is None:
self.zoom_factor0 = [0] * self.ndims
elif isinstance(self.zoom_factor, (list, tuple)):
self.zoom_factor0 = deepcopy(self.zoom_factor)
assert len(self.zoom_factor0) == self.ndims, \
'zoom factor length {} does not match number of dimensions {}'.format(len(self.zoom_factor), self.ndims)
else:
raise Exception('zoom_factor should be an int or a list/tuple of int (or None if size is not set to None)')
# check size
if isinstance(self.size, int):
self.size0 = [self.size] * self.ndims
elif self.size is None:
self.size0 = [0] * self.ndims
elif isinstance(self.size, (list, tuple)):
self.size0 = deepcopy(self.size)
assert len(self.size0) == self.ndims, \
'size length {} does not match number of dimensions {}'.format(len(self.size0), self.ndims)
else:
raise Exception('size should be an int or a list/tuple of int (or None if zoom_factor is not set to None)')
# confirm built
self.built = True
super(Resize, self).build(input_shape) # Be sure to call this somewhere!
def call(self, inputs, **kwargs):
"""
Parameters
inputs: volume or list of 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:]])
# set value of missing size or zoom_factor
if not any(self.zoom_factor0):
self.zoom_factor0 = [self.size0[i] / self.inshape[i+1] for i in range(self.ndims)]
else:
self.size0 = [int(self.inshape[f+1] * self.zoom_factor0[f]) for f in range(self.ndims)]
# map transform across batch
return tf.map_fn(self._single_resize, vol, dtype=vol.dtype)
def compute_output_shape(self, input_shape):
output_shape = [input_shape[0]]
output_shape += [int(input_shape[1:-1][f] * self.zoom_factor0[f]) for f in range(self.ndims)]
output_shape += [input_shape[-1]]
return tuple(output_shape)
def _single_resize(self, inputs):
return resize(inputs, self.zoom_factor0, self.size0, interp_method=self.interp_method)
# Zoom naming of resize, to match scipy's naming
Zoom = Resize
#########################################################
# "Local" layers -- layers with parameters at each voxel
#########################################################
class LocalBias(Layer):
"""
Local bias layer: each pixel/voxel has its own bias operation (one parameter)
out[v] = in[v] + b
"""
def __init__(self, my_initializer='RandomNormal', biasmult=1.0, **kwargs):
self.initializer = my_initializer
self.biasmult = biasmult
self.kernel = None
super(LocalBias, self).__init__(**kwargs)
def get_config(self):
config = super().get_config()
config["my_initializer"] = self.initializer
config["biasmult"] = self.biasmult
return config
def build(self, input_shape):
# Create a trainable weight variable for this layer.
self.kernel = self.add_weight(name='kernel',
shape=input_shape[1:],
initializer=self.initializer,
trainable=True)
super(LocalBias, self).build(input_shape) # Be sure to call this somewhere!
def call(self, x, **kwargs):
return x + self.kernel * self.biasmult # weights are difference from input
def compute_output_shape(self, input_shape):
return input_shape
