-
-
Notifications
You must be signed in to change notification settings - Fork 8
/
utils.py
212 lines (168 loc) · 8.33 KB
/
utils.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
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
import os
import tensorflow as tf
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from IPython.display import clear_output
IMAGE_SIZE = 224
BACKGROUND_LABEL = 0
BORDER_LABEL = 1
CONTENT_LABEL = 2
def display(display_list):
clear_output(wait=True)
plt.figure(figsize=(15, 15))
title = ['Input Image', 'True Mask', 'Predicted Mask']
for i in range(len(display_list)):
plt.subplot(1, len(display_list), i + 1)
plt.title(title[i])
plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
plt.axis('off')
plt.show()
def show_predictions_compared_to_real_data(images, true_masks, predictions):
for image_index in range(len(predictions)):
image = images[image_index]
true_mask = true_masks[image_index]
labeled_prediction = predictions[image_index]
display([image, true_mask, labeled_prediction])
def parse_image(img_path):
image = tf.io.read_file(img_path)
image = tf.image.decode_jpeg(image, channels=3)
image = tf.image.convert_image_dtype(image, tf.uint8)
mask_path = tf.strings.regex_replace(img_path, "raw", "segmentation_mask")
mask_path = tf.strings.regex_replace(mask_path, "jpg", "png")
mask = tf.io.read_file(mask_path)
mask = tf.image.decode_png(mask, channels=1)
# Transform mask colors into labels
# We will assume whites 0 which should be assigned to the background label
mask = tf.where(mask == 255, np.dtype('uint8').type(BACKGROUND_LABEL), mask)
# Dark values will use label the background label
mask = tf.where(mask == 29, np.dtype('uint8').type(BACKGROUND_LABEL), mask)
# Intermediate values will act as the border
mask = tf.where(mask == 76, np.dtype('uint8').type(BORDER_LABEL), mask)
mask = tf.where(mask == 134, np.dtype('uint8').type(BORDER_LABEL), mask)
# Brighter values will act as the content
mask = tf.where(mask == 149, np.dtype('uint8').type(CONTENT_LABEL), mask)
return {'image': image, 'segmentation_mask': mask}
def load_images_from_folder(folder, shuffle=True):
files = tf.data.Dataset.list_files(folder + "raw/*.jpg", shuffle=shuffle)
return files.map(parse_image)
def load_data_set():
return {
# Even when shuffle is recommended we don't want to shuffle the test dataset in order to be
# able to easily interpret the prediction result using the order as the index and assign the
# prediction index to the test image position inside the test folder.
'test': load_images_from_folder("./dataset/test/", shuffle=False),
'train': load_images_from_folder("./dataset/training/")
}
def normalize(input_image, input_mask):
input_image = tf.cast(input_image, tf.float32) / 255.0
return input_image, input_mask
def predicted_pixel_to_class(x):
return np.argmax(x)
# TODO: This makes us go really slow when checking accuracy.
# Review how we can make this faster in the future using arrays.
def map_prediction_to_mask(predicted_image):
predicted_mask = list()
for x in predicted_image:
predicted_mask_per_x = []
for y in x:
predicted_mask_per_x.append(predicted_pixel_to_class(y))
predicted_mask.append(predicted_mask_per_x)
return np.array(predicted_mask)
def tf_count(t, val):
elements_equal_to_value = tf.equal(t, val)
as_ints = tf.cast(elements_equal_to_value, tf.int32)
count = tf.reduce_sum(as_ints)
return count
@tf.function
def load_image_train(datapoint):
input_image = tf.image.resize_with_pad(datapoint['image'], target_height=IMAGE_SIZE, target_width=IMAGE_SIZE)
input_mask = tf.image.resize_with_pad(datapoint['segmentation_mask'], target_height=IMAGE_SIZE,
target_width=IMAGE_SIZE)
if tf.random.uniform(()) > 0.5:
input_image = tf.image.flip_left_right(input_image)
input_mask = tf.image.flip_left_right(input_mask)
input_image, input_mask = normalize(input_image, input_mask)
number_of_pixels_per_image = IMAGE_SIZE * IMAGE_SIZE
percentage_of_background_labels = tf_count(input_mask, BACKGROUND_LABEL) / number_of_pixels_per_image
percentage_of_content_labels = tf_count(input_mask, CONTENT_LABEL) / number_of_pixels_per_image
percentage_of_border_labels = tf_count(input_mask, BORDER_LABEL) / number_of_pixels_per_image
background_weight = tf.cast(0.33 / percentage_of_background_labels, tf.float32)
content_weight = tf.cast(0.34 / percentage_of_content_labels, tf.float32)
border_weight = tf.cast(0.33 / percentage_of_border_labels, tf.float32)
weights = tf.where(input_mask == BACKGROUND_LABEL, background_weight, input_mask)
weights = tf.where(input_mask == BORDER_LABEL, border_weight, weights)
weights = tf.where(input_mask == CONTENT_LABEL, content_weight, weights)
return input_image, input_mask, weights
def load_image_test(datapoint):
input_image = tf.image.resize_with_pad(datapoint['image'], target_height=IMAGE_SIZE, target_width=IMAGE_SIZE)
input_mask = tf.image.resize_with_pad(datapoint['segmentation_mask'], target_height=IMAGE_SIZE,
target_width=IMAGE_SIZE)
input_image, input_mask = normalize(input_image, input_mask)
return input_image, input_mask
def display(display_list):
plt.figure(figsize=(15, 15))
title = ['Input Image', 'True Mask', 'Predicted Mask']
for i in range(len(display_list)):
plt.subplot(1, len(display_list), i + 1)
plt.title(title[i])
plt.imshow(tf.keras.preprocessing.image.array_to_img(display_list[i]))
plt.axis('off')
plt.show()
def files_in_folder(folder):
return sorted([f for f in os.listdir(folder) if os.path.isfile(os.path.join(folder, f))])
def count_files_in_folder(folder):
return len(files_in_folder(folder))
# TODO: This is really slow. Should we avoid index access and use iterators instead?
def compare_accuracy_per_label(true_mask, predicted_mask):
total_background_labels = 0
total_border_labels = 0
total_content_labels = 0
properly_predicted_background_pixels = 0
properly_predicted_border_pixels = 0
properly_predicted_content_pixels = 0
for x in range(IMAGE_SIZE):
for y in range(IMAGE_SIZE):
prediction_label_per_pixel = int(predicted_mask[x][y])
mask_label_per_pixel = int(true_mask[x][y])
correct_prediction = prediction_label_per_pixel == mask_label_per_pixel
if mask_label_per_pixel == BACKGROUND_LABEL:
total_background_labels += 1
if correct_prediction:
properly_predicted_background_pixels += 1
if mask_label_per_pixel == BORDER_LABEL:
total_border_labels += 1
if correct_prediction:
properly_predicted_border_pixels += 1
if mask_label_per_pixel == CONTENT_LABEL:
total_content_labels += 1
if correct_prediction:
properly_predicted_content_pixels += 1
background_accuracy = properly_predicted_background_pixels / total_background_labels
border_accuracy = properly_predicted_border_pixels / total_border_labels
content_accuracy = properly_predicted_content_pixels / total_content_labels
return background_accuracy, border_accuracy, content_accuracy
def compare_accuracy(true_masks, predictions):
background_acc = 0.0
border_acc = 0.0
content_acc = 0.0
for index in range(len(predictions)):
print(f" - Checking accuracy for image with index {index}")
partial_back_acc, partial_border_acc, partial_content_acc = compare_accuracy_per_label(true_masks[index],
predictions[index])
background_acc += partial_back_acc
border_acc += partial_border_acc
content_acc += partial_content_acc
index += 1
pred_num = len(predictions)
return background_acc / pred_num, border_acc / pred_num, content_acc / pred_num
def label_to_rgb(labeled_pixel):
if labeled_pixel == BACKGROUND_LABEL:
return 0
if labeled_pixel == CONTENT_LABEL:
return 127
if labeled_pixel == BORDER_LABEL:
return 255
def labeled_prediction_to_image(predicted_result):
color_matrix = np.vectorize(label_to_rgb)(predicted_result)
return Image.fromarray(np.uint8(color_matrix))