def randflow_ronneberger_model(img_shape,
model,
model_name='randflow_ronneberger_model',
flow_sigma=None,
flow_amp=None,
blur_sigma=5,
interp_mode='linear',
indexing='xy',
):
n_dims = len(img_shape) - 1
x_in = Input(img_shape, name='img_input_randwarp')
if n_dims == 3:
n_pools = 5
flow = MaxPooling3D(2)(x_in)
for i in range(n_pools-1):
flow = MaxPooling3D(2)(flow)
# reduce flow by a factor of 64 until we have roughly 3x3x3
flow_shape = tuple([int(s/(2**n_pools)) for s in img_shape[:-1]] + [n_dims])
print('Smallest flow shape: {}'.format(flow_shape))
else:
#flow = flow_placeholder
flow = x_in
flow_shape = img_shape[:-1] + (n_dims,)
# random flow field
if flow_amp is None:
# sigmas and blurring are hand-tuned to be similar to gaussian with stddev = 10, with smooth upsampling
flow = RandFlow(name='randflow', img_shape=flow_shape, blur_sigma=0., flow_sigma=flow_sigma * 8)(flow)
if n_dims == 3:
print(flow_shape)
print(flow.get_shape())
flow = Reshape(flow_shape)(flow)
flow_shape = flow_shape[:-1]
for i in range(n_pools):
flow_shape = [fs * 2 for fs in flow_shape]
flow = Lambda(interp_upsampling, output_shape=tuple(flow_shape) + (n_dims,))(flow)
if i > 0 and i < 4:
print(flow_shape)
flow = BlurFlow(img_shape=tuple(flow_shape) + (n_dims,), blur_sigma=5,
)(flow)#min(7, flow_shape[0]/4.))(flow)
flow = basic_networks._pad_or_crop_to_shape(flow, flow_shape, img_shape)
flow = BlurFlow(img_shape[:-1] + (n_dims,), blur_sigma=3)(flow)
flow = Reshape(img_shape[:-1] + (n_dims,), name='randflow_out')(flow)
else:
flow = Reshape(img_shape[:-1] + (n_dims,), name='randflow_out')(flow)
x_warped = SpatialTransformer(
indexing=indexing, interp_method=interp_mode, name='densespatialtransformer_img')([x_in, flow])
if model is not None:
model_outputs = model(x_warped)
if not isinstance(model_outputs, list):
model_outputs = [model_outputs]
else:
model_outputs = [x_warped, flow]
return Model(inputs=[x_in], outputs=model_outputs, name=model_name)
def generator_net(vol_size, enc_nf, dec_nf, vel_resize, ti_flow, int_steps):
ndims = len(vol_size)
assert ndims in [
1, 2, 3
], "ndims should be one of 1, 2, or 3. found: {}".format(ndims)
unet = unet_core(vol_size, enc_nf, dec_nf, vel_resize, ti_flow)
# target delta in binary representation
b_in = Input(shape=(16, ))
x_in = unet.inputs[0]
x_out, x_ti = unet.outputs
inputs = [x_in, b_in]
Conv = getattr(KL, 'Conv{}D'.format(ndims))
# time dependent flow component (i.e. aging)
flow_mean = Conv(ndims,
kernel_size=3,
padding='same',
name='flow_mean',
kernel_initializer=RandomNormal(mean=0.0,
stddev=1e-5))(x_out)
flow_log_sigma = Conv(
ndims,
kernel_size=3,
padding='same',
name='flow_log_sigma',
kernel_initializer=RandomNormal(mean=0.0, stddev=1e-10),
bias_initializer=keras.initializers.Constant(value=-10))(x_out)
flow_params = concatenate([flow_mean, flow_log_sigma])
# sample velocity field (using reparameterization)
z_sample = Sample(name='z_sample')([flow_mean, flow_log_sigma])
# integrate flow
flow = VecInt(method='ss', name='flow_int',
int_steps=int_steps)([z_sample, b_in])
# resize to full size
if vel_resize != 1.0:
flow = trf_resize(flow, vel_resize, name='flow_resize')
# time independent flow component
flow_ti = Conv(ndims, kernel_size=3, padding='same', name='flow_ti')(x_ti)
flow_ti = trf_resize(flow_ti, 0.25, name='flow_ti_resize')
if ti_flow:
flow = add([flow, flow_ti])
y = SpatialTransformer(interp_method='linear', indexing='ij')([x_in, flow])
outputs = [y, flow_params, flow_ti]
return Model(inputs=inputs, outputs=outputs)
def warp_model(img_shape, interp_mode='linear', indexing='ij'):
n_dims = len(img_shape) - 1
img_in = Input(img_shape, name='input_img')
flow_in = Input(img_shape[:-1] + (n_dims,), name='input_flow')
img_warped = SpatialTransformer(
interp_mode, indexing=indexing, name='densespatialtransformer_img')([img_in, flow_in])
return Model(inputs=[img_in, flow_in], outputs=img_warped, name='warp_model')
def color_delta_unet_model(img_shape,
n_output_chans,
model_name='color_delta_unet',
enc_params=None,
include_aux_input=False,
aux_input_shape=None,
do_warp_to_target_space=False
):
x_src = Input(img_shape, name='input_src')
x_tgt = Input(img_shape, name='input_tgt')
if aux_input_shape is None:
aux_input_shape = img_shape
x_seg = Input(aux_input_shape, name='input_src_aux')
inputs = [x_src, x_tgt, x_seg]
if do_warp_to_target_space: # warp transformed vol to target space in the end
n_dims = len(img_shape) - 1
flow_srctotgt = Input(img_shape[:-1] + (n_dims,), name='input_flow')
inputs += [flow_srctotgt]
if include_aux_input:
unet_inputs = [x_src, x_tgt, x_seg]
unet_input_shape = img_shape[:-1] + (img_shape[-1] * 2 + aux_input_shape[-1],)
else:
unet_inputs = [x_src, x_tgt]
unet_input_shape = img_shape[:-1] + (img_shape[-1] * 2,)
x_stacked = Concatenate(axis=-1)(unet_inputs)
n_dims = len(img_shape) - 1
if n_dims == 2:
color_delta = basic_networks.unet2D(x_stacked, unet_input_shape, n_output_chans,
nf_enc=enc_params['nf_enc'],
nf_dec=enc_params['nf_dec'],
n_convs_per_stage=enc_params['n_convs_per_stage'],
include_residual=False)
conv_fn = Conv2D
else:
color_delta = basic_networks.unet3D(x_stacked, unet_input_shape, n_output_chans,
nf_enc=enc_params['nf_enc'],
nf_dec=enc_params['nf_dec'],
n_convs_per_stage=enc_params['n_convs_per_stage'],
include_residual=False)
conv_fn = Conv3D
# last conv to get the output shape that we want
color_delta = conv_fn(n_output_chans, kernel_size=3, padding='same', name='color_delta')(color_delta)
transformed_out = Add(name='add_color_delta')([x_src, color_delta])
if do_warp_to_target_space:
transformed_out = SpatialTransformer(indexing='xy')([transformed_out, flow_srctotgt])
# kind of silly, but do a reshape so keras doesnt complain about returning an input
x_seg = Reshape(aux_input_shape, name='aux')(x_seg)
return Model(inputs=inputs, outputs=[color_delta, transformed_out, x_seg], name=model_name)
def perform_test_affine_transform(arr):
"""
perform an identity affine transformation for testing
:param arr: image array
:return:
"""
tx = create_identity_transform_stn()
ST = SpatialTransformer(interp_method='linear', indexing='ij')
x = ST([tf.convert_to_tensor(arr), tf.convert_to_tensor(tx)])
plot_side2side(arr, x, 'before after')
return 1
def build_morph(self, enc_nf, dec_nf):
# 实例化morph main body
morph = self.morph_main_body(enc_nf, dec_nf)
# 指定输入
moving_images, fixed_images = morph.input
# 实例化空间变化 并且指定输出
moved_images = SpatialTransformer(interp_method='linear',
indexing='ij',
name='moved_images')(
[moving_images, morph.output])
# 返回模型
return Model([moving_images, fixed_images],
[moved_images, morph.output])
def interp_upsampling(V):
"""
upsample a field by a factor of 2
TODO: should switch this to use neuron.utils.interpn()
"""
V = tf.reshape(V, [-1] + V.get_shape().as_list()[1:])
grid = 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 = SpatialTransformer(interp_method='linear')([V, offset])
return V
def create_morph(self, enc_nf, dec_nf):
# 输入层
moving_images = Input(shape=[*self.vol_shape, self.channels],
name='moving_images')
fixed_images = Input(shape=[*self.vol_shape, self.channels],
name='fixed_images')
# 链接
x_in = concatenate([moving_images, fixed_images])
# 卷积实现下采样
x_enc = [x_in]
for i in range(len(enc_nf)):
if i == 0:
x_enc.append(self.conv_block(x_enc[-1], enc_nf[i]))
else:
x_enc.append(self.conv_block(x_enc[-1], enc_nf[i], 2))
# 上采样
x = UpSampling2D()(x_enc[-1])
x = concatenate([x, x_enc[-2]])
x = self.conv_block(x, dec_nf[0])
x = UpSampling2D()(x)
x = concatenate([x, x_enc[-3]])
x = self.conv_block(x, dec_nf[1])
x = UpSampling2D()(x)
x = concatenate([x, x_enc[-4]])
x = self.conv_block(x, dec_nf[2])
x = self.conv_block(x, dec_nf[3])
x = UpSampling2D()(x)
x = concatenate([x, x_enc[-5]])
x = self.conv_block(x, dec_nf[4])
# 特征形成形变场
Fai = Conv2D(filters=2,
kernel_size=3,
padding='same',
kernel_initializer='he_normal',
strides=1,
name='Fai')(x)
# 恢复
moved_images = SpatialTransformer(interp_method='linear',
indexing='ij',
name='moved_images')(
[moving_images, Fai])
return Model([moving_images, fixed_images], [moved_images, Fai])
def randflow_model(img_shape,
model,
model_name='randflow_model',
flow_sigma=None,
flow_amp=None,
blur_sigma=5,
interp_mode='linear',
indexing='xy',
):
n_dims = len(img_shape) - 1
x_in = Input(img_shape, name='img_input_randwarp')
if n_dims == 3:
flow = MaxPooling3D(2)(x_in)
flow = MaxPooling3D(2)(flow)
blur_sigma = int(np.ceil(blur_sigma / 4.))
flow_shape = tuple([int(s/4) for s in img_shape[:-1]] + [n_dims])
else:
flow = x_in
flow_shape = img_shape[:-1] + (n_dims,)
# random flow field
if flow_amp is None:
flow = RandFlow(name='randflow', img_shape=flow_shape, blur_sigma=blur_sigma, flow_sigma=flow_sigma)(flow)
elif flow_sigma is None:
flow = RandFlow_Uniform(name='randflow', img_shape=flow_shape, blur_sigma=blur_sigma, flow_amp=flow_amp)(flow)
if n_dims == 3:
flow = Reshape(flow_shape)(flow)
# upsample with linear interpolation
flow = Lambda(interp_upsampling)(flow)
flow = Lambda(interp_upsampling, output_shape=img_shape[:-1] + (n_dims,))(flow)
flow = Reshape(img_shape[:-1] + (n_dims,), name='randflow_out')(flow)
else:
flow = Reshape(img_shape[:-1] + (n_dims,), name='randflow_out')(flow)
x_warped = SpatialTransformer(interp_method=interp_mode, name='densespatialtransformer_img', indexing=indexing)(
[x_in, flow])
if model is not None:
model_outputs = model(x_warped)
if not isinstance(model_outputs, list):
model_outputs = [model_outputs]
else:
model_outputs = [x_warped, flow]
return Model(inputs=[x_in], outputs=model_outputs, name=model_name)
def __init__(self, fixed, moving, x, y, z, masked_loss=True):
"""
parameters of the class are
:param fixed: fixed image array
:param moving: moving image array to be registered
:param x: translation in X as a trainable tf variable
:param y: translation in Y as a trainable tf variable
:param z: translation in Z as a trainable tf variable
:param masked_loss: whether to calculate the loss wrt to the masked head only (remove bg)
"""
self.stn = SpatialTransformer(interp_method='linear', indexing='ij')
self.iteartion = 0
self.fixed = fixed
self.moving = moving
self.masked_loss = masked_loss
self.x = x
self.y = y
self.z = z
self.losses = []
def apply_transformation(
x_source,
x_transformation,
output_shape,
conditioning_input_shape,
transform_name,
apply_flow_transform=True,
apply_color_transform=False,
flow_indexing='xy',
color_transform_type='WB',
):
n_dims = len(conditioning_input_shape) - 1
transformation_shape = x_transformation.get_shape().as_list()[1:]
x_transformation = Reshape(
transformation_shape,
name='{}_dec_out'.format(transform_name))(x_transformation)
if apply_flow_transform:
# apply flow transform
im_out = SpatialTransformer(name='spatial_transformer',
indexing=flow_indexing)(
[x_source, x_transformation])
elif apply_color_transform:
# apply color transform
print('Applying color transform {}'.format(color_transform_type))
if color_transform_type == 'delta':
x_color_out = Add()([x_source, x_transformation])
elif color_transform_type == 'mult':
x_color_out = Multiply()([x_source, x_transformation])
else:
raise NotImplementedError(
'Only color transform types delta and mult are supported!')
im_out = Reshape(output_shape, name='color_transformer')(x_color_out)
else:
im_out = x_transformation
return im_out, x_transformation
def cvae_learned_prior_trainer_wrapper(
ae_input_shapes,
ae_input_names,
conditioning_input_shapes,
conditioning_input_names,
latent_distribution='normal',
output_shape=None,
model_name='transformer_trainer',
seg_model=None,
transform_encoder_model=None,
transformer_model=None,
prior_encoder_model=None,
transform_type='flow',
transform_latent_shape=(50, ),
include_aug_matrix=False,
n_outputs=1,
):
'''''' '''''' '''''' '''
VTE transformer train model
- takes I, I+J as input
- encodes I+J to z
- condition_on_image = True means that the transform is decoded from the transform+image embedding,
otherwise it is decoded from only the transform embedding
- decodes latent embedding into transform and applies it
''' '''''' '''''' ''''''
ae_inputs, ae_stack, conditioning_inputs, cond_stack = _collect_inputs(
ae_input_shapes, ae_input_names, conditioning_input_shapes,
conditioning_input_names)
conditioning_input_shape = cond_stack.get_shape().as_list()[1:]
inputs = ae_inputs + conditioning_inputs
if include_aug_matrix:
T_in = Input((3, 3), name='transform_input')
inputs += [T_in]
# ae_stack = SpatialTransformer(name='st_affine_stack')([ae_stack, T_in])
# cond_stack = SpatialTransformer(name='st_affine_img')([cond_stack, T_in])
ae_inputs = [
SpatialTransformer(name='st_affine_{}'.format(ae_input_names[ii]))(
[ae_input, T_in]) for ii, ae_input in enumerate(ae_inputs)
]
conditioning_inputs = [
SpatialTransformer(
name='st_affine_{}'.format(conditioning_input_names[ii]))(
[cond_input, T_in])
for ii, cond_input in enumerate(conditioning_inputs)
]
segs = seg_model(conditioning_inputs[0])
if latent_distribution == 'normal':
z_mean_prior, z_logvar_prior = prior_encoder_model(
conditioning_inputs) # + [segs])
z_mean_prior = Reshape(transform_latent_shape,
name='latent_mean_prior')(z_mean_prior)
z_logvar_prior = Reshape(transform_latent_shape,
name='latent_logvar_prior')(z_logvar_prior)
# encode x_stacked into z
z_mean, z_logvar = transform_encoder_model(ae_inputs) # + [segs])
z_mean = Reshape(transform_latent_shape, name='latent_mean')(z_mean)
z_logvar = Reshape(transform_latent_shape,
name='latent_logvar')(z_logvar)
z_sampled = Lambda(transform_network_utils.sampling,
output_shape=transform_latent_shape,
name='lambda_sampling')([z_mean, z_logvar])
z_outputs = [z_mean, z_logvar, z_mean_prior, z_logvar_prior]
elif latent_distribution == 'categorical':
z_prior = prior_encoder_model(conditioning_inputs) # + [segs])
z_post = transform_encoder_model(ae_inputs)
temperature = 5.
z_sampled = Lambda(
lambda x: gumbel_softmax(x, temperature=temperature, hard=False),
name='lambda_sampling_gumbel')(z_post)
# actually normalize outputs to a probability, since this is what the KL loss expects
z_prior = Activation('softmax',
name='latent_categorical_prior')(z_prior)
z_post = Activation('softmax', name='latent_categorical')(z_post)
z_outputs = [z_post, z_prior]
decoder_out = transformer_model(conditioning_inputs + [segs, z_sampled])
if transform_type == 'flow':
im_out, transform_out = decoder_out
transform_shape = transform_out.get_shape().as_list()[1:]
transform_out = Reshape(transform_shape,
name='decoder_flow_out')(transform_out)
im_out = Reshape(output_shape, name='spatial_transformer')(im_out)
elif transform_type == 'color':
im_out, transform_out = decoder_out
transform_out = Reshape(output_shape,
name='decoder_color_out')(transform_out)
im_out = Reshape(output_shape, name='color_transformer')(im_out)
else:
im_out = decoder_out
if transform_type is not None:
return Model(inputs=inputs,
outputs=[im_out] * n_outputs + [transform_out] +
z_outputs,
name=model_name)
else:
return Model(inputs=inputs,
outputs=[im_out] * n_outputs + z_outputs,
name=model_name)
def cvae_trainer_wrapper(
ae_input_shapes,
ae_input_names,
conditioning_input_shapes,
conditioning_input_names,
output_shape=None,
model_name='transformer_trainer',
transform_encoder_model=None,
transformer_model=None,
transform_type='flow',
transform_latent_shape=(50, ),
include_aug_matrix=False,
n_outputs=1,
):
'''''' '''''' '''''' '''
VTE transformer train model
- takes I, I+J as input
- encodes I+J to z
- condition_on_image = True means that the transform is decoded from the transform+image embedding,
otherwise it is decoded from only the transform embedding
- decodes latent embedding into transform and applies it
''' '''''' '''''' ''''''
ae_inputs, ae_stack, conditioning_inputs, cond_stack = _collect_inputs(
ae_input_shapes, ae_input_names, conditioning_input_shapes,
conditioning_input_names)
conditioning_input_shape = cond_stack.get_shape().as_list()[1:]
inputs = ae_inputs + conditioning_inputs
if include_aug_matrix:
T_in = Input((3, 3), name='transform_input')
inputs += [T_in]
# ae_stack = SpatialTransformer(name='st_affine_stack')([ae_stack, T_in])
# cond_stack = SpatialTransformer(name='st_affine_img')([cond_stack, T_in])
ae_inputs = [
SpatialTransformer(name='st_affine_{}'.format(ae_input_names[ii]))(
[ae_input, T_in]) for ii, ae_input in enumerate(ae_inputs)
]
conditioning_inputs = [
SpatialTransformer(
name='st_affine_{}'.format(conditioning_input_names[ii]))(
[cond_input, T_in])
for ii, cond_input in enumerate(conditioning_inputs)
]
# encode x_stacked into z
z_mean, z_logvar = transform_encoder_model(ae_inputs)
z_mean = Reshape(transform_latent_shape, name='latent_mean')(z_mean)
z_logvar = Reshape(transform_latent_shape, name='latent_logvar')(z_logvar)
z_sampled = Lambda(transform_network_utils.sampling,
output_shape=transform_latent_shape,
name='lambda_sampling')([z_mean, z_logvar])
decoder_out = transformer_model(conditioning_inputs + [z_sampled])
if transform_type == 'flow':
im_out, transform_out = decoder_out
transform_shape = transform_out.get_shape().as_list()[1:]
transform_out = Reshape(transform_shape,
name='decoder_flow_out')(transform_out)
im_out = Reshape(output_shape, name='spatial_transformer')(im_out)
elif transform_type == 'color':
im_out, transform_out = decoder_out
transform_out = Reshape(output_shape,
name='decoder_color_out')(transform_out)
im_out = Reshape(output_shape, name='color_transformer')(im_out)
else:
im_out = decoder_out
if transform_type is not None:
return Model(inputs=inputs,
outputs=[im_out] * n_outputs +
[transform_out, z_mean, z_logvar],
name=model_name)
else:
return Model(inputs=inputs,
outputs=[im_out] * n_outputs + [z_mean, z_logvar],
name=model_name)