Code-Note/segmentation-dl/vnet.py at master · jizhang02/Code-Note · GitHub

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
"""
Diogo Amorim, 2018-07-10
V-Net implementation in Keras 2
https://arxiv.org/pdf/1606.04797.pdf
"""
from keras.layers import *
from keras import backend as K
from keras.layers import Conv3D, Input
from keras.models import Model
from keras.layers.advanced_activations import PReLU

def vnet(input_size=(128, 128, 64, 1),filters = 16, batch_norm = True):
    if batch_norm:
        # Layer 1
        input = Input(input_size)
        print("the input tensor:",input.get_shape())
        conv1 = Conv3D(filters, kernel_size=(5, 5, 5), strides=1, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(input)
        print("the conv1 tensor:",conv1.get_shape())
        conv1 = PReLU()(conv1)
        print("the conv1 tensor after PReLU:",conv1.get_shape())
        add1 = add([input,conv1])
        down1 = Conv3D(filters*2, 2, strides=2, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(add1)
        down1 = PReLU()(down1)
        down1 = BatchNormalization()(down1)
        print("the down1 tensor:",down1.get_shape())
        # Layer 2
        conv2_1 = Conv3D(filters*2, kernel_size=5, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(down1)
        conv2 = PReLU()(conv2_1)
        conv2 = Conv3D(filters*2, kernel_size=5, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv2)
        conv2 = PReLU()(conv2)
        add2 = add([conv2_1,conv2])
        down2 = Conv3D(filters*4,2, strides=2, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(add2)
        down2 = PReLU()(down2)
        down2 = BatchNormalization()(down2)
        print("the down2 tensor:", down2.get_shape())
        # Layer 3
        conv3_1 = Conv3D(filters*4, kernel_size=5, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(down2)
        conv3 = PReLU()(conv3_1)
        conv3 = Conv3D(filters*4, kernel_size=5, padding='same',data_format = 'channels_last', kernel_initializer='he_normal')(conv3)
        conv3 = PReLU()(conv3)
        conv3 = Conv3D(filters*4, kernel_size=5, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv3)
        conv3 = PReLU()(conv3)
        add3 = add([conv3_1,conv3])
        down3 = Conv3D(filters*8,2, strides=2, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(add3)
        down3 = PReLU()(down3)
        down3 = BatchNormalization()(down3)
        print("the down3 tensor:", down3.get_shape())

        # Layer 4
        conv4_1 = Conv3D(filters*8, kernel_size=5, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(down3)
        conv4 = PReLU()(conv4_1)
        conv4 = Conv3D(filters*8, kernel_size=5, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv4)
        conv4 = PReLU()(conv4)
        conv4 = Conv3D(filters*8, kernel_size=5, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv4)
        conv4 = PReLU()(conv4)
        add4 = add([conv4_1,conv4])
        down4 = Conv3D(filters*16,2, strides=2, padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(add4)
        down4 = PReLU()(down4)
        down4 = BatchNormalization()(down4)
        print("the down4 tensor:", down4.get_shape())

        # Layer 5:Bottom
        conv5_1 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(down4)
        conv5 = PReLU()(conv5_1)
        conv5 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv5)
        conv5 = PReLU()(conv5)
        conv5 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv5)
        conv5 = PReLU()(conv5)
        add5 = add([conv5_1, conv5])
        up5 = Deconvolution3D(filters*16, kernel_size=(2, 2, 2), padding='same', data_format = 'channels_last',strides=(2, 2, 2))(add5)
        up5 = PReLU()(up5)
        up5 = BatchNormalization()(up5)
        print("the up5 tensor:", up5.get_shape())

        # Layer 6
        merge6 = concatenate([add4,up5], axis=4)
        conv6_1 = Conv3D(filters*16,kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(merge6)
        conv6 = PReLU()(conv6_1)
        conv6 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv6)
        conv6 = PReLU()(conv6)
        conv6 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv6)
        conv6 = PReLU()(conv6)
        add6 = add([conv6_1,conv6])
        up6 = Deconvolution3D(filters*8, kernel_size=(2, 2, 2), padding='same', data_format = 'channels_last',strides=(2, 2, 2))(add6)
        up6 = PReLU()(up6)
        up6 = BatchNormalization()(up6)
        print("the up6 tensor:", up6.get_shape())

        # Layer7
        merge7 = concatenate([add3,up6],axis=4)
        conv7_1 = Conv3D(filters*8, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(merge7)
        conv7 = PReLU()(conv7_1)
        conv7 = Conv3D(filters*8, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv7)
        conv7 = PReLU()(conv7)
        conv7 = Conv3D(filters*8, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv7)
        conv7 = PReLU()(conv7)
        add7 = add([conv7_1, conv7])
        up7 = Deconvolution3D(filters*4, kernel_size=(2, 2, 2), padding='same', data_format = 'channels_last',strides=(2, 2, 2))(add7)
        up7 = PReLU()(up7)
        up7 = BatchNormalization()(up7)
        print("the up7 tensor:", up7.get_shape())

        # Layer8
        merge8 = concatenate([add2,up7],axis=4)
        conv8_1 = Conv3D(filters*4, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(merge8)
        conv8 = PReLU()(conv8_1)
        conv8 = Conv3D(filters*4, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(conv8)
        conv8 = PReLU()(conv8)
        add8 = add([conv8_1, conv8])
        up8 = Deconvolution3D(filters*2, kernel_size=(2, 2, 2), padding='same', data_format = 'channels_last',strides=(2, 2, 2))(add8)
        up8 = PReLU()(up8)
        up8 = BatchNormalization()(up8)
        print("the up8 tensor:", up8.get_shape())

        # Layer 9
        merged9 = concatenate([add1,up8], axis=4)
        conv9_1 = Conv3D(filters*2, kernel_size=(5, 5, 5), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(merged9)
        conv9 = PReLU()(conv9_1)
        add9 = add([conv9_1,conv9])
        conv9 = Conv3D(2, kernel_size=(1, 1, 1), padding='same', data_format = 'channels_last',kernel_initializer='he_normal')(add9)
        conv9 = BatchNormalization()(conv9)

        sigmoid = Conv3D(1, kernel_size=(1, 1, 1), padding='same', data_format = 'channels_last',kernel_initializer='he_normal',activation='sigmoid')(conv9)
        #softmax = Conv3D(5, kernel_size=(1, 1, 1), padding='same', kernel_initializer='he_normal',activation='softmax')(conv9)
        print("the final sigmoid tensor:", sigmoid.get_shape())
        #print("the final softmax tensor:", softmax.get_shape())


        model = Model(inputs = input, outputs = sigmoid)
    else:
        # Layer 1
        input = Input(input_size)
        print("the input tensor:", input.get_shape())
        conv1 = Conv3D(filters, kernel_size=(5, 5, 5), strides=1, padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(input)
        print("the conv1 tensor:", conv1.get_shape())
        conv1 = PReLU()(conv1)
        print("the conv1 tensor after PReLU:", conv1.get_shape())
        add1 = add([input, conv1])
        down1 = Conv3D(filters*2, 2, strides=2, padding='same', data_format='channels_last', kernel_initializer='he_normal')(
            add1)
        down1 = PReLU()(down1)
        print("the down1 tensor:", down1.get_shape())
        # Layer 2
        conv2_1 = Conv3D(filters*2, kernel_size=5, padding='same', data_format='channels_last',
                         kernel_initializer='he_normal')(down1)
        conv2 = PReLU()(conv2_1)
        conv2 = Conv3D(filters*2, kernel_size=5, padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv2)
        conv2 = PReLU()(conv2)
        add2 = add([conv2_1, conv2])
        down2 = Conv3D(filters*4, 2, strides=2, padding='same', data_format='channels_last', kernel_initializer='he_normal')(
            add2)
        down2 = PReLU()(down2)
        print("the down2 tensor:", down2.get_shape())
        # Layer 3
        conv3_1 = Conv3D(filters*4, kernel_size=5, padding='same', data_format='channels_last',
                         kernel_initializer='he_normal')(down2)
        conv3 = PReLU()(conv3_1)
        conv3 = Conv3D(filters*4, kernel_size=5, padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv3)
        conv3 = PReLU()(conv3)
        conv3 = Conv3D(filters*4, kernel_size=5, padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv3)
        conv3 = PReLU()(conv3)
        add3 = add([conv3_1, conv3])
        down3 = Conv3D(filters*8, 2, strides=2, padding='same', data_format='channels_last', kernel_initializer='he_normal')(
            add3)
        down3 = PReLU()(down3)
        print("the down3 tensor:", down3.get_shape())

        # Layer 4
        conv4_1 = Conv3D(filters*8, kernel_size=5, padding='same', data_format='channels_last',
                         kernel_initializer='he_normal')(down3)
        conv4 = PReLU()(conv4_1)
        conv4 = Conv3D(filters*8, kernel_size=5, padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv4)
        conv4 = PReLU()(conv4)
        conv4 = Conv3D(filters*8, kernel_size=5, padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv4)
        conv4 = PReLU()(conv4)
        add4 = add([conv4_1, conv4])
        down4 = Conv3D(filters*16, 2, strides=2, padding='same', data_format='channels_last', kernel_initializer='he_normal')(
            add4)
        down4 = PReLU()(down4)
        print("the down4 tensor:", down4.get_shape())

        # Layer 5
        conv5_1 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                         kernel_initializer='he_normal')(down4)
        conv5 = PReLU()(conv5_1)
        conv5 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv5)
        conv5 = PReLU()(conv5)
        conv5 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv5)
        conv5 = PReLU()(conv5)
        add5 = add([conv5_1, conv5])
        up5 = Deconvolution3D(filters*16, kernel_size=(2, 2, 2), padding='same', data_format='channels_last',
                              strides=(2, 2, 2))(add5)
        up5 = PReLU()(up5)
        print("the up5 tensor:", up5.get_shape())

        # Layer 6
        merge6 = concatenate([add4, up5], axis=4)
        conv6_1 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                         kernel_initializer='he_normal')(merge6)
        conv6 = PReLU()(conv6_1)
        conv6 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv6)
        conv6 = PReLU()(conv6)
        conv6 = Conv3D(filters*16, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv6)
        conv6 = PReLU()(conv6)
        add6 = add([conv6_1, conv6])
        up6 = Deconvolution3D(filters*8, kernel_size=(2, 2, 2), padding='same', data_format='channels_last',
                              strides=(2, 2, 2))(add6)
        up6 = PReLU()(up6)
        print("the up6 tensor:", up6.get_shape())

        # Layer7
        merge7 = concatenate([add3, up6], axis=4)
        conv7_1 = Conv3D(filters*8, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                         kernel_initializer='he_normal')(merge7)
        conv7 = PReLU()(conv7_1)
        conv7 = Conv3D(filters*8, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv7)
        conv7 = PReLU()(conv7)
        conv7 = Conv3D(filters*8, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv7)
        conv7 = PReLU()(conv7)
        add7 = add([conv7_1, conv7])
        up7 = Deconvolution3D(filters*4, kernel_size=(2, 2, 2), padding='same', data_format='channels_last',
                              strides=(2, 2, 2))(add7)
        up7 = PReLU()(up7)
        print("the up7 tensor:", up7.get_shape())

        # Layer8
        merge8 = concatenate([add2, up7], axis=4)
        conv8_1 = Conv3D(filters*4, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                         kernel_initializer='he_normal')(merge8)
        conv8 = PReLU()(conv8_1)
        conv8 = Conv3D(filters*4, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(conv8)
        conv8 = PReLU()(conv8)
        add8 = add([conv8_1, conv8])
        up8 = Deconvolution3D(filters*4, kernel_size=(2, 2, 2), padding='same', data_format='channels_last',
                              strides=(2, 2, 2))(add8)
        up8 = PReLU()(up8)
        print("the up8 tensor:", up8.get_shape())

        # Layer 9
        merged9 = concatenate([add1, up8], axis=4)
        conv9_1 = Conv3D(filters*2, kernel_size=(5, 5, 5), padding='same', data_format='channels_last',
                         kernel_initializer='he_normal')(merged9)
        conv9 = PReLU()(conv9_1)
        add9 = add([conv9_1, conv9])
        conv9 = Conv3D(2, kernel_size=(1, 1, 1), padding='same', data_format='channels_last',
                       kernel_initializer='he_normal')(add9)

        sigmoid = Conv3D(1, kernel_size=(1, 1, 1), padding='same', data_format='channels_last',
                         kernel_initializer='he_normal', activation='sigmoid')(conv9)
        #softmax = Conv3D(1, kernel_size=(1, 1, 1), padding='same', data_format='channels_last',kernel_initializer='he_normal',activation='softmax')(conv9)
        #print("the final sigmoid tensor:", sigmoid.get_shape())

        model = Model(inputs=input, outputs=sigmoid)

    return model