-
Notifications
You must be signed in to change notification settings - Fork 0
/
mini_unet.py
71 lines (42 loc) · 3.55 KB
/
mini_unet.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
# based on https://github.com/YBIGTA/DL_Models/blob/master/models/unet/Keras%20(on%20Google%20Colaboratory)/mini%20u-net.ipynb
from keras.models import Model
from keras.layers import Input, Conv2D, MaxPooling2D, concatenate, UpSampling2D, Dropout, Cropping2D, Softmax
from keras.preprocessing import image
image_shape = [256, 256, 3]
keep_prob = 0.5 # dropout rate
num_of_classes = 8
def mini_unet_arch(image_shape, keep_prob, num_of_classes):
# Contracting Path
input_image = Input(image_shape)
conv1_1 = Conv2D(filters = 32, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv1_1')(input_image)
conv1_2 = Conv2D(filters = 32, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv1_2')(conv1_1)
pool_1 = MaxPooling2D(name = 'pool_1')(conv1_2)
conv2_1 = Conv2D(filters = 64, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv2_1')(pool_1)
conv2_2 = Conv2D(filters = 64, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv2_2')(conv2_1)
pool_2 = MaxPooling2D(name = 'pool_2')(conv2_2)
conv3_1 = Conv2D(filters = 128, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv3_1')(pool_2)
conv3_2 = Conv2D(filters = 128, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv3_2')(conv3_1)
pool_3 = MaxPooling2D(name = 'pool_3')(conv3_2)
conv4_1 = Conv2D(filters = 256, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv4_1')(pool_3)
conv4_2 = Conv2D(filters = 256, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv4_2')(conv4_1)
# Expanding Path
upconv5_1 = UpSampling2D(name = 'upconv5_1')(conv4_2)
upconv5_2 = Conv2D(filters = 128, kernel_size = (2, 2), activation = 'relu', padding = 'same', name = 'upconv5_2')(upconv5_1)
concat_5 = concatenate([upconv5_2, conv3_2], axis = 3, name = 'concat_5')
conv5_1 = Conv2D(filters = 128, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv5_1')(concat_5)
conv5_2 = Conv2D(filters = 128, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv5_2')(conv5_1)
upconv6_1 = UpSampling2D(name = 'upconv6_1')(conv5_2)
upconv6_2 = Conv2D(filters = 64, kernel_size = (2, 2), activation = 'relu', padding = 'same', name = 'upconv6_2')(upconv6_1)
concat_6 = concatenate([upconv6_2, conv2_2], axis = 3, name = 'concat_6')
conv6_1 = Conv2D(filters = 64, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv6_1')(concat_6)
conv6_2 = Conv2D(filters = 64, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv6_2')(conv6_1)
upconv7_1 = UpSampling2D(name = 'upconv7_1')(conv6_2)
upconv7_2 = Conv2D(filters = 32, kernel_size = (2, 2), activation = 'relu', padding = 'same', name = 'upconv7_2')(upconv7_1)
concat_7 = concatenate([upconv7_2, conv1_2], axis = 3, name = 'concat_7')
conv7_1 = Conv2D(filters = 32, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv7_1')(concat_7)
conv7_2 = Conv2D(filters = 32, kernel_size = (3, 3), padding = 'same', activation = 'relu', name = 'conv7_2')(conv7_1)
conv8 = Conv2D(filters = num_of_classes, kernel_size = (1, 1), activation = 'softmax', name = 'conv8')(conv7_2)
model = Model(inputs = input_image, outputs = conv8, name = 'model')
return model
mini_unet = mini_unet_arch(image_shape, keep_prob, num_of_classes)
# u_net.compile(loss = categoricalcrossentropy, metrics = ohe_mean_iou)