-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmodel.py
105 lines (82 loc) · 4.49 KB
/
model.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
from modules import *
import neuron.layers as nrn_layers
class nn_trf(object):
def __init__(self, name='nn_trf'):
self.name = name
def __call__(self, images_1_seg, flow, indexing='ij'):
out = nrn_layers.SpatialTransformer(interp_method='nearest', indexing=indexing)([images_1_seg, flow])
return out
class BiLevelNet(object):
GRAD_IS_ZERO = 1e-12
def __init__(self, num_levels=2, name='pro-net'):
self.num_levels = num_levels
self.output_level = num_levels - 1
self.name = name
self.fp_extractor = Pro_FeatureLearning(self.num_levels)
self.of_estimator = [Estimator_1(name=f'estimator1_{l}') for l in range(self.num_levels)]
self.conv_block = conv_block(name='conv_block')
self.int_steps = 7
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.append(flows)
@property
def vars(self):
return [var for var in tf.global_variables() if self.name in var.name]
class BiLevelNet_train(object):
GRAD_IS_ZERO = 1e-12
def __init__(self, num_levels=2, name='pro-net'):
self.num_levels = num_levels
self.output_level = num_levels - 1
self.name = name
self.fp_extractor = Pro_FeatureLearning(self.num_levels)
self.of_estimator = [Estimator_1(name=f'estimator1_{l}') for l in range(self.num_levels)]
self.conv_block = conv_block(name='conv_block')
self.int_steps = 7
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)
@property
def vars(self):
return [var for var in tf.global_variables() if self.name in var.name]