CelestialSurveyor

1
# Slow Fast Model implementation taken from https://github.com/xuzheyuan624/slowfast-keras with some modifications
2
from keras.layers import (Conv3D, BatchNormalization, ReLU, Add, MaxPool3D, GlobalAveragePooling3D,
3
                          Concatenate, Dropout, Dense, Lambda, Input)
4
from keras.models import Model
5
from keras import Sequential
6

7

8
def Conv_BN_ReLU(planes, kernel_size, strides=(1, 1, 1), padding='same', use_bias=False):
9
    return Sequential([
10
        Conv3D(planes, kernel_size, strides=strides, padding=padding, use_bias=use_bias),
11
        BatchNormalization(),
12
        ReLU()
13
    ])
14

15

16
def bottleneck(x, planes, stride=1, downsample=None, head_conv=1, use_bias=False):
17
    residual = x
18
    if head_conv == 1:
19
        x = Conv_BN_ReLU(planes, kernel_size=1, use_bias=use_bias)(x)
20
    elif head_conv == 3:
21
        x = Conv_BN_ReLU(planes, kernel_size=(3, 1, 1), use_bias=use_bias)(x)
22
    else:
23
        raise ValueError('Unsupported head_conv!!!')
24
    x = Conv_BN_ReLU(planes, kernel_size=(1, 3, 3), strides=(1, stride, stride), use_bias=use_bias)(x)
25
    x = Conv3D(planes*4, kernel_size=1, use_bias=use_bias)(x)
26
    x = BatchNormalization()(x)
27
    if downsample is not None:
28
        residual = downsample(residual)
29
    x = Add()([x, residual])
30
    x = ReLU()(x)
31
    return x
32

33

34
def datalayer(x, stride):
35
    return x[:, ::stride, :, :, :]
36

37

38
def SlowFast_body(layers, block, dropout=0.5):
39
    inputs = Input(shape=(None, 64, 64, 1))
40
    inputs_fast = Lambda(datalayer, name='data_fast', arguments={'stride':1})(inputs)
41
    inputs_slow = Lambda(datalayer, name='data_slow', arguments={'stride':4})(inputs)
42
    fast, lateral = Fast_body(inputs_fast, layers, block)
43
    slow = Slow_body(inputs_slow, lateral, layers, block)
44
    x = Concatenate()([slow, fast])
45
    x = Dropout(dropout)(x)
46
    x = Dense(64, activation='relu')(x)
47
    x = Dense(32, activation='relu')(x)
48
    out = Dense(1, activation='sigmoid')(x)
49
    return Model(inputs, out)
50

51

52
def Fast_body(x, layers, block):
53
    fast_inplanes = 8
54
    lateral = []
55
    x = Conv_BN_ReLU(8, kernel_size=(5, 3, 3), strides=(1, 2, 2))(x)
56
    x = MaxPool3D(pool_size=(1, 3, 3), strides=(1, 2, 2), padding='same')(x)
57
    lateral_p1 = Conv3D(8*2, kernel_size=(5, 1, 1), strides=(4, 1, 1), padding='same', use_bias=False)(x)
58
    lateral.append(lateral_p1)
59
    x, fast_inplanes = make_layer_fast(x, block, 8, layers[0], head_conv=3, fast_inplanes=fast_inplanes)
60
    lateral_res2 = Conv3D(32*2, kernel_size=(5, 1, 1), strides=(4, 1, 1), padding='same', use_bias=False)(x)
61
    lateral.append(lateral_res2)
62
    x, fast_inplanes = make_layer_fast(x, block, 16, layers[1], stride=2, head_conv=3, fast_inplanes=fast_inplanes)
63
    lateral_res3 = Conv3D(64*2, kernel_size=(5, 1, 1), strides=(4, 1, 1), padding='same', use_bias=False)(x)
64
    lateral.append(lateral_res3)
65
    x, fast_inplanes = make_layer_fast(x, block, 32, layers[2], stride=2, head_conv=3, fast_inplanes=fast_inplanes)
66
    lateral_res4 = Conv3D(128*2, kernel_size=(5, 1, 1), strides=(4, 1, 1), padding='same', use_bias=False)(x)
67
    lateral.append(lateral_res4)
68
    x, fast_inplanes = make_layer_fast(x, block, 64, layers[3], stride=2, head_conv=3, fast_inplanes=fast_inplanes)
69
    x = GlobalAveragePooling3D()(x)
70
    return x, lateral
71

72

73
def Slow_body(x, lateral, layers, block):
74
    slow_inplanes = 64 + 64//8*2
75
    x = Conv_BN_ReLU(32, kernel_size=(1, 3, 3), strides=(1, 2, 2))(x)
76
    x = MaxPool3D(pool_size=(1, 3, 3), strides=(1, 2, 2), padding='same')(x)
77
    x = Concatenate()([x, lateral[0]])
78
    x, slow_inplanes = make_layer_slow(x, block, 32, layers[0], head_conv=1, slow_inplanes=slow_inplanes)
79
    x = Concatenate()([x, lateral[1]])
80
    x, slow_inplanes = make_layer_slow(x, block, 64, layers[1], stride=2, head_conv=1, slow_inplanes=slow_inplanes)
81
    x = Concatenate()([x, lateral[2]])
82
    x, slow_inplanes = make_layer_slow(x, block, 128, layers[2], stride=2, head_conv=1, slow_inplanes=slow_inplanes)
83
    x = Concatenate()([x, lateral[3]])
84
    x, slow_inplanes = make_layer_slow(x, block, 256, layers[3], stride=2, head_conv=1, slow_inplanes=slow_inplanes)
85
    x = GlobalAveragePooling3D()(x)
86
    return x
87

88

89
def make_layer_fast(x, block, planes, blocks, stride=1, head_conv=1, fast_inplanes=8, block_expansion=4):
90
    downsample = None
91
    if stride != 1 or fast_inplanes != planes * block_expansion:
92
        downsample = Sequential([
93
            Conv3D(planes*block_expansion, kernel_size=1, strides=(1, stride, stride), use_bias=False),
94
            BatchNormalization()
95
        ])
96
    fast_inplanes = planes * block_expansion
97
    x = block(x, planes, stride, downsample=downsample, head_conv=head_conv)
98
    for _ in range(1, blocks):
99
        x = block(x, planes, head_conv=head_conv)
100
    return x, fast_inplanes
101

102

103
def make_layer_slow(x, block, planes, blocks, stride=1, head_conv=1, slow_inplanes=80, block_expansion=4):
104
    downsample = None
105
    if stride != 1 or slow_inplanes != planes * block_expansion:
106
        downsample = Sequential([
107
            Conv3D(planes*block_expansion, kernel_size=1, strides = (1, stride, stride), use_bias=False),
108
            BatchNormalization()
109
        ])
110
    x = block(x, planes, stride, downsample, head_conv=head_conv)
111
    for _ in range(1, blocks):
112
        x = block(x, planes, head_conv=head_conv)
113
    slow_inplanes = planes * block_expansion + planes * block_expansion//8*2
114
    return x, slow_inplanes
115