Spaces:

akhaliq
/

deeplab2

Runtime error

deeplab2 / trainer /vis_utils.py

akhaliq3

spaces demo

506da10 almost 3 years ago

No virus

20.1 kB

	# coding=utf-8
	# Copyright 2021 The Deeplab2 Authors.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	"""Utility functions for the visualizer."""
	from absl import logging
	import matplotlib.pyplot as plt
	import numpy as np
	import PIL
	import tensorflow as tf

	from deeplab2.data import coco_constants

	# Amount of color perturbation added to colormap.
	_COLOR_PERTURBATION = 60


	def bit_get(val, idx):
	"""Gets the bit value.

	Args:
	val: Input value, int or numpy int array.
	idx: Which bit of the input val.

	Returns:
	The "idx"-th bit of input val.
	"""
	return (val >> idx) & 1


	def create_pascal_label_colormap():
	"""Creates a label colormap used in PASCAL VOC segmentation benchmark.

	Returns:
	A colormap for visualizing segmentation results.
	"""
	colormap = np.zeros((512, 3), dtype=int)
	ind = np.arange(512, dtype=int)

	for shift in reversed(list(range(8))):
	for channel in range(3):
	colormap[:, channel] \|= bit_get(ind, channel) << shift
	ind >>= 3

	return colormap


	def create_rgb_from_instance_map(instance_map):
	"""Creates an RGB image from an instance map for visualization.

	To assign a color to each instance, if the maximum value of the instance
	labels is smaller than the maximum allowed value of Pascal's colormap, we use
	Pascal's colormap. Otherwise, we use random and non-repeated colors.

	Args:
	instance_map: Numpy array of shape `[height, width]`, the instance map.

	Returns:
	instance_image: Numpy array of shape `[height, width, 3]`, the visualized
	RGB instance image.
	"""
	# pylint: disable=protected-access
	if np.max(instance_map) < 512:
	colormap = create_pascal_label_colormap()
	instance_image = colormap[instance_map]
	else:
	np.random.seed(0)

	used_colors = [(0, 0, 0)]
	instanc_map_shape = instance_map.shape
	instance_image = np.zeros([instanc_map_shape[0], instanc_map_shape[1], 3],
	np.uint8)
	instance_ids = np.unique(instance_map)
	for instance_id in instance_ids:
	# We preserve the id "0" for stuff.
	if instance_id == 0:
	continue
	r = np.random.randint(0, 256, dtype=np.uint8)
	g = np.random.randint(0, 256, dtype=np.uint8)
	b = np.random.randint(0, 256, dtype=np.uint8)
	while (r, g, b) in used_colors:
	r = np.random.randint(0, 256, dtype=np.uint8)
	g = np.random.randint(0, 256, dtype=np.uint8)
	b = np.random.randint(0, 256, dtype=np.uint8)
	instance_image[instance_map == instance_id, :] = (r, g, b)
	used_colors.append((r, g, b))
	instance_image[instance_map == 0, :] = (0, 0, 0)

	return instance_image


	def _generate_color(used_colors):
	""""Generates a non-repeated color.

	This function first uses the pascal colormap to generate the color. If more
	colors are requested, it randomly generates a non-repeated color.

	Args:
	used_colors: A list, where each element is a tuple in the format of
	(r, g, b).

	Returns:
	A tuple representing a color in the format of (r, g, b).
	A list, which is the updated `used_colors` with the returned color tuple
	appended to it.
	"""

	pascal_colormap = create_pascal_label_colormap()

	if len(used_colors) < len(pascal_colormap):
	color = tuple(pascal_colormap[len(used_colors)])
	else:
	r = np.random.randint(0, 256, dtype=np.uint8)
	g = np.random.randint(0, 256, dtype=np.uint8)
	b = np.random.randint(0, 256, dtype=np.uint8)
	while (r, g, b) in used_colors:
	r = np.random.randint(0, 256, dtype=np.uint8)
	g = np.random.randint(0, 256, dtype=np.uint8)
	b = np.random.randint(0, 256, dtype=np.uint8)
	color = (r, g, b)
	used_colors.append(color)

	return color, used_colors


	def overlay_heatmap_on_image(heatmap,
	input_image,
	dpi=80.0,
	add_color_bar=False):
	"""Overlays a heatmap on top of an image.

	Args:
	heatmap: A numpy array (float32) of shape `[height, width]`,
	which is the heatmap of keypoints.
	input_image: A numpy array (float32 or uint8) of shape
	`[height, width, 3]`, which is an image and all the pixel values are in
	the range of [0.0, 255.0].
	dpi: Float, the dpi of the output image.
	add_color_bar: Boolean, whether to add a colorbar to the output image.

	Returns:
	A numpy array (uint8) of the same shape as the `input_image`.
	"""

	# Generate the cmap.
	cmap = plt.cm.Reds
	# pylint: disable=protected-access
	cmap._init()
	# pylint: disable=protected-access
	cmap._lut[:, -1] = np.linspace(0, 1.0, 259)

	# Plot.
	image = input_image.astype(np.float32) / 255.0
	image_height, image_width, _ = image.shape
	fig, ax = plt.subplots(
	1,
	1,
	facecolor='white',
	figsize=(image_width / dpi, image_height / dpi),
	dpi=dpi)
	grid_y, grid_x = np.mgrid[0:image_height, 0:image_width]
	cb = ax.contourf(grid_x, grid_y, heatmap, 10, cmap=cmap)
	ax.imshow(image)
	ax.grid(False)
	plt.axis('off')
	if add_color_bar:
	plt.colorbar(cb)
	fig.subplots_adjust(bottom=0)
	fig.subplots_adjust(top=1)
	fig.subplots_adjust(right=1)
	fig.subplots_adjust(left=0)

	# Get the output image.
	fig.canvas.draw()
	# pylint: disable=protected-access
	output_image = np.array(fig.canvas.renderer._renderer)[:, :, :-1]
	plt.close()

	return output_image


	# pylint: disable=invalid-name
	def make_colorwheel():
	"""Generates a color wheel for optical flow visualization.

	Reference implementation:
	https://github.com/tomrunia/OpticalFlow_Visualization

	Returns:
	flow_image: A numpy array of output image.
	"""

	RY = 15
	YG = 6
	GC = 4
	CB = 11
	BM = 13
	MR = 6

	ncols = RY + YG + GC + CB + BM + MR
	colorwheel = np.zeros((ncols, 3))
	col = 0

	# RY
	colorwheel[0:RY, 0] = 255
	colorwheel[0:RY, 1] = np.floor(255 * np.arange(0, RY) / RY)
	col = col + RY
	# YG
	colorwheel[col:col + YG, 0] = 255 - np.floor(255 * np.arange(0, YG) / YG)
	colorwheel[col:col + YG, 1] = 255
	col = col + YG
	# GC
	colorwheel[col:col + GC, 1] = 255
	colorwheel[col:col + GC, 2] = np.floor(255 * np.arange(0, GC) / GC)
	col = col + GC
	# CB
	colorwheel[col:col+CB, 1] = 255 - np.floor(255*np.arange(CB)/CB)
	colorwheel[col:col+CB, 2] = 255
	col = col+CB
	# BM
	colorwheel[col:col + BM, 2] = 255
	colorwheel[col:col + BM, 0] = np.floor(255 * np.arange(0, BM) / BM)
	col = col + BM
	# MR
	colorwheel[col:col+MR, 2] = 255 - np.floor(255*np.arange(MR)/MR)
	colorwheel[col:col+MR, 0] = 255
	return colorwheel
	# pylint: enable=invalid-name


	def flow_compute_color(u, v):
	"""Computes color for 2D flow field.

	Reference implementation:
	https://github.com/tomrunia/OpticalFlow_Visualization

	Args:
	u: A numpy array of horizontal flow.
	v: A numpy array of vertical flow.

	Returns:
	flow_image: A numpy array of output image.
	"""

	flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)

	colorwheel = make_colorwheel() # shape [55x3]
	ncols = colorwheel.shape[0]

	rad = np.sqrt(np.square(u) + np.square(v))
	a = np.arctan2(-v, -u) / np.pi

	fk = (a + 1) / 2 * (ncols - 1)
	k0 = np.floor(fk).astype(np.int32)
	k1 = k0 + 1
	k1[k1 == ncols] = 0
	f = fk - k0

	for i in range(colorwheel.shape[1]):
	tmp = colorwheel[:, i]
	color0 = tmp[k0] / 255.0
	color1 = tmp[k1] / 255.0
	color = (1 - f) * color0 + f * color1

	idx = (rad <= 1)
	color[idx] = 1 - rad[idx] * (1 - color[idx])
	color[~idx] = color[~idx] * 0.75

	# The order is RGB.
	ch_idx = i
	flow_image[:, :, ch_idx] = np.floor(255 * color)

	return flow_image


	def flow_to_color(flow_uv, clip_flow=None):
	"""Applies color to 2D flow field.

	Reference implementation:
	https://github.com/tomrunia/OpticalFlow_Visualization

	Args:
	flow_uv: A numpy array of flow with shape [Height, Width, 2].
	clip_flow: A float to clip the maximum value for the flow.

	Returns:
	flow_image: A numpy array of output image.

	Raises:
	ValueError: Input flow does not have dimension equals to 3.
	ValueError: Input flow does not have shape [H, W, 2].
	"""

	if flow_uv.ndim != 3:
	raise ValueError('Input flow must have three dimensions.')
	if flow_uv.shape[2] != 2:
	raise ValueError('Input flow must have shape [H, W, 2].')

	if clip_flow is not None:
	flow_uv = np.clip(flow_uv, 0, clip_flow)

	u = flow_uv[:, :, 0]
	v = flow_uv[:, :, 1]

	rad = np.sqrt(np.square(u) + np.square(v))
	rad_max = np.max(rad)

	epsilon = 1e-5
	u = u / (rad_max + epsilon)
	v = v / (rad_max + epsilon)

	return flow_compute_color(u, v)


	def squeeze_batch_dim_and_convert_to_numpy(input_dict):
	for key in input_dict:
	input_dict[key] = tf.squeeze(input_dict[key], axis=0).numpy()
	return input_dict


	def create_cityscapes_label_colormap():
	"""Creates a label colormap used in CITYSCAPES segmentation benchmark.

	Returns:
	A colormap for visualizing segmentation results.
	"""
	colormap = np.zeros((256, 3), dtype=np.uint8)
	colormap[0] = [128, 64, 128]
	colormap[1] = [244, 35, 232]
	colormap[2] = [70, 70, 70]
	colormap[3] = [102, 102, 156]
	colormap[4] = [190, 153, 153]
	colormap[5] = [153, 153, 153]
	colormap[6] = [250, 170, 30]
	colormap[7] = [220, 220, 0]
	colormap[8] = [107, 142, 35]
	colormap[9] = [152, 251, 152]
	colormap[10] = [70, 130, 180]
	colormap[11] = [220, 20, 60]
	colormap[12] = [255, 0, 0]
	colormap[13] = [0, 0, 142]
	colormap[14] = [0, 0, 70]
	colormap[15] = [0, 60, 100]
	colormap[16] = [0, 80, 100]
	colormap[17] = [0, 0, 230]
	colormap[18] = [119, 11, 32]
	return colormap


	def create_motchallenge_label_colormap():
	"""Creates a label colormap used in MOTChallenge-STEP benchmark.

	Returns:
	A colormap for visualizing segmentation results.
	"""
	colormap = np.zeros((256, 3), dtype=np.uint8)
	colormap[0] = [244, 35, 232]
	colormap[1] = [70, 70, 70]
	colormap[2] = [107, 142, 35]
	colormap[3] = [70, 130, 180]
	colormap[4] = [220, 20, 60]
	colormap[5] = [255, 0, 0]
	colormap[6] = [119, 11, 32]
	return colormap


	def create_coco_label_colormap():
	"""Creates a label colormap used in COCO dataset.

	Returns:
	A colormap for visualizing segmentation results.
	"""
	# Obtain the dictionary mapping original category id to contiguous ones.
	coco_categories = coco_constants.get_coco_reduced_meta()
	colormap = np.zeros((256, 3), dtype=np.uint8)
	for category in coco_categories:
	colormap[category['id']] = category['color']
	return colormap


	def label_to_color_image(label, colormap_name='cityscapes'):
	"""Adds color defined by the colormap derived from the dataset to the label.

	Args:
	label: A 2D array with integer type, storing the segmentation label.
	colormap_name: A string specifying the name of the dataset. Used for
	choosing the right colormap. Currently supported: 'cityscapes',
	'motchallenge'. (Default: 'cityscapes')

	Returns:
	result: A 2D array with floating type. The element of the array
	is the color indexed by the corresponding element in the input label
	to the cityscapes colormap.

	Raises:
	ValueError: If label is not of rank 2 or its value is larger than color
	map maximum entry.
	"""
	if label.ndim != 2:
	raise ValueError('Expect 2-D input label. Got {}'.format(label.shape))

	if np.max(label) >= 256:
	raise ValueError(
	'label value too large: {} >= 256.'.format(np.max(label)))

	if colormap_name == 'cityscapes':
	colormap = create_cityscapes_label_colormap()
	elif colormap_name == 'motchallenge':
	colormap = create_motchallenge_label_colormap()
	elif colormap_name == 'coco':
	colormap = create_coco_label_colormap()
	else:
	raise ValueError('Could not find a colormap for dataset %s.' %
	colormap_name)
	return colormap[label]


	def save_parsing_result(parsing_result,
	label_divisor,
	thing_list,
	save_dir,
	filename,
	id_to_colormap=None,
	colormap_name='cityscapes'):
	"""Saves the parsing results.

	The parsing result encodes both semantic segmentation and instance
	segmentation results. In order to visualize the parsing result with only
	one png file, we adopt the following procedures, similar to the
	`visualization.py` provided in the COCO panoptic segmentation evaluation
	codes.

	1. Pixels predicted as `stuff` will take the same semantic color defined
	in the colormap.
	2. Pixels of a predicted `thing` instance will take similar semantic color
	defined in the colormap. For example, `car` class takes blue color in
	the colormap. Predicted car instance 1 will then be colored with the
	blue color perturbed with a small amount of RGB noise.

	Args:
	parsing_result: The numpy array to be saved. The data will be converted
	to uint8 and saved as png image.
	label_divisor: Integer, encoding the semantic segmentation and instance
	segmentation results as value = semantic_label * label_divisor +
	instance_label.
	thing_list: A list containing the semantic indices of the thing classes.
	save_dir: String, the directory to which the results will be saved.
	filename: String, the image filename.
	id_to_colormap: An optional mapping from track ID to color.
	colormap_name: A string specifying the dataset to choose the corresponding
	color map. Currently supported: 'cityscapes', 'motchallenge'. (Default:
	'cityscapes').

	Raises:
	ValueError: If parsing_result is not of rank 2 or its value in semantic
	segmentation result is larger than color map maximum entry.
	ValueError: If provided colormap_name is not supported.

	Returns:
	If id_to_colormap is passed, the updated id_to_colormap will be returned.
	"""
	if parsing_result.ndim != 2:
	raise ValueError('Expect 2-D parsing result. Got {}'.format(
	parsing_result.shape))
	semantic_result = parsing_result // label_divisor
	instance_result = parsing_result % label_divisor
	colormap_max_value = 256
	if np.max(semantic_result) >= colormap_max_value:
	raise ValueError('Predicted semantic value too large: {} >= {}.'.format(
	np.max(semantic_result), colormap_max_value))
	height, width = parsing_result.shape
	colored_output = np.zeros((height, width, 3), dtype=np.uint8)
	if colormap_name == 'cityscapes':
	colormap = create_cityscapes_label_colormap()
	elif colormap_name == 'motchallenge':
	colormap = create_motchallenge_label_colormap()
	elif colormap_name == 'coco':
	colormap = create_coco_label_colormap()
	else:
	raise ValueError('Could not find a colormap for dataset %s.' %
	colormap_name)
	# Keep track of used colors.
	used_colors = set()
	if id_to_colormap is not None:
	used_colors = set([tuple(val) for val in id_to_colormap.values()])
	np_state = None
	else:
	# Use random seed 0 in order to reproduce the same visualization.
	np_state = np.random.RandomState(0)

	unique_semantic_values = np.unique(semantic_result)
	for semantic_value in unique_semantic_values:
	semantic_mask = semantic_result == semantic_value
	if semantic_value in thing_list:
	# For `thing` class, we will add a small amount of random noise to its
	# correspondingly predefined semantic segmentation colormap.
	unique_instance_values = np.unique(instance_result[semantic_mask])
	for instance_value in unique_instance_values:
	instance_mask = np.logical_and(semantic_mask,
	instance_result == instance_value)
	if id_to_colormap is not None:
	if instance_value in id_to_colormap:
	colored_output[instance_mask] = id_to_colormap[instance_value]
	continue
	random_color = perturb_color(
	colormap[semantic_value],
	_COLOR_PERTURBATION,
	used_colors,
	random_state=np_state)
	colored_output[instance_mask] = random_color
	if id_to_colormap is not None:
	id_to_colormap[instance_value] = random_color
	else:
	# For `stuff` class, we use the defined semantic color.
	colored_output[semantic_mask] = colormap[semantic_value]
	used_colors.add(tuple(colormap[semantic_value]))

	pil_image = PIL.Image.fromarray(colored_output.astype(dtype=np.uint8))
	with tf.io.gfile.GFile('{}/{}.png'.format(save_dir, filename), mode='w') as f:
	pil_image.save(f, 'PNG')
	if id_to_colormap is not None:
	return id_to_colormap


	def perturb_color(color,
	noise,
	used_colors=None,
	max_trials=50,
	random_state=None):
	"""Pertrubs the color with some noise.

	If `used_colors` is not None, we will return the color that has
	not appeared before in it.

	Args:
	color: A numpy array with three elements [R, G, B].
	noise: Integer, specifying the amount of perturbing noise.
	used_colors: A set, used to keep track of used colors.
	max_trials: An integer, maximum trials to generate random color.
	random_state: An optional np.random.RandomState. If passed, will be used to
	generate random numbers.

	Returns:
	A perturbed color that has not appeared in used_colors.
	"""
	for _ in range(max_trials):
	if random_state is not None:
	random_color = color + random_state.randint(
	low=-noise, high=noise + 1, size=3)
	else:
	random_color = color + np.random.randint(low=-noise,
	high=noise+1,
	size=3)
	random_color = np.maximum(0, np.minimum(255, random_color))
	if used_colors is None:
	return random_color
	elif tuple(random_color) not in used_colors:
	used_colors.add(tuple(random_color))
	return random_color
	logging.warning('Using duplicate random color.')
	return random_color


	def save_annotation(label,
	save_dir,
	filename,
	add_colormap=True,
	normalize_to_unit_values=False,
	scale_values=False,
	colormap_name='cityscapes'):
	"""Saves the given label to image on disk.

	Args:
	label: The numpy array to be saved. The data will be converted
	to uint8 and saved as png image.
	save_dir: String, the directory to which the results will be saved.
	filename: String, the image filename.
	add_colormap: Boolean, add color map to the label or not.
	normalize_to_unit_values: Boolean, normalize the input values to [0, 1].
	scale_values: Boolean, scale the input values to [0, 255] for visualization.
	colormap_name: A string specifying the dataset to choose the corresponding
	color map. Currently supported: 'cityscapes', 'motchallenge'. (Default:
	'cityscapes').
	"""
	# Add colormap for visualizing the prediction.
	if add_colormap:
	colored_label = label_to_color_image(label, colormap_name)
	else:
	colored_label = label
	if normalize_to_unit_values:
	min_value = np.amin(colored_label)
	max_value = np.amax(colored_label)
	range_value = max_value - min_value
	if range_value != 0:
	colored_label = (colored_label - min_value) / range_value

	if scale_values:
	colored_label = 255. * colored_label

	pil_image = PIL.Image.fromarray(colored_label.astype(dtype=np.uint8))
	with tf.io.gfile.GFile('%s/%s.png' % (save_dir, filename), mode='w') as f:
	pil_image.save(f, 'PNG')