/
model.py
80 lines (62 loc) · 4.01 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
from keras.layers import Activation, Reshape, Dropout
from keras.layers import AtrousConvolution2D, Convolution2D, MaxPooling2D, ZeroPadding2D
from keras.models import Sequential
#
# The VGG16 keras model is taken from here:
# https://gist.github.com/fchollet/f35fbc80e066a49d65f1688a7e99f069
# The (caffe) structure of DilatedNet is here:
# https://github.com/fyu/dilation/blob/master/models/dilation8_pascal_voc_deploy.prototxt
def get_frontend(input_width, input_height) -> Sequential:
model = Sequential()
# model.add(ZeroPadding2D((1, 1), input_shape=(input_width, input_height, 3)))
model.add(Convolution2D(64, 3, 3, activation='relu', name='conv1_1', input_shape=(input_width, input_height, 3)))
model.add(Convolution2D(64, 3, 3, activation='relu', name='conv1_2'))
model.add(MaxPooling2D((2, 2), strides=(2, 2)))
model.add(Convolution2D(128, 3, 3, activation='relu', name='conv2_1'))
model.add(Convolution2D(128, 3, 3, activation='relu', name='conv2_2'))
model.add(MaxPooling2D((2, 2), strides=(2, 2)))
model.add(Convolution2D(256, 3, 3, activation='relu', name='conv3_1'))
model.add(Convolution2D(256, 3, 3, activation='relu', name='conv3_2'))
model.add(Convolution2D(256, 3, 3, activation='relu', name='conv3_3'))
model.add(MaxPooling2D((2, 2), strides=(2, 2)))
model.add(Convolution2D(512, 3, 3, activation='relu', name='conv4_1'))
model.add(Convolution2D(512, 3, 3, activation='relu', name='conv4_2'))
model.add(Convolution2D(512, 3, 3, activation='relu', name='conv4_3'))
# Compared to the original VGG16, we skip the next 2 MaxPool layers,
# and go ahead with dilated convolutional layers instead
model.add(AtrousConvolution2D(512, 3, 3, atrous_rate=(2, 2), activation='relu', name='conv5_1'))
model.add(AtrousConvolution2D(512, 3, 3, atrous_rate=(2, 2), activation='relu', name='conv5_2'))
model.add(AtrousConvolution2D(512, 3, 3, atrous_rate=(2, 2), activation='relu', name='conv5_3'))
# Compared to the VGG16, we replace the FC layer with a convolution
model.add(AtrousConvolution2D(4096, 7, 7, atrous_rate=(4, 4), activation='relu', name='fc6'))
model.add(Dropout(0.5))
model.add(Convolution2D(4096, 1, 1, activation='relu', name='fc7'))
model.add(Dropout(0.5))
# Note: this layer has linear activations, not ReLU
model.add(Convolution2D(21, 1, 1, activation='linear', name='fc-final'))
# model.layers[-1].output_shape == (None, 16, 16, 21)
return model
def add_softmax(model: Sequential) -> Sequential:
""" Append the softmax layers to the frontend or frontend + context net. """
# The softmax layer doesn't work on the (width, height, channel)
# shape, so we reshape to (width*height, channel) first.
# https://github.com/fchollet/keras/issues/1169
_, curr_width, curr_height, curr_channels = model.layers[-1].output_shape
model.add(Reshape((curr_width * curr_height, curr_channels)))
model.add(Activation('softmax'))
# Technically, we need another Reshape here to reshape to 2d, but TF
# the complains when batch_size > 1. We're just going to reshape in numpy.
# model.add(Reshape((curr_width, curr_height, curr_channels)))
return model
def add_context(model: Sequential) -> Sequential:
""" Append the context layers to the frontend. """
model.add(ZeroPadding2D(padding=(33, 33)))
model.add(Convolution2D(42, 3, 3, activation='relu', name='ct_conv1_1'))
model.add(Convolution2D(42, 3, 3, activation='relu', name='ct_conv1_2'))
model.add(AtrousConvolution2D(84, 3, 3, atrous_rate=(2, 2), activation='relu', name='ct_conv2_1'))
model.add(AtrousConvolution2D(168, 3, 3, atrous_rate=(4, 4), activation='relu', name='ct_conv3_1'))
model.add(AtrousConvolution2D(336, 3, 3, atrous_rate=(8, 8), activation='relu', name='ct_conv4_1'))
model.add(AtrousConvolution2D(672, 3, 3, atrous_rate=(16, 16), activation='relu', name='ct_conv5_1'))
model.add(Convolution2D(672, 3, 3, activation='relu', name='ct_fc1'))
model.add(Convolution2D(21, 1, 1, name='ct_final'))
return model