-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathpredict.py
144 lines (131 loc) · 6.24 KB
/
predict.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
133
134
135
136
137
138
139
140
141
142
143
144
import math, os.path, sys
import numpy as np
import matplotlib.pyplot as plt
import tifffile as tiff
from train_unet import weights_path, get_model, normalize, PATCH_SZ, N_CLASSES
def predict(x, model, patch_sz=160, n_classes=5):
img_height = x.shape[0]
img_width = x.shape[1]
n_channels = x.shape[2]
# make extended img so that it contains integer number of patches
npatches_vertical = math.ceil(img_height / patch_sz)
npatches_horizontal = math.ceil(img_width / patch_sz)
extended_height = patch_sz * npatches_vertical
extended_width = patch_sz * npatches_horizontal
ext_x = np.zeros(shape=(extended_height, extended_width, n_channels), dtype=np.float32)
# fill extended image with mirrors:
ext_x[:img_height, :img_width, :] = x
for i in range(img_height, extended_height):
ext_x[i, :, :] = ext_x[2 * img_height - i - 1, :, :]
for j in range(img_width, extended_width):
ext_x[:, j, :] = ext_x[:, 2 * img_width - j - 1, :]
# now we assemble all patches in one array
patches_list = []
for i in range(0, npatches_vertical):
for j in range(0, npatches_horizontal):
x0, x1 = i * patch_sz, (i + 1) * patch_sz
y0, y1 = j * patch_sz, (j + 1) * patch_sz
patches_list.append(ext_x[x0:x1, y0:y1, :])
# model.predict() needs numpy array rather than a list
patches_array = np.asarray(patches_list)
# predictions:
patches_predict = model.predict(patches_array, batch_size=4)
prediction = np.zeros(shape=(extended_height, extended_width, n_classes), dtype=np.float32)
for k in range(patches_predict.shape[0]):
#only print out debug info if flag is set
if x0_x1_debug == True:
print("K: {}".format(k))
i = k // npatches_horizontal
j = k % npatches_vertical
x0, x1 = i * patch_sz, (i + 1) * patch_sz
#only print out debug info if flag is set
if x0_x1_debug == True:
print("x0: {}, x1: {}".format(x0, x1))
#only print out debug info if flag is set
y0, y1 = j * patch_sz, (j + 1) * patch_sz
#only print out debug info if flag is set
if x0_x1_debug == True:
print("y0: {}, y1: {}".format(y0, y1))
prediction[x0:x1, y0:y1, :] = patches_predict[k, :, :, :]
return prediction[:img_height, :img_width, :]
def picture_from_mask(mask, threshold=0):
colors = {
0: [150, 150, 150], # Buildings
1: [223, 194, 125], # Roads & Tracks
2: [27, 120, 55], # Trees
3: [166, 219, 160], # Crops
4: [116, 173, 209] # Water
}
z_order = {
1: 3,
2: 4,
3: 0,
4: 1,
5: 2
}
pict = 255*np.ones(shape=(3, mask.shape[1], mask.shape[2]), dtype=np.uint8)
for i in range(1, 6):
cl = z_order[i]
for ch in range(3):
pict[ch,:,:][mask[cl,:,:] > threshold] = colors[cl][ch]
return pict
if __name__ == '__main__':
#set debug flag for additional output to help fix predict function
x0_x1_debug = False
model = get_model()
model.load_weights(weights_path)
overwrite_check = ['output/planet_classtest.tif', 'output/result.tif', 'output/map.tif', 'output/planet_result.tif', 'output/planet_map.tif']
for file in overwrite_check:
if os.path.exists(file):
print('ERROR: file {0} already exists. Please rename or delete the following files before creating predictions:'.format(file))
print(*overwrite_check, sep = '\n')
sys.exit()
planet_imagedir = 'data/planet_training/predict/'
image_id = '20180412_143154_1003_1B_AnalyticMS'
img = normalize(tiff.imread(planet_imagedir+'{}.tif'.format(image_id)))
# add 4 channels of 0 to array to predict planet image
# img_pad = ((0,0), (0,0), (0,4))
# img_fixed2 = np.pad(img, pad_width=img_pad, mode='constant', constant_values=0)
#trim the planet image to the same dimensions as training data
# img_fixed2 = img_fixed2[:848, :837, :]
# tiff.imsave('output/planet_classtest.tif', img)
# img = img_fixed2
for i in range(7):
if i == 0: # reverse first dimension
mymat = predict(img[::-1,:,:], model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
print("Case 1",img.shape, mymat.shape)
elif i == 1: # reverse second dimension
temp = predict(img[:,::-1,:], model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
print("Case 2", temp.shape, mymat.shape)
mymat = np.mean( np.array([ temp[:,::-1,:], mymat ]), axis=0 )
elif i == 2: # transpose(interchange) first and second dimensions
#transpose removed to hopefully unbreak script 4/16/19
temp = predict(img, model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
print("Case 3", temp.shape, mymat.shape)
mymat = np.mean( np.array([ temp, mymat ]), axis=0 )
elif i == 3:
#was previously rotating by 90 deg. This and 270 deg rotation does not work with
#rectangular images like ours.
#180 works fine bc its a flip
#circle back and add augmentation after run completes
temp = predict(img, model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
print("Case 4", temp.shape, mymat.shape)
mymat = np.mean( np.array([ temp, mymat ]), axis=0 )
elif i == 4:
temp = predict(np.rot90(img,2), model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
print("Case 5", temp.shape, mymat.shape)
mymat = np.mean( np.array([ np.rot90(temp,-2), mymat ]), axis=0 )
elif i == 5:
#was a 270 degree rotation, which doesn't work with rectangular image
#removing for now, will replace
temp = predict(img, model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
print("Case 6", temp.shape, mymat.shape)
mymat = np.mean( np.array([ temp, mymat ]), axis=0 )
else:
temp = predict(img, model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
print("Case 7", temp.shape, mymat.shape)
mymat = np.mean( np.array([ temp, mymat ]), axis=0 )
#create classified map
map = picture_from_mask(mymat, 0.5)
tiff.imsave('output/result.tif', (255*mymat).astype('uint8'))
tiff.imsave('output/map.tif', map)