def segmentation_for_folder(input_folder, output_folder, color=None):
""" Draws segmentation images for a given folder of labels
Parameters
----------
input_folder: str
path with json files / labels
output_folder: str
folder to store segmentation images, cannot exist
color: int, gray color value
(int, int, int), BGR values
None for default colors
Returns nothing
"""
# TODO Add color image option
# TODO keep input name and folders
if os.path.exists(output_folder):
raise IOError(
'Output folder already exists, stopping to not mess things up')
os.makedirs(output_folder)
input_labels = helper_scripts.get_files_from_folder(input_folder, '.json')
for i, label_path in tqdm.tqdm(enumerate(input_labels)):
segmentation_image = visualize_labels.create_segmentation_image(
label_path, image='gray', color=color)
cv2.imwrite(os.path.join(output_folder,
str(i) + '.png'), segmentation_image)
def create_multi_class_segmentation_label(json_path):
""" Creates pixel-level label of markings color coded by lane association
Only for the for closest lane dividers, i.e. l1, l0, r0, r1
Parameters
----------
json_path: str
path to label file
Returns
-------
numpy.array
pixel level segmentation with lane association (717, 1276, 5)
Notes
-----
Only draws 4 classes, can easily be extended for to a given number of lanes
"""
debug_image = visualize_labels.create_segmentation_image(json_path, image='blank')
l1 = (debug_image == dc.DICT_COLORS['l1']).all(axis=2).astype(numpy.uint8)
l0 = (debug_image == dc.DICT_COLORS['l0']).all(axis=2).astype(numpy.uint8)
r0 = (debug_image == dc.DICT_COLORS['r0']).all(axis=2).astype(numpy.uint8)
r1 = (debug_image == dc.DICT_COLORS['r1']).all(axis=2).astype(numpy.uint8)
no_marker = (l1 + l0 + r0 + r1) == 0
return numpy.stack((no_marker, l1, l0, r0, r1), axis=2)
def create_binary_segmentation_label(json_path):
""" Creates binary segmentation image from label
Parameters
----------
json_path: str
path to label file
Returns
-------
numpy.array
binary image, 0 for background or 1 for marker, (716, 1276), numpy.uint8
"""
blank_image = numpy.zeros((717, 1276), dtype=numpy.uint8)
blank_image = visualize_labels.create_segmentation_image(
json_path, color=1, image=blank_image)
return blank_image
def create_deeplab_tfrecords(input_folder, tfrecord_file):
"""Creates a tfrecord file for a given folder
Parameters:
input_folder: str, path to samples for a given dataset
tfrecord_file: str, path to tfrecord that will be created
Flags:
See docstring for more information
color_input: whether to use gray or color images
multi_class: binary or multi-class segmentation
location_gradients: location information as extra channels
"""
label_paths = helper_scripts.get_files_from_folder(input_folder, '.json')
shuffle(label_paths)
print('{} label files in {}'.format(len(label_paths), input_folder))
loc_grad_x = list(
map(lambda z: z / constants.IMAGE_WIDTH * 255,
range(constants.IMAGE_WIDTH)))
loc_grad_y = list(
map(lambda z: z / constants.IMAGE_HEIGHT * 255,
range(constants.IMAGE_HEIGHT)))
loc_grad_x = numpy.asarray([loc_grad_x] * constants.IMAGE_HEIGHT)
loc_grad_y = numpy.asarray([loc_grad_y] *
constants.IMAGE_WIDTH).transpose()
loc_grad_x = numpy.round(loc_grad_x).astype(numpy.uint8)
loc_grad_y = numpy.round(loc_grad_y).astype(numpy.uint8)
os.makedirs(os.path.dirname(tfrecord_file), exist_ok=True)
with tf.python_io.TFRecordWriter(tfrecord_file) as writer:
for label_path in tqdm.tqdm(label_paths,
total=len(label_paths),
desc='Creating ' + tfrecord_file):
image_name = os.path.basename(label_path).replace('.json', '')
if FLAGS.color_input:
image_data = label_file_scripts.read_image(label_path,
image_type='color')
else:
image_data = label_file_scripts.read_image(label_path,
image_type='gray')
if FLAGS.location_gradients:
image_data = numpy.stack(
[image_data, loc_grad_x, loc_grad_y], -1)
image_data = cv2.imencode('.png', image_data)[1].tostring()
if FLAGS.multi_class:
segmentation_label = segmentation_labels.create_multi_class_segmentation_label(
label_path)
segmentation = numpy.zeros(segmentation_label.shape[0:2],
numpy.uint8)
for class_index in range(1, 5):
segmentation[segmentation_label[:, :, class_index] >
0] = class_index
else:
segmentation = visualize_labels.create_segmentation_image(
label_path, image='blank')
segmentation = cv2.cvtColor(segmentation, cv2.COLOR_BGR2GRAY)
segmentation = segmentation > 0
segmentation = segmentation.astype(numpy.uint8)
segmentation = cv2.imencode('.png', segmentation)[1].tostring()
example = build_data.image_seg_to_tfexample(
image_data, image_name, constants.IMAGE_HEIGHT,
constants.IMAGE_WIDTH, segmentation)
writer.write(example.SerializeToString())