-
Notifications
You must be signed in to change notification settings - Fork 0
/
cityscapes_utils.py
346 lines (288 loc) · 14.4 KB
/
cityscapes_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
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
"""
File name: cityscapes_utils.py
Author: Benjamin Planche
Date created: 14.02.2019
Date last modified: 14:49 14.02.2019
Python Version: 3.6
Copyright = "Copyright (C) 2018-2019 of Packt"
Credits = ["Eliot Andres, Benjamin Planche"]
License = "MIT"
Version = "1.0.0"
Maintainer = "non"
Status = "Prototype" # "Prototype", "Development", or "Production"
"""
#==============================================================================
# Imported Modules
#==============================================================================
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import tensorflow as tf
import cityscapesscripts.helpers.labels as cityscapes_labels
import glob
import numpy as np
import functools
#==============================================================================
# Constant Definitions
#==============================================================================
CITYSCAPES_FOLDER = os.path.expanduser('~/datasets/cityscapes')
CITYSCAPES_IGNORE_VALUE = 255
CITYSCAPES_LABELS = [label for label in cityscapes_labels.labels
if -1 < label.trainId < CITYSCAPES_IGNORE_VALUE]
CITYSCAPES_COLORS = np.asarray([label.color for label in CITYSCAPES_LABELS])
CITYSCAPES_COLORS_TF = tf.constant(CITYSCAPES_COLORS, dtype=tf.int32)
CITYSCAPES_IMG_RATIO = 2
CITYSCAPES_INT_FILL = 6
CITYSCAPES_FILE_TEMPLATE = os.path.join(
'{root}', '{type}', '{split}', '{city}',
'{city}_{seq:{filler}>{len_fill}}_{frame:{filler}>{len_fill}}_{type}{type2}{ext}')
#==============================================================================
# Function Definitions
#==============================================================================
# -----------------------------------------------------------------------------
# DATA FUNCTIONS
# -----------------------------------------------------------------------------
def get_cityscapes_file_pairs(split='train', city='*', sequence='*',
frame='*', ext='.*', gt_type='labelTrainIds', type='leftImg8bit',
root_folder=CITYSCAPES_FOLDER, file_template=CITYSCAPES_FILE_TEMPLATE):
"""
Fetch pairs of filenames for the Cityscapes dataset.
Note: wildcards accepted for the parameters (e.g. city='*' to return image pairs from every city)
:param split: Name of the split to return pairs from ("train", "val", ...)
:param city: Name of the city(ies)
:param sequence: Name of the video sequence(s)
:param frame: Name of the frame
:param ext: File extension
:param gt_type: Cityscapes GT type
:param type: Cityscapes image type
:param root_folder: Cityscapes root folder
:param file_template: File template to be applied (default corresponds to Cityscapes original format)
:return: List of input files, List of corresponding GT files
"""
input_file_template = file_template.format(
root=root_folder, type=type, type2='', len_fill=1, filler='*',
split=split, city=city, seq=sequence, frame=frame, ext=ext)
input_files = glob.glob(input_file_template)
gt_file_template = file_template.format(
root=root_folder, type='gtFine', type2='_'+gt_type, len_fill=1, filler='*',
split=split, city=city, seq=sequence, frame=frame, ext=ext)
gt_files = glob.glob(gt_file_template)
assert(len(input_files) == len(gt_files))
return sorted(input_files), sorted(gt_files)
def parse_function(filenames, resize_to=[226, 226], augment=True):
"""
Parse files into input/label image pair.
:param filenames: Dict containing the file(s) (filenames['image'], filenames['label'])
:param resize_to: H x W Dimensions to resize the image and label to
:param augment: Flag to augment the pair
:return: Input tensor, Label tensor
"""
img_filename, gt_filename = filenames['image'], filenames.get('label', None)
# Reading the file and returning its content as bytes:
image_string = tf.io.read_file(img_filename)
# Decoding into an image:
image_decoded = tf.io.decode_jpeg(image_string, channels=3)
# Converting image to float:
image = tf.image.convert_image_dtype(image_decoded, tf.float32)
# Resizing:
image = tf.image.resize(image, resize_to)
if gt_filename is not None:
# Same for GT image:
gt_string = tf.io.read_file(gt_filename)
gt_decoded = tf.io.decode_png(gt_string, channels=1)
gt = tf.cast(gt_decoded, dtype=tf.int32)
gt = tf.image.resize(gt, resize_to, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
# Opt. augmenting the image:
if augment:
image, gt = _augmentation_fn(image, gt)
return image, gt
else:
if augment:
image = _augmentation_fn(image)
return image
def _augmentation_fn(image, gt_image=None):
"""
Apply random transformations to augment the training images.
:param images: Images
:return: Augmented Images
"""
original_shape = tf.shape(image)[-3:-1]
num_image_channels = tf.shape(image)[-1]
# If we decide to randomly flip or resize/crop the image, the same should be applied to
# the label one so they still match. Therefore, to simplify the procedure, we stack the
# two images together along the channel axis, before these random operations:
if gt_image is None:
stacked_images = image
num_stacked_channels = num_image_channels
else:
stacked_images = tf.concat([image, tf.cast(gt_image, dtype=image.dtype)], axis=-1)
num_stacked_channels = tf.shape(stacked_images)[-1]
# Randomly applied horizontal flip:
stacked_images = tf.image.random_flip_left_right(stacked_images)
# Random cropping:
random_scale_factor = tf.random.uniform([], minval=.8, maxval=1., dtype=tf.float32)
crop_shape = tf.cast(tf.cast(original_shape, tf.float32) * random_scale_factor, tf.int32)
if len(stacked_images.shape) == 3: # single image:
crop_shape = tf.concat([crop_shape, [num_stacked_channels]], axis=0)
else: # batched images:
batch_size = tf.shape(stacked_images)[0]
crop_shape = tf.concat([[batch_size], crop_shape, [num_stacked_channels]], axis=0)
stacked_images = tf.image.random_crop(stacked_images, crop_shape)
# The remaining transformations should be applied either differently to the input and GT images
# (nearest-neighbor resizing for the label image VS interpolated resizing for the image),
# or only to the input image, not the GT one (color changes, etc.). Therefore, we split them back:
image = stacked_images[..., :num_image_channels]
# Resizing back to expected dimensions:
image = tf.image.resize(image, original_shape)
# Random B/S changes:
image = tf.image.random_brightness(image, max_delta=0.15)
image = tf.image.random_saturation(image, lower=0.5, upper=1.75)
image = tf.clip_by_value(image, 0.0, 1.0) # keeping pixel values in check
if gt_image is not None:
gt_image = tf.cast(stacked_images[..., num_image_channels:], dtype=gt_image.dtype)
gt_image = tf.image.resize(gt_image, original_shape, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
return image, gt_image
else:
return image
def segmentation_input_fn(image_files, gt_files=None, resize_to=[256, 256],
shuffle=False, batch_size=32, num_epochs=None, augment=False,
seed=None):
"""
Set up an input data pipeline for semantic segmentation applications.
:param image_files: List of input image files
:param gt_files: (opt.) List of corresponding label image files
:param resize_to: H x W Dimensions to resize the image and label to
:param shuffle: Flag to shuffle the dataset
:param batch_size: Batch size
:param num_epochs: Number of epochs the dataset would be iterated over
:param augment: Flag to augment the pairs
:param seed: (opt) Seed
:return: tf.data.Dataset
"""
# Converting to TF dataset:
image_files = tf.constant(image_files)
data_dict = {'image': image_files}
if gt_files is not None:
gt_files = tf.constant(gt_files)
data_dict['label'] = gt_files
dataset = tf.data.Dataset.from_tensor_slices(data_dict)
if shuffle:
dataset = dataset.shuffle(buffer_size=1000, seed=seed)
dataset = dataset.prefetch(1)
# Batching + adding parsing operation:
parse_fn = functools.partial(parse_function, resize_to=resize_to, augment=augment)
dataset = dataset.map(parse_fn, num_parallel_calls=4)
dataset = dataset.batch(batch_size)
dataset = dataset.repeat(num_epochs)
return dataset
def cityscapes_input_fn(split='train', root_folder=CITYSCAPES_FOLDER, resize_to=[256, 256],
shuffle=False, batch_size=32, num_epochs=None, augment=False,
seed=None, blurred=False):
"""
Set up an input data pipeline for semantic segmentation applications on Cityscapes dataset.
:param split: Split name ('train', 'val', 'test')
:param root_folder: Cityscapes root folder
:param resize_to: H x W Dimensions to resize the image and label to
:param shuffle: Flag to shuffle the dataset
:param batch_size: Batch size
:param num_epochs: Number of epochs the dataset would be iterated over
:param augment: Flag to augment the pairs
:param seed: (opt) Seed
:param blurred: Flag to use images with faces and immatriculation plates blurred
(for display)
:return: tf.data.Dataset
"""
type = "leftImg8bit_blurred" if blurred else "leftImg8bit"
input_files, gt_files = get_cityscapes_file_pairs(split=split, root_folder=root_folder, type=type)
return segmentation_input_fn(input_files, gt_files,
resize_to, shuffle, batch_size, num_epochs, augment, seed)
# -----------------------------------------------------------------------------
# DISPLAY FUNCTIONS
# -----------------------------------------------------------------------------
def change_ratio(image=None, pred=None, gt=None, ratio=CITYSCAPES_IMG_RATIO):
"""
Resze the images to the corresponding ratio.
:param image: (opt) Input image
:param pred: (opt) Predicted label image
:param gt: (opt) Target image
:param ratio: Ratio
:return: 3 resized images
"""
valid_input = image if image is not None else pred if pred is not None else gt
current_size = tf.shape(valid_input)[-3:-1]
width_with_ratio = tf.cast(tf.cast(current_size[1], tf.float32) * ratio, tf.int32)
size_with_ratio = tf.stack([current_size[0], width_with_ratio], axis=0)
if image is not None:
image = tf.image.resize(image, size_with_ratio)
if pred is not None:
pred = tf.image.resize(pred, size_with_ratio, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
if gt is not None:
gt = tf.image.resize(gt, size_with_ratio, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
return image, pred, gt
def convert_label_to_colors(label, one_hot=True, num_classes=len(CITYSCAPES_LABELS),
color_tensor=CITYSCAPES_COLORS_TF):
"""
Convert label images into color ones for display (for Tensors).
:param label: Label image (Tensor)
:param one_hot: Flag if the label image hasn't been one-hot yet and therefore should
:param num_classes: Number of classes (for one-hotting)
:param color_tensor: Tensor mapping labels to colors
:return: Color map
"""
label_shape = tf.shape(label)
color_channels = tf.shape(color_tensor)[-1]
if one_hot:
label = tf.one_hot(label, num_classes)
else:
label_shape = label_shape[:-1]
label = tf.reshape(tf.cast(label, tf.int32), (-1, num_classes))
colors = tf.matmul(label, color_tensor)
return tf.reshape(colors, tf.concat([label_shape, [color_channels]], axis=0))
def postprocess_to_show(image=None, pred=None, gt=None, one_hot=True, ratio=CITYSCAPES_IMG_RATIO):
"""
Post-process the training results of a segmentation model (as Tensors), for display.
:param image: (opt.) Input image tensor
:param pred: (opt.) Predicted label map tensor
:param gt: (opt.) Target label map tensor
:param one_hot: Flag if the predicted label image hasn't been one-hot yet and therefore should
:param ratio: Original image ratio
:return: Processed image tensor(s)
"""
out = []
image_show, pred_show, gt_show = change_ratio(image, pred, gt,
ratio)
if image is not None:
out.append(image_show)
if pred is not None:
if one_hot:
pred_show = tf.squeeze(pred_show, -1) # removing unnecessary channel dimension
pred_show = convert_label_to_colors(pred_show, one_hot=one_hot)
out.append(pred_show)
if gt is not None:
gt_show = tf.squeeze(gt_show, -1) # removing unnecessary channel dimension
gt_show = convert_label_to_colors(gt_show)
out.append(gt_show)
return out if len(out) > 1 else out[0]
def convert_labels_to_colors_numpy(label, one_hot=True, num_classes=len(CITYSCAPES_LABELS),
color_array=CITYSCAPES_COLORS, ignore_value=CITYSCAPES_IGNORE_VALUE):
"""
Convert label images into color ones for display (for numpy objects).
:param label: Label image (numpy array)
:param one_hot: Flag if the label image hasn't been one-hot yet and therefore should
:param num_classes: Number of classes (for one-hotting)
:param color_array: Array mapping labels to colors
:param ignore_value: Value of label to be ignored (for one-hotting)
:return: Color map
"""
if one_hot:
label_shape = label.shape
label = label.reshape(-1)
label[label == ignore_value] = num_classes
label = np.eye(num_classes + 1, dtype=np.int32)[label]
label = label[..., :num_classes]
else:
label_shape = label.shape[:-1]
label = label.reshape(-1, label.shape[-1])
colors = np.matmul(label, color_array)
return colors.reshape(list(label_shape) + [colors.shape[1]])