-
Notifications
You must be signed in to change notification settings - Fork 4
/
predict.py
executable file
·165 lines (150 loc) · 8.09 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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
import os
os.environ["CUDA_VISIBLE_DEVICES"]="1" # second gpu
import math
import numpy as np
import matplotlib.pyplot as plt
import tifffile as tiff
from os import sys
import itertools
from shapely.geometry.polygon import Polygon
from PIL import Image, ImageDraw
from train_unet import weights_path, get_model, normalize, PATCH_SZ, N_CLASSES
from custom_utils import get_4bands
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]):
i = k // npatches_horizontal
j = k % npatches_vertical
x0, x1 = i * patch_sz, (i + 1) * patch_sz
y0, y1 = j * patch_sz, (j + 1) * patch_sz
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
}
# original z-order = 3,4,0,1,2
z_order = {
1: 3,
2: 2,
3: 4,
4: 0,
5: 1
}
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__':
if len(sys.argv)<2:
print("Enter weight path (eg. ./weights/w28.py)")
sys.exit()
weights_path = sys.argv[1]
print ("----------------------------------------------------------------------------")
model = get_model()
print ("...model loaded")
print ("...loading weights from ", weights_path)
model.load_weights(weights_path)
all_sliced_tifs = os.listdir('/home/ekbana/computer_vision/satellite-image/Planet.com/Planet_Data_Sliced/tif/all-sliced-tiff-850px/')
result_path = "../Planet.com/Planet_Data_Sliced/tif/result/"
all_sliced_tifs = [file for file in all_sliced_tifs if file[-4:] == ".tif"][0:3]
total_files_count = len(all_sliced_tifs)
for current_file_count, test_file in enumerate(all_sliced_tifs[0:]):
test_filename = test_file
print ("...running inference for {}".format(test_filename))
test_file = "../../satellite-image/Planet.com/Planet_Data_Sliced/tif/all-sliced-tiff-850px/" + test_file
class_id = 4
# img = take_4bands(normalize(tiff.imread('data/mband/{}.tif'.format(test_id)).transpose([1,2,0]))) # make channels last
# img = tiff.imread(test_file).transpose([1,2,0])
# print (img)
img = tiff.imread(test_file).transpose([1,2,0])
img, _ = get_4bands(normalize(tiff.imread(test_file).transpose([1,2,0]))) # make channels last
band_list = [0, 1, 2, 3] # to permute all the bands
perm_list = list(itertools.permutations(band_list))
print ("...the shape of given image is: ", img.shape)
res_shape = tuple([8]+list(img.shape[0:2])+[N_CLASSES]) # [8] here comes from the 8 augmented images
results = np.zeros(shape=res_shape)
for i in range(8):
print ("...running inference for augmentation {}".format(i))
# TODO: The average taken in the following way isn't mathematically correct. Refer to the issue in original repo
if i == 0: # reverse first dimension
mymat = predict(img[::-1,:,:], model, patch_sz=PATCH_SZ, n_classes=N_CLASSES).transpose([2,0,1])
# print(mymat[class_id][0][0], mymat[3][12][13])
# 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).transpose([2,0,1])
# print(temp[class_id][0][0], temp[3][12][13])
# 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
temp = predict(img.transpose([1,0,2]), model, patch_sz=PATCH_SZ, n_classes=N_CLASSES).transpose([2,0,1])
# print(temp[class_id][0][0], temp[3][12][13])
# print("Case 3", temp.shape, mymat.shape)
mymat = np.mean( np.array([ temp.transpose(0,2,1), mymat ]), axis=0 )
elif i == 3:
temp = predict(np.rot90(img, 1), model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
# print(temp.transpose([2,0,1])[class_id][0][0], temp.transpose([2,0,1])[3][12][13])
# print("Case 4", temp.shape, mymat.shape)
mymat = np.mean( np.array([ np.rot90(temp, -1).transpose([2,0,1]), mymat ]), axis=0 )
elif i == 4:
temp = predict(np.rot90(img,2), model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
# print(temp.transpose([2,0,1])[class_id][0][0], temp.transpose([2,0,1])[3][12][13])
# print("Case 5", temp.shape, mymat.shape)
mymat = np.mean( np.array([ np.rot90(temp,-2).transpose([2,0,1]), mymat ]), axis=0 )
elif i == 5:
temp = predict(np.rot90(img,3), model, patch_sz=PATCH_SZ, n_classes=N_CLASSES)
# print(temp.transpose([2,0,1])[class_id][0][0], temp.transpose([2,0,1])[3][12][13])
# print("Case 6", temp.shape, mymat.shape)
mymat = np.mean( np.array([ np.rot90(temp, -3).transpose(2,0,1), mymat ]), axis=0 )
elif i ==6:
#added by me
temp = predict(img, model, patch_sz=PATCH_SZ, n_classes=N_CLASSES).transpose([2,0,1])
# print ("Case 7", temp.shape, mymat.shape)
else:
temp = predict(img, model, patch_sz=PATCH_SZ, n_classes=N_CLASSES).transpose([2,0,1])
# print(temp[class_id][0][0], temp[3][12][13])
# print("Case 8", temp.shape, mymat.shape)
mymat = np.mean( np.array([ temp, mymat ]), axis=0 )
print ("...prediction complete. Results Shape = ", results.shape)
#print(mymat[class_id][0][0], mymat[3][12][13])
print ("...creating map from result")
map = picture_from_mask(mymat, 0.5)
#mask = predict(img, model, patch_sz=PATCH_SZ, n_classes=N_CLASSES).transpose([2,0,1]) # make channels first
#map = picture_from_mask(mask, 0.5)
print ("...saving result for {}".format(test_filename))
tiff.imsave(result_path + test_filename + '_result.tif', (255*mymat).astype('uint8'))
print ("...saving map for {}".format(test_filename))
tiff.imsave(result_path + test_filename + '_map.tif', map)
print ("Prediction Completed for file {} : {} out of {}".format(test_filename, current_file_count+1, total_files_count))