/
SegNet.py
132 lines (102 loc) · 3.82 KB
/
SegNet.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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import numpy as np
import pandas as pd
import json
import sys
from skimage.io import imread
from matplotlib import pyplot as plt
import os
os.environ['KERAS_BACKEND'] = 'theano'
os.environ['THEANO_FLAGS'] = 'mode=FAST_RUN, device=gpu0, floatX=float32, optimizer=fast_compile'
from keras import models
from keras.optimizers import SGD
path = 'Data/'
img_w = 256
img_h = 256
n_labels = 2
n_train = 6
n_test = 3
def label_map(labels):
label_map = np.zeros([img_h, img_w, n_labels])
for r in range(img_h):
for c in range(img_w):
label_map[r, c, labels[r][c]] = 1
return label_map
def prep_data(mode):
assert mode in {'test', 'train'}, \
'mode should be either \'test\' or \'train\''
data = []
label = []
df = pd.read_csv(path + mode + '.csv')
n = n_train if mode == 'train' else n_test
for i, item in df.iterrows():
if i >= n:
break
img, gt = [imread(path + item[0])], np.clip(imread(path + item[1]), 0, 1)
data.append(img)
label.append(label_map(gt))
sys.stdout.write('\r')
sys.stdout.write(mode + ": [%-20s] %d%%" % ('=' * int(20. * (i + 1) / n - 1) + '>',
int(100. * (i + 1) / n)))
sys.stdout.flush()
sys.stdout.write('\r')
sys.stdout.flush()
data, label = np.array(data), np.array(label).reshape((n, img_h * img_w, n_labels))
print mode + ': OK'
print '\tshapes: {}, {}'.format(data.shape, label.shape)
print '\ttypes: {}, {}'.format(data.dtype, label.dtype)
print '\tmemory: {}, {} MB'.format(data.nbytes / 1048576, label.nbytes / 1048576)
return data, label
def plot_results(output):
gt = []
df = pd.read_csv(path + 'test.csv')
for i, item in df.iterrows():
gt.append(np.clip(imread(path + item[1]), 0, 1))
plt.figure(figsize=(15, 2 * n_test))
for i, item in df.iterrows():
plt.subplot(n_test, 4, 4 * i + 1)
plt.title('Ground Truth')
plt.axis('off')
gt = imread(path + item[1])
plt.imshow(np.clip(gt, 0, 1))
plt.subplot(n_test, 4, 4 * i + 2)
plt.title('Prediction')
plt.axis('off')
labeled = np.argmax(output[i], axis=-1)
plt.imshow(labeled)
plt.subplot(n_test, 4, 4 * i + 3)
plt.title('Heat map')
plt.axis('off')
plt.imshow(output[i][:, :, 1])
plt.subplot(n_test, 4, 4 * i + 4)
plt.title('Comparison')
plt.axis('off')
rgb = np.empty((img_h, img_w, 3))
rgb[:, :, 0] = labeled
rgb[:, :, 1] = imread(path + item[0])
rgb[:, :, 2] = gt
plt.imshow(rgb)
plt.savefig('result.png')
plt.show()
#########################################################################################################
with open('model_5l.json') as model_file:
autoencoder = models.model_from_json(model_file.read())
optimizer = SGD(lr=0.001, momentum=0.9, decay=0.0005, nesterov=False)
autoencoder.compile(loss="categorical_crossentropy", optimizer=optimizer, metrics=['accuracy'])
print 'Compiled: OK'
# Train model or load weights
# train_data, train_label = prep_data('train')
# nb_epoch = 50
# batch_size = 18
# history = autoencoder.fit(train_data, train_label, batch_size=batch_size, nb_epoch=nb_epoch, verbose=1)
# autoencoder.save_weights('model_5l_weight_ep50.hdf5')
autoencoder.load_weights('model_5l_weight_ep50.hdf5')
# Model visualization
from keras.utils.visualize_util import plot
plot(autoencoder, to_file='model.png', show_shapes=True)
test_data, test_label = prep_data('test')
score = autoencoder.evaluate(test_data, test_label, verbose=0)
print 'Test score:', score[0]
print 'Test accuracy:', score[1]
output = autoencoder.predict_proba(test_data, verbose=0)
output = output.reshape((output.shape[0], img_h, img_w, n_labels))
plot_results(output)