-
Notifications
You must be signed in to change notification settings - Fork 62
/
Copy pathrainnet.py
91 lines (74 loc) · 4.01 KB
/
rainnet.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
from keras.models import *
from keras.layers import *
def rainnet(input_shape=(928, 928, 4), mode="regression"):
"""
The function for building the RainNet (v1.0) model from scratch
using Keras functional API.
Parameters:
input size: tuple(W x H x C), where W (width) and H (height)
describe spatial dimensions of input data (e.g., 928x928 for RY data);
and C (channels) describes temporal (depth) dimension of
input data (e.g., 4 means accounting four latest radar scans at time
t-15, t-10, t-5 minutes, and t)
mode: "regression" (default) or "segmentation".
For "regression" mode the last activation function is linear,
while for "segmentation" it is sigmoid.
To train RainNet to predict continuous precipitation intensities use
"regression" mode.
RainNet could be trained to predict the exceedance of specific intensity
thresholds. For that purpose, use "segmentation" mode.
"""
inputs = Input(input_shape)
conv1f = Conv2D(64, 3, padding='same', kernel_initializer='he_normal')(inputs)
conv1f = Activation("relu")(conv1f)
conv1s = Conv2D(64, 3, padding='same', kernel_initializer='he_normal')(conv1f)
conv1s = Activation("relu")(conv1s)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1s)
conv2f = Conv2D(128, 3, padding='same', kernel_initializer='he_normal')(pool1)
conv2f = Activation("relu")(conv2f)
conv2s = Conv2D(128, 3, padding='same', kernel_initializer='he_normal')(conv2f)
conv2s = Activation("relu")(conv2s)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2s)
conv3f = Conv2D(256, 3, padding='same', kernel_initializer='he_normal')(pool2)
conv3f = Activation("relu")(conv3f)
conv3s = Conv2D(256, 3, padding='same', kernel_initializer='he_normal')(conv3f)
conv3s = Activation("relu")(conv3s)
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3s)
conv4f = Conv2D(512, 3, padding='same', kernel_initializer='he_normal')(pool3)
conv4f = Activation("relu")(conv4f)
conv4s = Conv2D(512, 3, padding='same', kernel_initializer='he_normal')(conv4f)
conv4s = Activation("relu")(conv4s)
drop4 = Dropout(0.5)(conv4s)
pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)
conv5f = Conv2D(1024, 3, padding='same', kernel_initializer='he_normal')(pool4)
conv5f = Activation("relu")(conv5f)
conv5s = Conv2D(1024, 3, padding='same', kernel_initializer='he_normal')(conv5f)
conv5s = Activation("relu")(conv5s)
drop5 = Dropout(0.5)(conv5s)
up6 = concatenate([UpSampling2D(size=(2, 2))(drop5), conv4s], axis=3)
conv6 = Conv2D(512, 3, padding='same', kernel_initializer='he_normal')(up6)
conv6 = Activation("relu")(conv6)
conv6 = Conv2D(512, 3, padding='same', kernel_initializer='he_normal')(conv6)
conv6 = Activation("relu")(conv6)
up7 = concatenate([UpSampling2D(size=(2, 2))(conv6), conv3s], axis=3)
conv7 = Conv2D(256, 3, padding='same', kernel_initializer='he_normal')(up7)
conv7 = Activation("relu")(conv7)
conv7 = Conv2D(256, 3, padding='same', kernel_initializer='he_normal')(conv7)
conv7 = Activation("relu")(conv7)
up8 = concatenate([UpSampling2D(size=(2, 2))(conv7), conv2s], axis=3)
conv8 = Conv2D(128, 3, padding='same', kernel_initializer='he_normal')(up8)
conv8 = Activation("relu")(conv8)
conv8 = Conv2D(128, 3, padding='same', kernel_initializer='he_normal')(conv8)
conv8 = Activation("relu")(conv8)
up9 = concatenate([UpSampling2D(size=(2, 2))(conv8), conv1s], axis=3)
conv9 = Conv2D(64, 3, padding='same', kernel_initializer='he_normal')(up9)
conv9 = Activation("relu")(conv9)
conv9 = Conv2D(64, 3, padding='same', kernel_initializer='he_normal')(conv9)
conv9 = Activation("relu")(conv9)
conv9 = Conv2D(2, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)
if mode == "regression":
outputs = Conv2D(1, 1, activation='linear')(conv9)
elif mode == "segmentation":
outputs = Conv2D(1, 1, activation='sigmoid')(conv9)
model = Model(inputs=inputs, outputs=outputs)
return model