-
Notifications
You must be signed in to change notification settings - Fork 2
/
networks.py
326 lines (266 loc) · 12.8 KB
/
networks.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
"""
Networks for voxelmorph model
In general, these are fairly specific architectures that were designed for the presented papers.
However, the VoxelMorph concepts are not tied to a very particular architecture, and we
encourage you to explore architectures that fit your needs.
see e.g. more powerful unet function in https://github.com/adalca/neuron/blob/master/neuron/models.py
"""
# main imports
import sys
# third party
import numpy as np
import keras.backend as K
from keras.models import Model
import keras.layers as KL
from keras.layers import Layer
from keras.layers import Conv3D, Activation, Input, UpSampling3D, concatenate
from keras.layers import LeakyReLU, Reshape, Lambda
from keras.initializers import RandomNormal
import keras.initializers
import tensorflow as tf
# import neuron layers, which will be useful for Transforming.
sys.path.append('../ext/neuron')
sys.path.append('../ext/pynd-lib')
sys.path.append('../ext/pytools-lib')
import neuron.layers as nrn_layers
import neuron.models as nrn_models
import neuron.utils as nrn_utils
# other vm functions
import losses
def unet_core(vol_size, enc_nf, dec_nf, full_size=True, src=None, tgt=None, src_feats=1, tgt_feats=1):
"""
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
upsample_layer = getattr(KL, 'UpSampling%dD' % ndims)
# inputs
if src is None:
src = Input(shape=[*vol_size, src_feats])
if tgt is None:
tgt = Input(shape=[*vol_size, tgt_feats])
x_in = concatenate([src, tgt])
# down-sample path (encoder)
x_enc = [x_in]
for i in range(len(enc_nf)):
x_enc.append(conv_block(x_enc[-1], enc_nf[i], 2))
# up-sample path (decoder)
x = conv_block(x_enc[-1], dec_nf[0])
x = upsample_layer()(x)
x = concatenate([x, x_enc[-2]])
x = conv_block(x, dec_nf[1])
x = upsample_layer()(x)
x = concatenate([x, x_enc[-3]])
x = conv_block(x, dec_nf[2])
x = upsample_layer()(x)
x = concatenate([x, x_enc[-4]])
x = conv_block(x, dec_nf[3])
x = conv_block(x, 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
if full_size:
x = upsample_layer()(x)
x = concatenate([x, x_enc[0]])
x = conv_block(x, dec_nf[5])
# optional convolution at output resolution (used in voxelmorph-2)
if len(dec_nf) == 7:
x = conv_block(x, dec_nf[6])
return Model(inputs=[src, tgt], outputs=[x])
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)
: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
model = Model(inputs=[src, tgt], outputs=[y, flow])
return model
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 nn_trf(vol_size, indexing='xy'):
"""
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=indexing)
nn_spatial_output = nn_output([subj_input, trf_input])
return keras.models.Model([subj_input, trf_input], nn_spatial_output)
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])