-
Notifications
You must be signed in to change notification settings - Fork 15
/
unet_keras.py
110 lines (91 loc) · 4.72 KB
/
unet_keras.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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import os
import matplotlib.pyplot as plt
from keras.layers.core import Dropout
from keras.layers.convolutional import Conv2D, Conv2DTranspose
from keras.layers.pooling import MaxPooling2D
from keras.layers import Input
from keras.models import Model
from keras.layers.merge import concatenate
from keras.callbacks import EarlyStopping, ModelCheckpoint
from utils import img_batch_generator, get_train_val_paths
from unet.loss import dice_loss
from unet.metrics import mean_iou
image_folder = "./data/2d_images/"
masks_folder = "./data/2d_masks/"
tr_paths, v_paths = get_train_val_paths(image_folder, masks_folder)
inputs = Input((512, 512, 1))
c1 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(inputs)
c1 = Dropout(0.1)(c1)
c1 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c1)
p1 = MaxPooling2D((2, 2))(c1)
c2 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(p1)
c2 = Dropout(0.1)(c2)
c2 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c2)
p2 = MaxPooling2D((2, 2))(c2)
c3 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(p2)
c3 = Dropout(0.2)(c3)
c3 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c3)
p3 = MaxPooling2D((2, 2))(c3)
c4 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(p3)
c4 = Dropout(0.2)(c4)
c4 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c4)
p4 = MaxPooling2D(pool_size=(2, 2))(c4)
c5 = Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(p4)
c5 = Dropout(0.3)(c5)
c5 = Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c5)
u6 = Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c5)
u6 = concatenate([u6, c4])
c6 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(u6)
c6 = Dropout(0.2)(c6)
c6 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c6)
u7 = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(c6)
u7 = concatenate([u7, c3])
c7 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(u7)
c7 = Dropout(0.2)(c7)
c7 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c7)
u8 = Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(c7)
u8 = concatenate([u8, c2])
c8 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(u8)
c8 = Dropout(0.1)(c8)
c8 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c8)
u9 = Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same')(c8)
u9 = concatenate([u9, c1], axis=3)
c9 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(u9)
c9 = Dropout(0.1)(c9)
c9 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same')(c9)
outputs = Conv2D(1, (1, 1), activation='sigmoid')(c9)
model = Model(inputs=[inputs], outputs=[outputs])
model.compile(optimizer='adam', loss=[dice_loss], metrics=[mean_iou])
print(model.summary())
batch_size = 4
model_save_path = "./models"
if not os.path.exists(model_save_path):
os.mkdir(model_save_path)
train_gen = img_batch_generator(tr_paths["train_imgs"], tr_paths["train_mask"], batch_size=batch_size)
val_gen = img_batch_generator(v_paths["val_imgs"], v_paths["val_mask"], batch_size=batch_size)
train_steps = len(tr_paths["train_imgs"]) // batch_size
val_steps = len(v_paths["val_imgs"]) // batch_size
early_stop = EarlyStopping(patience=10, verbose=1)
checkpoint = ModelCheckpoint(os.path.join(model_save_path, "keras_unet_model.h5"),
verbose=1, save_best_only=True)
history = model.fit_generator(train_gen, steps_per_epoch=train_steps,
epochs=50,
validation_data=val_gen,
validation_steps=val_steps,
verbose=1,
max_queue_size=4,
callbacks=[early_stop, checkpoint])
loss = history.history["loss"]
val_loss = history.history["val_loss"]
miou = history.history["mean_iou"]
val_miou = history.history["val_mean_iou"]
epochs = range(1, len(loss) + 1)
plt.plot(epochs, loss, color="red", label="training loss")
plt.plot(epochs, val_loss, color="blue", label="validation loss")
plt.title("Training and Validation Loss")
plt.legend()
plt.figure()
plt.plot(epochs, miou, color="red", label="training MeanIOU")
plt.plot(epochs, val_miou, color="blue", label="validation MeanIOU")
plt.title("Training and Validation IOU")
plt.legend()