def cvpr2018_net(vol_size,
enc_nf,
dec_nf,
full_size=True,
indexing='ij',
use_seg=False,
n_seg=2):
"""
unet architecture for voxelmorph models presented in the CVPR 2018 paper.
You may need to modify this code (e.g., number of layers) to suit your project needs.
:param vol_size: volume size. e.g. (256, 256, 256)
:param enc_nf: list of encoder filters. right now it needs to be 1x4.
e.g. [16,32,32,32]
:param dec_nf: list of decoder filters. right now it must be 1x6 (like voxelmorph-1) or 1x7 (voxelmorph-2)
:return: the keras model
"""
ndims = len(vol_size)
assert ndims in [1, 2,
3], "ndims should be one of 1, 2, or 3. found: %d" % ndims
# get the core model
unet_model = unet_core(vol_size, enc_nf, dec_nf, full_size=full_size)
[src, tgt] = unet_model.inputs
x = unet_model.output
# transform the results into a flow field.
Conv = getattr(KL, 'Conv%dD' % ndims)
flow = Conv(ndims,
kernel_size=3,
padding='same',
name='flow',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5))(x)
# warp the source with the flow
y = nrn_layers.SpatialTransformer(interp_method='linear',
indexing=indexing)([src, flow])
# prepare model
if not use_seg:
model = Model(inputs=[src, tgt], outputs=[y, flow])
else:
downfac = 2
src_seg = Input(shape=(vol_size[0] // downfac, vol_size[1] // downfac,
vol_size[2] // downfac, n_seg))
flow_dn = Lambda(interp_downsampling)(flow)
flow_dn = Lambda(lambda arg: arg / 2.0)(flow_dn)
y_seg = nrn_layers.SpatialTransformer(
interp_method='linear', indexing=indexing)([src_seg, flow_dn])
model = Model(inputs=[src, tgt, src_seg], outputs=[y, flow, y_seg])
return model
def cvpr2018_net(vol_size, enc_nf, dec_nf, full_size=True, indexing='ij'):
"""
unet architecture for voxelmorph models presented in the CVPR 2018 paper.
You may need to modify this code (e.g., number of layers) to suit your project needs.
:param vol_size: volume size. e.g. (256, 256, 256) -> 160x192x224
:param enc_nf: list of encoder filters. right now it needs to be 1x4.
e.g. [16,32,32,32]
:param dec_nf: list of decoder filters. right now it must be 1x6 (like voxelmorph-1) or 1x7 (voxelmorph-2)
:return: the keras model e.g.[32, 32, 32, 32, 32, 16, 16]
"""
ndims = len(vol_size) # -> 3
assert ndims in [1, 2, 3], "ndims should be one of 1, 2, or 3. found: %d" % ndims
# get the core model
unet_model = unet_core(vol_size, enc_nf, dec_nf, full_size=full_size)
[src, tgt] = unet_model.inputs
x = unet_model.output
# transform the results into a flow field.
Conv = getattr(KL, 'Conv%dD' % ndims)
flow = Conv(ndims, kernel_size=3, padding='same', name='flow',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5))(x)
# warp the source with the flow
# indexing = 'ij'
y = nrn_layers.SpatialTransformer(interp_method='linear', indexing=indexing)([src, flow])
# print(y.shape) (?, 160, 192, 224, 1)
# src: ?x160x192x224x1 flow: ?x160x192x224x3
# prepare model
model = Model(inputs=[src, tgt], outputs=[y, flow])
return model
def cvpr2018_net(vol_size, enc_nf, dec_nf, full_size=True, indexing='ij'):
ndims = len(vol_size)
assert ndims in [1, 2,
3], "ndims should be one of 1, 2, or 3. found: %d" % ndims
# get the core model
unet_model = unet_core(vol_size, enc_nf, dec_nf, full_size=full_size)
[src, tgt] = unet_model.inputs
x = unet_model.output
# transform the results into a flow field.
Conv = getattr(KL, 'Conv%dD' % ndims)
flow = Conv(ndims,
kernel_size=3,
padding='same',
name='flow',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5))(x)
# warp the source with the flow
y = nrn_layers.SpatialTransformer(interp_method='linear',
indexing=indexing)([src, flow])
# prepare model
model = Model(inputs=[src, tgt], outputs=[y, flow])
return model
def trf_core(vol_size,
interp_method='linear',
indexing='ij',
nb_feats=1,
int_steps=0):
"""
Simple transform model
Note: this is essentially a wrapper for the neuron.utils.transform(..., interp_method='nearest')
"""
ndims = len(vol_size)
# nn warp model
subj_input = Input((*vol_size, nb_feats), name='subj_input')
trf_input = Input((*vol_size, ndims), name='trf_input')
if int_steps > 0:
trf = nrn_layers.VecInt(method='ss',
name='trf-int',
int_steps=int_steps)(trf_input)
else:
trf = trf_input
# note the nearest neighbour interpolation method
# use xy indexing when Guha's original code switched x and y dimensions
nn_output = nrn_layers.SpatialTransformer(interp_method=interp_method,
indexing=indexing)
nn_spatial_output = nn_output([subj_input, trf])
return keras.models.Model([subj_input, trf_input], nn_spatial_output)
def interp_upsampling(V):
"""
upsample a field by a factor of 2
TODO: should switch this to use neuron.utils.interpn()
"""
#print(V.shape)
grid = nrn_utils.volshape_to_ndgrid(
[f * 2 for f in V.get_shape().as_list()[1:-1]])
grid = [tf.cast(f, 'float32') for f in grid]
grid = [tf.expand_dims(f / 2 - f, 0) for f in grid]
offset = tf.stack(grid, len(grid) + 1)
#xx = tf.cast(xx, 'float32')
#yy = tf.cast(yy, 'float32')
#zz = tf.cast(zz, 'float32')
#xx = tf.expand_dims(xx/2-xx, 0)
#yy = tf.expand_dims(yy/2-yy, 0)
#zz = tf.expand_dims(zz/2-zz, 0)
#print(xx.shape)
#offset = tf.stack([xx, yy], 3)
#print(offset.shape)
#print(V.shape)
# V = nrn_utils.transform(V, offset)
V = nrn_layers.SpatialTransformer(interp_method='linear')([V, offset])
#print(V.shape)
return V
def unet(vol_size, enc_nf, dec_nf, full_size=True):
"""
unet architecture for voxelmorph models presented in the CVPR 2018 paper.
You may need to modify this code (e.g., number of layers) to suit your project needs.
:param vol_size: volume size. e.g. (256, 256, 256)
:param enc_nf: list of encoder filters. right now it needs to be 1x4.
e.g. [16,32,32,32]
:param dec_nf: list of decoder filters. right now it must be 1x6 (like voxelmorph-1) or 1x7 (voxelmorph-2)
:return: the keras model
"""
# get the core model
unet_model = unet_core(vol_size, enc_nf, dec_nf, full_size=full_size)
[src, tgt] = unet_model.inputs
x = unet_model.output
# transform the results into a flow field.
flow = Conv3D(3, kernel_size=3, padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5), name='flow')(x)
# warp the source with the flow
y = nrn_layers.SpatialTransformer(interp_method='linear', indexing='xy')([src, flow])
# prepare model
model = Model(inputs=[src, tgt], outputs=[y, flow])
return model
def nn_trf(vol_size, indexing='xy'):
ndims = len(vol_size)
# nn warp model
subj_input = Input((*vol_size, 1), name='subj_input')
trf_input = Input((*vol_size, ndims), name='trf_input')
# note the nearest neighbour interpolation method
# note xy indexing because Guha's original code switched x and y dimensions
nn_output = nrn_layers.SpatialTransformer(interp_method='nearest',
indexing=indexing)
nn_spatial_output = nn_output([subj_input, trf_input])
return keras.models.Model([subj_input, trf_input], nn_spatial_output)
def interp_upsampling(V):
"""
upsample a field by a factor of 2
TODO: should switch this to use neuron.utils.interpn()
"""
grid = nrn_utils.volshape_to_ndgrid([f*2 for f in V.get_shape().as_list()[1:-1]])
grid = [tf.cast(f, 'float32') for f in grid]
grid = [tf.expand_dims(f/2 - f, 0) for f in grid]
offset = tf.stack(grid, len(grid) + 1)
# V = nrn_utils.transform(V, offset)
V = nrn_layers.SpatialTransformer(interp_method='linear')([V, offset])
return V
def _lseg(self, warp):
print('warp is:', warp)
Sm_ten = tf.convert_to_tensor(self.Sm_)
Sm_tensor = tf.to_float(Sm_ten, name='ToFloat')
# print(Sm_tensor)
y = nrn_layers.SpatialTransformer(interp_method='nearest',
indexing='ij')([Sm_tensor,
warp]) # nearest
# print(y)
Sf_ten = tf.convert_to_tensor(self.Sf_)
Sf_tensor = tf.to_float(Sf_ten, name='ToFloat')
# print(Sf_tensor)
dicem = binary_dice(Sf_tensor, y)
return dicem
def nn_trf(vol_size):
"""
Simple transform model for nearest-neighbor based transformation
Note: this is essentially a wrapper for the neuron.utils.transform(..., interp_method='nearest')
"""
ndims = len(vol_size)
# nn warp model
subj_input = Input((*vol_size, 1), name='subj_input')
trf_input = Input((*vol_size, ndims) , name='trf_input')
# note the nearest neighbour interpolation method
# note xy indexing because Guha's original code switched x and y dimensions
nn_output = nrn_layers.SpatialTransformer(interp_method='nearest', indexing='xy')
nn_spatial_output = nn_output([subj_input, trf_input])
return keras.models.Model([subj_input, trf_input], nn_spatial_output)
def make_flow(self):
[Atlas, Moving] = self.unet_model.inputs
x = self.unet_model.output
flow = Conv3D(3,
kernel_size=3,
padding='same',
name='flow',
kernel_initializer=RandomNormal(mean=0.0,
stddev=1e-5))(x)
# warp the source with the flow
y = nrn_layers.SpatialTransformer(
interp_method='linear', indexing=self.indexing)([Moving, flow])
# prepare model
self.model = Model(inputs=[Atlas, Moving], outputs=[y, flow])
def ncc_l(I, J, flow, l):
"""
local (over window) normalized cross correlation
"""
eps = 1e-5
J = nrn_layers.SpatialTransformer(interp_method='linear', indexing='ij')([J, flow])
ndims = len(I.get_shape().as_list()) - 2
assert ndims in [1, 2, 3], "volumes should be 1 to 3 dimensions. found: %d" % ndims
# set window size
win = [l] * ndims
# get convolution function
conv_fn = getattr(tf.nn, 'conv%dd' % ndims)
# compute CC squares
I2 = I * I
J2 = J * J
IJ = I * J
# compute filters
sum_filt = tf.ones([*win, 1, 1])
strides = [1] * (ndims + 2)
padding = 'SAME'
# compute local sums via convolution
I_sum = conv_fn(I, sum_filt, strides, padding)
J_sum = conv_fn(J, sum_filt, strides, padding)
I2_sum = conv_fn(I2, sum_filt, strides, padding)
J2_sum = conv_fn(J2, sum_filt, strides, padding)
IJ_sum = conv_fn(IJ, sum_filt, strides, padding)
# compute cross correlation
win_size = np.prod(win)
u_I = I_sum / win_size
u_J = J_sum / win_size
cross = IJ_sum - u_J * I_sum - u_I * J_sum + u_I * u_J * win_size
I_var = I2_sum - 2 * u_I * I_sum + u_I * u_I * win_size
J_var = J2_sum - 2 * u_J * J_sum + u_J * u_J * win_size
cc = cross * cross / (I_var * J_var + eps)
# return negative cc.
return -tf.reduce_mean(cc)
def __call__(self, images_0, images_1, reuse=False):# reuse=False
with tf.variable_scope(self.name, reuse=reuse) as vs:
x_in = tf.concat([images_1, images_0], axis=-1)
# down-sample path (encoder)
x_enc = [x_in]
for i in range(len(self.enc_nf)):
x_enc.append(self.conv_block(x_enc[-1], self.enc_nf[i], 2))
# up-sample path (decoder)
x = self.conv_block(x_enc[-1], self.dec_nf[0])
x = UpSampling3D()(x)
x = tf.concat([x, x_enc[-2]], axis=-1)
x = self.conv_block(x, self.dec_nf[1])
x = UpSampling3D()(x)
x = tf.concat([x, x_enc[-3]], axis=-1)
x = self.conv_block(x, self.dec_nf[2])
x = UpSampling3D()(x)
x = tf.concat([x, x_enc[-4]], axis=-1)
x = self.conv_block(x, self.dec_nf[3])
x = self.conv_block(x, self.dec_nf[4])
# only upsampleto full dim if full_size
# here we explore architectures where we essentially work with flow fields
# that are 1/2 size
x = UpSampling3D()(x)
x = tf.concat([x, x_enc[0]], axis=-1)
x = self.conv_block(x, self.dec_nf[5])
# optional convolution at output resolution (used in voxelmorph-2)
if len(self.dec_nf) == 7:
x = self.conv_block(x, self.dec_nf[6])
# transform the results into a flow field.
# flow = tf.layers.Conv3D(3, kernel_size=3, padding='same', name='flow',
# kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5))(x)
flow = tf.layers.Conv3D(3, (3, 3, 3), (1, 1, 1), 'same', kernel_initializer='he_normal')(x)
# warp the source with the flow
y = nrn_layers.SpatialTransformer(interp_method='linear', indexing='ij')([images_1, flow])
return flow, y
def interp_upsampling(V):
"""
upsample a field by a factor of 2
TODO: should switch this to use neuron.utils.interpn()
"""
[xx, yy, zz] = nrn_utils.volshape_to_ndgrid([f*2 for f in V.get_shape().as_list()[1:4]])
xx = tf.cast(xx, 'float32')
yy = tf.cast(yy, 'float32')
zz = tf.cast(zz, 'float32')
xx = tf.expand_dims(xx/2-xx, 0)
yy = tf.expand_dims(yy/2-yy, 0)
zz = tf.expand_dims(zz/2-zz, 0)
offset = tf.stack([xx, yy, zz], 4)
# V = nrn_utils.transform(V, offset)
V = nrn_layers.SpatialTransformer(interp_method='linear')([V, offset])
return V
def __call__(self, images_0, images_1, reuse=False): # reuse=False
with tf.variable_scope(self.name, reuse=reuse) as vs:
pyramid_0, pyramid_params_0 = self.fp_extractor(images_0,
reuse=reuse)
pyramid_1, pyramid_params_1 = self.fp_extractor(images_1)
flows_pyramid = []
flows_up, features_up = None, None
for l, (features_0,
features_1) in enumerate(zip(pyramid_0, pyramid_1)):
# Flow estimation
flows = self.of_estimator[l](features_0, features_1, flows_up)
# Integrate if diffeomorphic (i.e. treating 'flow' above as stationary velocity field)
z_sample = flows
flows = nrn_layers.VecInt(method='ss',
name='flow-int',
int_steps=self.int_steps)(z_sample)
if l < self.output_level:
# up-sample
flows_up = nrn_layers.Resize(zoom_factor=2,
interp_method='linear')(
flows * 2)
else:
# At output level
flows_pyramid.append(flows)
# Obtain finally scale-adjusted flow
upscale = 2**(self.num_levels - self.output_level)
flows_final = nrn_layers.Resize(zoom_factor=upscale,
interp_method='linear')(
flows * upscale)
y = nrn_layers.SpatialTransformer(interp_method='linear',
indexing='ij')([
images_1, flows_final
])
return flows_final, y, flows_pyramid, pyramid_params_0, pyramid_params_1
flows_pyramid.append(flows)
def miccai2018_net(vol_size, enc_nf, dec_nf, use_miccai_int=True, int_steps=7, indexing='xy'):
"""
architecture for probabilistic diffeomoprhic VoxelMorph presented in the MICCAI 2018 paper.
You may need to modify this code (e.g., number of layers) to suit your project needs.
The stationary velocity field operates in a space (0.5)^3 of vol_size for computational reasons.
:param vol_size: volume size. e.g. (256, 256, 256)
:param enc_nf: list of encoder filters. right now it needs to be 1x4.
e.g. [16,32,32,32]
:param dec_nf: list of decoder filters. right now it must be 1x6, see unet function.
:param use_miccai_int: whether to use the manual miccai implementation of scaling and squaring integration
note that the 'velocity' field outputted in that case was
since then we've updated the code to be part of a flexible layer. see neuron.layers.VecInt
:param int_steps: the number of integration steps
:param indexing: xy or ij indexing. we recommend ij indexing if training from scratch.
miccai 2018 runs were done with xy indexing.
:return: the keras model
"""
# get unet
unet_model = unet_core(vol_size, enc_nf, dec_nf, full_size=False)
[src,tgt] = unet_model.inputs
x_out = unet_model.outputs[-1]
# velocity mean and logsigma layers
flow_mean = Conv3D(3, kernel_size=3, padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5), name='flow')(x_out)
flow_log_sigma = Conv3D(3, kernel_size=3, padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-10),
bias_initializer=keras.initializers.Constant(value=-10), name='log_sigma')(x_out)
flow_params = concatenate([flow_mean, flow_log_sigma])
# velocity sample
flow = Lambda(sample, name="z_sample")([flow_mean, flow_log_sigma])
# integrate if diffeomorphic (i.e. treating 'flow' above as stationary velocity field)
if use_miccai_int:
# for the miccai2018 submission, the scaling and squaring layer
# was manually composed of a Transform and and Add Layer.
flow = Lambda(lambda x: x, name='flow-fix')(flow) # remanant of old code
v = flow
for _ in range(int_steps):
v1 = nrn_layers.SpatialTransformer(interp_method='linear', indexing=indexing)([v, v])
v = keras.layers.add([v, v1])
flow = v
else:
# new implementation in neuron is cleaner.
# the 2**int_steps is a correcting factor left over from the miccai implementation.
# * (2**int_steps)
flow = Lambda(lambda x: x, name='flow-fix')(flow)
flow = nrn_layers.VecInt(method='ss', name='flow-int', int_steps=7)(flow)
# get up to final resolution
flow = Lambda(interp_upsampling, output_shape=vol_size+(3,), name='pre_diffflow')(flow)
flow = Lambda(lambda arg: arg*2, name='diffflow')(flow)
# transform
y = nrn_layers.SpatialTransformer(interp_method='linear', indexing=indexing)([src, flow])
# prepare outputs and losses
outputs = [y, flow_params]
# build the model
return Model(inputs=[src, tgt], outputs=outputs)
def __call__(self, images_1_seg, flow, indexing='ij'):
out = nrn_layers.SpatialTransformer(
interp_method='nearest', indexing=indexing)([images_1_seg, flow])
return out
def miccai2018_net(vol_size,
enc_nf,
dec_nf,
int_steps=7,
use_miccai_int=False,
indexing='ij',
bidir=False,
vel_resize=1 / 2):
"""
architecture for probabilistic diffeomoprhic VoxelMorph presented in the MICCAI 2018 paper.
You may need to modify this code (e.g., number of layers) to suit your project needs.
The stationary velocity field operates in a space (0.5)^3 of vol_size for computational reasons.
:param vol_size: volume size. e.g. (256, 256, 256)
:param enc_nf: list of encoder filters. right now it needs to be 1x4.
e.g. [16,32,32,32]
:param dec_nf: list of decoder filters. right now it must be 1x6, see unet function.
:param use_miccai_int: whether to use the manual miccai implementation of scaling and squaring integration
note that the 'velocity' field outputted in that case was
since then we've updated the code to be part of a flexible layer. see neuron.layers.VecInt
**This param will be phased out (set to False behavior)**
:param int_steps: the number of integration steps
:param indexing: xy or ij indexing. we recommend ij indexing if training from scratch.
miccai 2018 runs were done with xy indexing.
**This param will be phased out (set to 'ij' behavior)**
:return: the keras model
"""
ndims = len(vol_size)
assert ndims in [1, 2,
3], "ndims should be one of 1, 2, or 3. found: %d" % ndims
# get unet
unet_model = unet_core(vol_size, enc_nf, dec_nf, full_size=False)
[src, tgt] = unet_model.inputs
x_out = unet_model.outputs[-1]
# velocity mean and logsigma layers
Conv = getattr(KL, 'Conv%dD' % ndims)
flow_mean = Conv(ndims,
kernel_size=3,
padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5),
name='flow')(x_out)
# we're going to initialize the velocity variance very low, to start stable.
flow_log_sigma = Conv(
ndims,
kernel_size=3,
padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-10),
bias_initializer=keras.initializers.Constant(value=-10),
name='log_sigma')(x_out)
flow_params = concatenate([flow_mean, flow_log_sigma])
# velocity sample
flow = Sample(name="z_sample")([flow_mean, flow_log_sigma])
# integrate if diffeomorphic (i.e. treating 'flow' above as stationary velocity field)
if use_miccai_int:
# for the miccai2018 submission, the squaring layer
# scaling was essentially built in by the network
# was manually composed of a Transform and and Add Layer.
v = flow
for _ in range(int_steps):
v1 = nrn_layers.SpatialTransformer(interp_method='linear',
indexing=indexing)([v, v])
v = keras.layers.add([v, v1])
flow = v
else:
# new implementation in neuron is cleaner.
z_sample = flow
flow = nrn_layers.VecInt(method='ss',
name='flow-int',
int_steps=int_steps)(z_sample)
if bidir:
rev_z_sample = Negate()(z_sample)
neg_flow = nrn_layers.VecInt(method='ss',
name='neg_flow-int',
int_steps=int_steps)(rev_z_sample)
# get up to final resolution
flow = trf_resize(flow, vel_resize, name='diffflow')
if bidir:
neg_flow = trf_resize(neg_flow, vel_resize, name='neg_diffflow')
# transform
y = nrn_layers.SpatialTransformer(interp_method='linear',
indexing=indexing)([src, flow])
if bidir:
y_tgt = nrn_layers.SpatialTransformer(
interp_method='linear', indexing=indexing)([tgt, neg_flow])
# prepare outputs and losses
outputs = [y, flow_params]
if bidir:
outputs = [y, y_tgt, flow_params]
# build the model
return Model(inputs=[src, tgt], outputs=outputs)
def rigid_net(vol_size, enc_nf, dec_nf):
"""
architecture for rigid registration.
rigid registration: CNN output ND x ND+1 affine matrix, affine matrix is passed to affine_to_shift function
:param vol_size: volume size
:param enc_nf: list of encoder
:param dec_nf: list of decoder
:return: model
"""
unet_model = unet_core(vol_size, enc_nf, dec_nf, full_size=False)
[src, tgt] = unet_model.inputs
x_out = unet_model.outputs[-1]
# affine transform matrix
# build full connected layer into the model, output the ND x ND+1 affine matrix
flow = Conv3D(3,
kernel_size=3,
padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5),
name='flow')(x_out)
#flow = LeakyReLU(0.2)(flow)
flow1 = Lambda(reduce_dim1)(flow)
flow1 = Lambda(my_reshape)(flow1)
print("flow1's shape is: " + str(flow1.shape))
flow2 = Lambda(reduce_dim2)(flow)
flow2 = Lambda(my_reshape)(flow2)
flow3 = Lambda(reduce_dim3)(flow)
flow3 = Lambda(my_reshape)(flow3)
# add convolutinal layer into the model, which outputs affine matrix.
affine_matrix1 = Conv3D(filters=4,
kernel_size=(80, 96, 112),
padding='valid',
kernel_initializer=RandomNormal(mean=0.0,
stddev=1e-5),
name='flow1')(flow1)
affine_matrix1 = Reshape((1, 4))(affine_matrix1)
#affine_matrix1 = LeakyReLU(0.2)(affine_matrix1)
print("affine_matrix1's shape is: " + str(affine_matrix1.shape))
affine_matrix2 = Conv3D(filters=4,
kernel_size=(80, 96, 112),
padding='valid',
kernel_initializer=RandomNormal(mean=0.0,
stddev=1e-5),
name='flow2')(flow2)
affine_matrix2 = Reshape((1, 4))(affine_matrix2)
#affine_matrix2 = LeakyReLU(0.2)(affine_matrix2)
affine_matrix3 = Conv3D(filters=4,
kernel_size=(80, 96, 112),
padding='valid',
kernel_initializer=RandomNormal(mean=0.0,
stddev=1e-5),
name='flow3')(flow3)
affine_matrix3 = Reshape((1, 4))(affine_matrix3)
#affine_matrix3 = LeakyReLU(0.2)(affine_matrix3)
affine_matrix = concatenate([affine_matrix1, affine_matrix2])
affine_matrix = concatenate([affine_matrix, affine_matrix3])
affine_matrix = Lambda(reduce_dim4)(affine_matrix)
print("affine_matrix's shape is :" + str(affine_matrix.shape))
# spatial transform
y = nrn_layers.SpatialTransformer(interp_method='linear',
indexing='xy')([src, affine_matrix])
model = Model(inputs=[src, tgt], outputs=[y, flow])
print("the output's shape is:" + str(y.shape))
return model
def keras_block(flow, src, interp_method='linear'):
y = nrn_layers.SpatialTransformer(interp_method=interp_method,
indexing=indexing)([src, flow])
return y
def miccai2018_net(vol_size, enc_nf, dec_nf, int_steps=6, indexing='ij', vel_resize=0.5):
"""
architecture for probabilistic diffeomoprhic VoxelMorph presented in the MICCAI 2018 paper.
You may need to modify this code (e.g., number of layers) to suit your project needs.
The stationary velocity field operates in a space (0.5)^3 of vol_size for computational reasons.
:param vol_size: volume size. e.g. (256, 256, 256)
:param enc_nf: list of encoder filters. right now it needs to be 1x4.
e.g. [16,32,32,32]
:param dec_nf: list of decoder filters. right now it must be 1x6, see unet function.
:param use_miccai_int: whether to use the manual miccai implementation of scaling and squaring integration
note that the 'velocity' field outputted in that case was
since then we've updated the code to be part of a flexible layer. see neuron.layers.VecInt
**This param will be phased out (set to False behavior)**
:param int_steps: the number of integration steps
:param indexing: xy or ij indexing. we recommend ij indexing if training from scratch.
miccai 2018 runs were done with xy indexing.
**This param will be phased out (set to 'ij' behavior)**
:return: the keras model
"""
ndims = len(vol_size)
assert ndims in [1, 2, 3], "ndims should be one of 1, 2, or 3. found: %d" % ndims
# get unet
unet_model = unet_core(vol_size, enc_nf, dec_nf, vel_resize)
# target delta in binary representation
b_in = Input(shape=(16,))
x_in = unet_model.inputs[0]
x_out = unet_model.outputs[-1]
# velocity mean and logsigma layers
Conv = getattr(KL, 'Conv%dD' % ndims)
flow_mean = Conv(ndims, kernel_size=3, padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5), name='flow')(x_out)
# we're going to initialize the velocity variance very low, to start stable.
flow_log_sigma = Conv(ndims, kernel_size=3, padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-10),
bias_initializer=keras.initializers.Constant(value=-10),
name='log_sigma')(x_out)
flow_params = concatenate([flow_mean, flow_log_sigma])
# velocity sample
flow = Sample(name="z_sample")([flow_mean, flow_log_sigma])
# integrate if diffeomorphic (i.e. treating 'flow' above as stationary velocity field)
z_sample = flow
flow = nrn_layers.VecInt(method='ss', name='flow-int', int_steps=int_steps)([z_sample, b_in])
# get up to final resolution
if vel_resize != 1.0:
flow = trf_resize(flow, vel_resize, name='diff-flow')
# transform
y = nrn_layers.SpatialTransformer(interp_method='linear', indexing=indexing)([x_in, flow])
# prepare outputs
outputs = [y, flow_params]
# build the model
return Model(inputs=[x_in, b_in], outputs=outputs)
def cvpr2018_net_probatlas(vol_size, enc_nf, dec_nf, nb_labels,
diffeomorphic=True,
full_size=True,
indexing='ij',
init_mu=None,
init_sigma=None,
stat_post_warp=False, # compute statistics post warp?
network_stat_weight=0.001,
warp_method='WARP',
stat_nb_feats=16):
"""
Network to do unsupervised segmentation with probabilistic atlas
(Dalca et al., submitted to MICCAI 2019)
"""
# print(warp_method)
ndims = len(vol_size)
assert ndims in [1, 2, 3], "ndims should be one of 1, 2, or 3. found: %d" % ndims
weaknorm = RandomNormal(mean=0.0, stddev=1e-5)
# get the core model
unet_model = unet_core(vol_size, enc_nf, dec_nf, full_size=full_size, tgt_feats=nb_labels)
[src_img, src_atl] = unet_model.inputs
x = unet_model.output
# transform the results into a flow field.
Conv = getattr(KL, 'Conv%dD' % ndims)
flow1 = Conv(ndims, kernel_size=3, padding='same', name='flow', kernel_initializer=weaknorm)(x)
if diffeomorphic:
flow2 = nrn_layers.VecInt(method='ss', name='flow-int', int_steps=8)(flow1)
else:
flow2 = flow1
if full_size:
flow = flow2
else:
flow = trf_resize(flow2, 1/2, name='diffflow')
# warp atlas
if warp_method == 'WARP':
warped_atlas = nrn_layers.SpatialTransformer(interp_method='linear', indexing=indexing, name='warped_atlas')([src_atl, flow])
else:
warped_atlas = src_atl
if stat_post_warp:
assert warp_method == 'WARP', "if computing stat post warp, must do warp... :) set warp_method to 'WARP' or stat_post_warp to False?"
# combine warped atlas and warpedimage and output mu and log_sigma_squared
combined = concatenate([warped_atlas, src_img])
else:
combined = unet_model.layers[-2].output
conv1 = conv_block(combined, stat_nb_feats)
conv2 = conv_block(conv1, nb_labels)
stat_mu_vol = Conv(nb_labels, kernel_size=3, name='mu_vol',
kernel_initializer=weaknorm, bias_initializer=weaknorm)(conv2)
stat_mu = keras.layers.GlobalMaxPooling3D()(stat_mu_vol)
stat_logssq_vol = Conv(nb_labels, kernel_size=3, name='logsigmasq_vol',
kernel_initializer=weaknorm, bias_initializer=weaknorm)(conv2)
stat_logssq = keras.layers.GlobalMaxPooling3D()(stat_logssq_vol)
# combine mu with initializtion
if init_mu is not None:
init_mu = np.array(init_mu)
stat_mu = Lambda(lambda x: network_stat_weight * x + init_mu, name='comb_mu')(stat_mu)
# combine sigma with initializtion
if init_sigma is not None:
init_logsigmasq = np.array([2*np.log(f) for f in init_sigma])
stat_logssq = Lambda(lambda x: network_stat_weight * x + init_logsigmasq, name='comb_sigma')(stat_logssq)
# unnorm log-lik
def unnorm_loglike(I, mu, logsigmasq, uselog=True):
P = tf.distributions.Normal(mu, K.exp(logsigmasq/2))
if uselog:
return P.log_prob(I)
else:
return P.prob(I)
uloglhood = KL.Lambda(lambda x:unnorm_loglike(*x), name='unsup_likelihood')([src_img, stat_mu, stat_logssq])
# compute data loss as a layer, because it's a bit easier than outputting a ton of things, etc.
# def logsum(ll, atl):
# pdf = ll * atl
# return tf.log(tf.reduce_sum(pdf, -1, keepdims=True) + K.epsilon())
def logsum_safe(prob_ll, atl):
"""
safe computation using the log sum exp trick
e.g. https://www.xarg.org/2016/06/the-log-sum-exp-trick-in-machine-learning/
where x = logpdf
note does not normalize p
"""
logpdf = prob_ll + K.log(atl + K.epsilon())
alpha = tf.reduce_max(logpdf, -1, keepdims=True)
return alpha + tf.log(tf.reduce_sum(K.exp(logpdf-alpha), -1, keepdims=True) + K.epsilon())
loss_vol = Lambda(lambda x: logsum_safe(*x))([uloglhood, warped_atlas])
return Model(inputs=[src_img, src_atl], outputs=[loss_vol, flow])
def miccai2018_net(vol_size,
enc_nf,
dec_nf,
int_steps=7,
use_miccai_int=False,
indexing='ij',
bidir=False,
vel_resize=1 / 2):
ndims = len(vol_size)
assert ndims in [1, 2,
3], "ndims should be one of 1, 2, or 3. found: %d" % ndims
# get unet
unet_model = unet_core(vol_size, enc_nf, dec_nf, full_size=False)
[src, tgt] = unet_model.inputs
x_out = unet_model.outputs[-1]
# velocity mean and logsigma layers
Conv = getattr(KL, 'Conv%dD' % ndims)
flow_mean = Conv(ndims,
kernel_size=3,
padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5),
name='flow')(x_out)
# we're going to initialize the velocity variance very low, to start stable.
flow_log_sigma = Conv(
ndims,
kernel_size=3,
padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-10),
bias_initializer=keras.initializers.Constant(value=-10),
name='log_sigma')(x_out)
flow_params = concatenate([flow_mean, flow_log_sigma])
# velocity sample
flow = Sample(name="z_sample")([flow_mean, flow_log_sigma])
# integrate if diffeomorphic (i.e. treating 'flow' above as stationary velocity field)
if use_miccai_int:
# for the miccai2018 submission, the squaring layer
# scaling was essentially built in by the network
# was manually composed of a Transform and and Add Layer.
v = flow
for _ in range(int_steps):
v1 = nrn_layers.SpatialTransformer(interp_method='linear',
indexing=indexing)([v, v])
v = keras.layers.add([v, v1])
flow = v
else:
# new implementation in neuron is cleaner.
z_sample = flow
flow = nrn_layers.VecInt(method='ss',
name='flow-int',
int_steps=int_steps)(z_sample)
if bidir:
rev_z_sample = Negate()(z_sample)
neg_flow = nrn_layers.VecInt(method='ss',
name='neg_flow-int',
int_steps=int_steps)(rev_z_sample)
# get up to final resolution
flow = trf_resize(flow, vel_resize, name='diffflow')
if bidir:
neg_flow = trf_resize(neg_flow, vel_resize, name='neg_diffflow')
# transform
y = nrn_layers.SpatialTransformer(interp_method='linear',
indexing=indexing)([src, flow])
if bidir:
y_tgt = nrn_layers.SpatialTransformer(
interp_method='linear', indexing=indexing)([tgt, neg_flow])
# prepare outputs and losses
outputs = [y, flow_params]
if bidir:
outputs = [y, y_tgt, flow_params]
# build the model
return Model(inputs=[src, tgt], outputs=outputs)
def labels_to_image_model(labels_shape,
crop_shape,
generation_label_list,
segmentation_label_list,
n_channels=1,
labels_res=(1,1,1),
target_res=None,
padding_margin=None,
apply_affine_trans=False,
apply_nonlin_trans=True,
nonlin_shape_factor=0.0625,
apply_bias_field=True,
bias_shape_factor=0.025,
blur_background=True,
normalise=True,
out_div_32=False,
convert_back=False,
id=0, # For different layer names if several models.
rand_blur=True):
"""
This function builds a keras/tensorflow model to generate brain images from supplied labels.
It returns the model as well as the shape ouf the output images without batch and channel dimensions
(height*width*depth).
The model takes as inputs:
-a label image
-a vector containing the means of the Gaussian distributions to sample for each label,
-a similar vector for the associated standard deviations.
-if apply_affine_deformation=True: a (n_dims+1)x(n_dims+1) affine matrix
-if apply_non_linear_deformation=True: a small non linear field of size batch*x*y*z*n_dims that will be
resampled to labels size
-if apply_bias_field=True: a small bias field of size batch*x*y*z*1 that will be resampled to labels size
The model returns:
-the generated image
-the corresponding label map
:param labels_shape: should be a list or tensor with image dimension plus channel size at the end
:param n_channels: number of channels to be synthetised
:param labels_res: list of dimension resolutions of model's inputs
:param target_res: list of dimension resolutions of model's outputs
:param crop_shape: list, shape of model's outputs
:param generation_label_list: list of all the labels in the dataset (internally converted to [0...N-1] and converted
back to original values at the end of model)
:param segmentation_label_list: list of all the labels in the output labels (internally converted to [0...N-1] and
converted back to original values at the end of model)
:param padding_margin: margin by which the input labels will be 0-padded. This step happens
before an eventual cropping. Default is None, no padding.
:param apply_affine_trans: whether to apply affine deformation during generation
:param apply_nonlin_trans: whether to apply non linear deformation during generation
:param nonlin_shape_factor: if apply_non_linear_deformation=True, factor between the shapes of the labels and of
the non-linear field that will be sampled
:param apply_bias_field: whether to apply a bias field to the created image during generation
:param bias_shape_factor: if apply_bias_field=True, factor between the shapes of the labels and of the bias field
that will be sampled
:param blur_background: Whether background is a regular label, thus blurred with the others.
:param normalise: whether to normalise data. Default is False.
:param out_div_32: whether model's outputs must be of shape divisible by 32
"""
# get shapes
n_dims = len(labels_shape)
target_res = format_target_res(target_res, n_dims)
crop_shape, resample_shape, output_shape, padding_margin = get_shapes(labels_shape,
crop_shape,
labels_res,
target_res,
padding_margin,
out_div_32)
# create new_label_list and corresponding LUT to make sure that labels go from 0 to N-1
n_generation_labels = generation_label_list.shape[0]
new_generation_label_list = np.arange(n_generation_labels)
lut = np.zeros(np.max(generation_label_list).astype('int') + 1)
for n in range(n_generation_labels):
lut[generation_label_list[n].astype('int')] = n
# define mandatory inputs
labels_input = KL.Input(shape=(*labels_shape, 1), name=f'labels_input_{id}')
means_input = KL.Input(shape=(*new_generation_label_list.shape, n_channels), name=f'means_input_{id}')
std_devs_input = KL.Input(shape=(*new_generation_label_list.shape, n_channels), name=f'std_devs_input_{id}')
list_inputs = [labels_input, means_input, std_devs_input]
# convert labels to new_label_list
labels = KL.Lambda(lambda x: tf.gather(tf.convert_to_tensor(lut, dtype='int32'),
tf.cast(x, dtype='int32')), name=f'convert_labels_{id}')(labels_input)
# pad labels
if padding_margin is not None:
pad = np.transpose(np.array([[0] + padding_margin + [0]] * 2))
labels = KL.Lambda(lambda x: tf.pad(x, tf.cast(tf.convert_to_tensor(pad), dtype='int32')), name=f'pad_{id}')(labels)
labels_shape = labels.get_shape().as_list()[1:-1]
# cropping
if crop_shape is not None:
# get maximum cropping indices in each dimension
cropping_max_val = [labels_shape[i] - crop_shape[i] for i in range(n_dims)]
# prepare cropping indices and tensor's new shape
idx = KL.Lambda(lambda x: tf.zeros([1], dtype='int32'), name=f'no_cropping_batch_{id}')(means_input) # no cropping
for val_idx, val in enumerate(cropping_max_val): # draw cropping indices for image dimensions
if val > 0:
idx = KL.Lambda(lambda x: tf.concat(
[tf.cast(x, dtype='int32'), K.random_uniform([1], minval=0, maxval=val, dtype='int32')], axis=0),
name=f'pick_cropping_idx_{val_idx}_{id}')(idx)
else:
idx = KL.Lambda(lambda x: tf.concat([tf.cast(x, dtype='int32'), tf.zeros([1], dtype='int32')], axis=0),
name=f'pick_cropping_idx_{val_idx}_{id}')(idx)
idx = KL.Lambda(lambda x: tf.concat([tf.cast(x, dtype='int32'), tf.zeros([1], dtype='int32')], axis=0),
name=f'no_cropping_channel_{id}')(idx) # no cropping for channel dimension
patch_shape_tens = KL.Lambda(lambda x: tf.convert_to_tensor([-1] + crop_shape + [-1], dtype='int32'),
name=f'tensor_cropping_idx_{id}')(means_input)
# perform cropping
labels = KL.Lambda(
lambda x: tf.slice(x[0], begin=tf.cast(x[1], dtype='int32'), size=tf.cast(x[2], dtype='int32')),
name=f'cropping_{id}')([labels, idx, patch_shape_tens])
else:
crop_shape = labels_shape
labels = KL.Lambda(lambda x: tf.cast(x, dtype='float'))(labels)
# if necessary, resample image and labels at target resolution
if resample_shape is not None:
labels = KL.Lambda(lambda x: tf.cast(x, dtype='float32'))(labels)
zoom_fact = [r / l for r, l in zip(resample_shape, labels_shape)]
labels = nrn_layers.Resize(zoom_fact, interp_method='nearest', name=f'resample_labels_{id}')(labels)
# deform labels
if apply_affine_trans | apply_nonlin_trans:
labels._keras_shape = tuple(labels.get_shape().as_list())
trans_inputs = [labels]
# add affine deformation to inputs list
if apply_affine_trans:
aff_in = KL.Input(shape=(n_dims + 1, n_dims + 1), name=f'aff_input_{id}')
list_inputs.append(aff_in)
trans_inputs.append(aff_in)
# prepare non-linear deformation field and add it to inputs list
if apply_nonlin_trans:
def_field_size = get_nonlin_field_shape(crop_shape, nonlin_shape_factor)
nonlin_field_in = KL.Input(shape=def_field_size, name=f'nonlin_input_{id}')
list_inputs.append(nonlin_field_in)
int_at = 2.0
zoom = [o / d / int_at for o, d in zip(output_shape, def_field_size)]
vel_field = nonlin_field_in
vel_field = nrn_layers.Resize(zoom, interp_method='linear', name=f'resize_vel_{id}')(vel_field)
def_field = nrn_layers.VecInt(int_steps=5)(vel_field)
#def_field = nrn_layers.RescaleValues(int_at)(def_field)
def_field = nrn_layers.Resize(int_at, interp_method='linear', name=f'resize_def_{id}')(def_field)
trans_inputs.append(def_field)
# apply deformations
labels = nrn_layers.SpatialTransformer(interp_method='nearest', name=f'trans_{id}')(trans_inputs)
labels = KL.Lambda(lambda x: tf.cast(x, dtype='int32'))(labels)
# sample from normal distribution
image = KL.Lambda(lambda x: tf.random.normal(tf.shape(x)), name=f'sample_normal_{id}')(labels)
# build synthetic image
f_cat = lambda x: tf.concat([x+n_generation_labels*i for i in range(n_channels)], -1)
cat_labels = KL.Lambda(f_cat, name=f'cat_labels_{id}')(labels)
f_gather = lambda x: tf.gather(tf.reshape(x[0], [-1]), tf.cast(x[1], dtype='int32'))
f_map = lambda x: tf.map_fn(f_gather, x, dtype='float32')
means = KL.Lambda(f_map)([means_input, cat_labels])
std_devs = KL.Lambda(f_map)([std_devs_input, cat_labels])
image = KL.Multiply(name=f'mul_std_dev_{id}')([image, std_devs])
image = KL.Add(name=f'add_means_{id}')([image, means])
if rand_blur:
shape = [5] * n_dims
lim = [(s - 1) / 2 for s in shape]
lim = [np.arange(-l, l+1) for l in lim]
grid = np.meshgrid(*lim, indexing='ij')
grid = [g ** 2 for g in grid]
c_grid = KL.Lambda(lambda x: tf.constant(np.stack(grid), dtype='float32'))([])
sigma = KL.Lambda(lambda x: tf.random.uniform((n_dims,), minval=1e-6, maxval=1))([])
f = lambda x: x[0] / x[1]**2
kernel = KL.Lambda(lambda x: tf.map_fn(f, x, dtype='float32'))([c_grid, sigma])
kernel = KL.Lambda(lambda x: tf.exp( -tf.reduce_sum(x, axis=0) ))(kernel)
kernel = KL.Lambda(lambda x: x[..., None, None] / tf.reduce_sum(x))(kernel)
else:
if (target_res is None) | (labels_res == target_res):
sigma = [0.55] * n_dims
else:
sigma = [0.85 * labels_res[i] / target_res[i] for i in range(n_dims)]
kernel = KL.Lambda(lambda x: tf.convert_to_tensor(add_axis(add_axis(gauss_kernel(sigma, n_dims), -1), -1),
dtype=x.dtype), name=f'gauss_kernel_{id}')(image)
if n_channels == 1:
image = KL.Lambda(lambda x: tf.nn.convolution(x[0], x[1], padding='SAME', strides=[1] * n_dims),
name=f'blur_image_{id}')([image, kernel])
mask = KL.Lambda(lambda x: tf.where(tf.greater(x, 0), tf.ones_like(x, dtype='float32'),
tf.zeros_like(x, dtype='float32')), name=f'masking_{id}')(labels)
if not blur_background:
blurred_mask = KL.Lambda(lambda x: tf.nn.convolution(x[0], x[1], padding='SAME', strides=[1] * n_dims),
name=f'blur_mask_{id}')([mask, kernel])
image = KL.Lambda(lambda x: x[0] / (x[1] + K.epsilon()), name=f'masked_blurring_{id}')([image, blurred_mask])
bckgd_mean = KL.Lambda(lambda x: tf.random.uniform((1, 1), 0, 10), name=f'bckgd_mean_{id}')([])
bckgd_std = KL.Lambda(lambda x: tf.random.uniform((1, 1), 0, 5), name=f'bckgd_std_{id}')([])
rand_flip = KL.Lambda(lambda x: K.greater(tf.random.uniform((1, 1), 0, 1), 0.5), name=f'bool_{id}')([])
bckgd_mean = KL.Lambda(lambda y: K.switch(y[0],
KL.Lambda(lambda x: tf.zeros_like(x))(y[1]),
y[1]), name=f'switch_backgd_mean_{id}')([rand_flip, bckgd_mean])
bckgd_std = KL.Lambda(lambda y: K.switch(y[0],
KL.Lambda(lambda x: tf.zeros_like(x))(y[1]),
y[1]), name=f'switch_backgd_std_{id}')([rand_flip, bckgd_std])
background = KL.Lambda(lambda x: x[1] + x[2] * tf.random.normal(tf.shape(x[0])),
name=f'gaussian_bckgd_{id}')([image, bckgd_mean, bckgd_std])
image = KL.Lambda(lambda x: tf.where(tf.cast(x[1], dtype='bool'), x[0], x[2]),
name=f'mask_blurred_image_{id}')([image, mask, background])
else:
rand_flip = KL.Lambda(lambda x: K.greater(tf.random.uniform((1, 1), 0, 1), 0.8), name=f'bool_{id}')([])
image = KL.Lambda(lambda y: K.switch(y[0], KL.Lambda(
lambda x: tf.where(tf.cast(x[1], dtype='bool'), x[0], tf.zeros_like(x[0])), name=f'mask_image_{id}')(
[y[1], y[2]]), y[1]), name=f'switch_backgd_reset_{id}')([rand_flip, image, mask])
else:
# blur each image channel separately
split = KL.Lambda(lambda x: tf.split(x, [1]*n_channels, axis=-1))(image)
image = KL.Lambda(lambda x: tf.nn.convolution(x[0], x[1], padding='SAME', strides=[1] * n_dims),
name=f'blurring_0_{id}')([split[0], kernel])
for i in range(1, n_channels):
temp_blurred = KL.Lambda(lambda x: tf.nn.convolution(x[0], x[1], padding='SAME', strides=[1] * n_dims),
name=f'blurring_{i}_{id}')([split[i], kernel])
mask = KL.Lambda(lambda x: tf.where(tf.greater(x, 0), tf.ones_like(x, dtype='float32'),
tf.zeros_like(x, dtype='float32')), name=f'masking_{i}_{id}')(labels)
if not blur_background:
blurred_mask = KL.Lambda(lambda x: tf.nn.convolution(x[0], x[1], padding='SAME', strides=[1] * n_dims),
name=f'blur_mask_{i}_{id}')([mask, kernel])
temp_blurred = KL.Lambda(lambda x: x[0] / (x[1]+K.epsilon()),
name=f'masked_blurring_{i}_{id}')([temp_blurred, blurred_mask])
bckgd_mean = KL.Lambda(lambda x: tf.random.uniform((1, 1), 0, 10), name=f'bckgd_mean_{i}_{id}')([])
bckgd_std = KL.Lambda(lambda x: tf.random.uniform((1, 1), 0, 5), name=f'bckgd_std_{i}_{id}')([])
rand_flip = KL.Lambda(lambda x: K.greater(tf.random.uniform((1, 1), 0, 1), 0.5), name=f'bool{i}_{id}')([])
bckgd_mean = KL.Lambda(lambda y: K.switch(y[0],
KL.Lambda(lambda x: tf.zeros_like(x, dtype='float32'))(y[1]),
y[1]), name=f'switch_backgd_mean{i}_{id}')([rand_flip, bckgd_mean])
bckgd_std = KL.Lambda(lambda y: K.switch(y[0],
KL.Lambda(lambda x: tf.zeros_like(x, dtype='float32'))(y[1]),
y[1]), name=f'switch_backgd_std_{i}_{id}')([rand_flip, bckgd_std])
background = KL.Lambda(lambda x: x[1] + x[2] * tf.random.normal(tf.shape(x[0])),
name=f'gaussian_bckgd_{i}_{id}')([temp_blurred, bckgd_mean, bckgd_std])
temp_blurred = KL.Lambda(lambda x: tf.where(tf.cast(x[1], dtype='bool'), x[0], x[2]),
name=f'mask_blurred_image_{i}_{id}')([temp_blurred, mask, background])
else:
rand_flip = KL.Lambda(lambda x: K.greater(tf.random.uniform((1, 1), 0, 1), 0.8), name=f'boo{i}_{id}')([])
image = KL.Lambda(lambda y: K.switch(y[0], KL.Lambda(
lambda x: tf.where(tf.cast(x[1], dtype='bool'), x[0], tf.zeros_like(x[0])), name=f'mask_image_{i}_{id}')(
[y[1], y[2]]), y[1]), name=f'switch_backgd_reset_{i}_{id}')([rand_flip, image, mask])
image = KL.Lambda(lambda x: tf.concat([x[0], x[1]], -1),
name=f'cat_blurring_{i}_{id}')([image, temp_blurred])
# apply bias field
if apply_bias_field:
# format bias field and add it to inputs list
bias_field_size = get_bias_field_shape(output_shape, bias_shape_factor)
bias_field_in = KL.Input(shape=bias_field_size, name=f'bias_input_{id}')
list_inputs.append(bias_field_in)
# resize bias field and apply it to image
zoom_fact = [o / d for o, d in zip(output_shape, bias_field_size)]
bias_field = nrn_layers.Resize(zoom_fact, interp_method='linear', name=f'log_bias_{id}')(bias_field_in)
bias_field = KL.Lambda(lambda x: K.exp(x), name=f'bias_field_{id}')(bias_field)
image._keras_shape = tuple(image.get_shape().as_list())
bias_field._keras_shape = tuple(bias_field.get_shape().as_list())
image = KL.multiply([bias_field, image])
# make sure image's intensities are between 0 and 255
image = KL.Lambda(lambda x: K.clip(x, 0, 255), name=f'clipping_{id}')(image)
# contrast stretching
image = KL.Lambda(
lambda x: x * tf.random.uniform([1], minval=0.6, maxval=1.4) + tf.random.uniform([1], minval=-30, maxval=30),
name=f'stretching_{id}')(image)
# convert labels back to original values and remove unwanted labels
if convert_back:
out_lut = [x if x in segmentation_label_list else 0 for x in generation_label_list]
else:
# Rebase wanted indices into [0, N-1] for one-hot encoding.
n = 0
out_lut = [None] * len(generation_label_list)
for i, x in enumerate(generation_label_list):
out = -1
if x in segmentation_label_list:
out = n
n += 1
out_lut[i] = out
labels = KL.Lambda(lambda x: tf.gather(tf.cast(out_lut, dtype='int32'),
tf.cast(x, dtype='int32')), name=f'labels_back_{id}')(labels)
# normalise the produced image (include labels_out, so this layer is not removed when plugging in other keras model)
if normalise:
m = KL.Lambda(lambda x: K.min(x), name=f'min_{id}')(image)
M = KL.Lambda(lambda x: K.max(x), name=f'max_{id}')(image)
image = KL.Lambda(lambda x: (x[0]-x[1])/(x[2]-x[1]), name=f'normalisation_{id}')([image, m, M])
else:
image = KL.Lambda(lambda x: x[0] + K.zeros(1), name=f'dummy_{id}')([image])
# gamma augmentation
image = KL.Lambda(lambda x: tf.math.pow(x[0], tf.math.exp(tf.random.normal([1], mean=0, stddev=0.25))),
name=f'gamma_{id}')([image, labels])
outputs = [image, labels]
if apply_nonlin_trans:
outputs.append(vel_field)
brain_model = keras.Model(inputs=list_inputs, outputs=outputs)
return brain_model, def_field_size, bias_field_size
def diff_net(vol_size,
enc_nf,
dec_nf,
int_steps=7,
src_feats=1,
indexing='ij',
bidir=False,
ret_flows=False,
full_size=False,
vel_resize=1 / 2,
src=None,
tgt=None):
"""
diffeomorphic net, similar to miccai2018, but no sampling.
architecture for probabilistic diffeomoprhic VoxelMorph presented in the MICCAI 2018 paper.
You may need to modify this code (e.g., number of layers) to suit your project needs.
The stationary velocity field operates in a space (0.5)^3 of vol_size for computational reasons.
:param vol_size: volume size. e.g. (256, 256, 256)
:param enc_nf: list of encoder filters. right now it needs to be 1x4.
e.g. [16,32,32,32]
:param dec_nf: list of decoder filters. right now it must be 1x6, see unet function.
:param use_miccai_int: whether to use the manual miccai implementation of scaling and squaring integration
note that the 'velocity' field outputted in that case was
since then we've updated the code to be part of a flexible layer. see neuron.layers.VecInt
**This param will be phased out (set to False behavior)**
:param int_steps: the number of integration steps
:param indexing: xy or ij indexing. we recommend ij indexing if training from scratch.
miccai 2018 runs were done with xy indexing.
**This param will be phased out (set to 'ij' behavior)**
:return: the keras model
"""
ndims = len(vol_size)
assert ndims in [1, 2,
3], "ndims should be one of 1, 2, or 3. found: %d" % ndims
# get unet
unet_model = unet_core(vol_size,
enc_nf,
dec_nf,
full_size=full_size,
src=src,
tgt=tgt,
src_feats=src_feats)
[src, tgt] = unet_model.inputs
# velocity sample
# unet_model.layers[-1].name = 'vel'
# vel = unet_model.output
x_out = unet_model.outputs[-1]
# velocity mean and logsigma layers
Conv = getattr(KL, 'Conv%dD' % ndims)
vel = Conv(ndims,
kernel_size=3,
padding='same',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-5),
name='flow')(x_out)
if full_size and vel_resize != 1:
vel = trf_resize(vel, 1.0 / vel_resize, name='flow-resize')
# new implementation in neuron is cleaner.
flow = nrn_layers.VecInt(method='ss', name='flow-int',
int_steps=int_steps)(vel)
if bidir:
# rev_z_sample = Lambda(lambda x: -x)(z_sample)
neg_vel = Negate()(vel)
neg_flow = nrn_layers.VecInt(method='ss',
name='neg_flow-int',
int_steps=int_steps)(neg_vel)
# get up to final resolution
flow = trf_resize(flow, vel_resize, name='diffflow')
if bidir:
neg_flow = trf_resize(neg_flow, vel_resize, name='neg_diffflow')
# transform
y = nrn_layers.SpatialTransformer(interp_method='linear',
indexing=indexing,
name='warped_src')([src, flow])
if bidir:
y_tgt = nrn_layers.SpatialTransformer(interp_method='linear',
indexing=indexing,
name='warped_tgt')(
[tgt, neg_flow])
# prepare outputs and losses
outputs = [y, vel]
if bidir:
outputs = [y, y_tgt, vel]
model = Model(inputs=[src, tgt], outputs=outputs)
if ret_flows:
outputs += [
model.get_layer('diffflow').output,
model.get_layer('neg_diffflow').output
]
return Model(inputs=[src, tgt], outputs=outputs)
else:
return model
