Spaces:

akhaliq
/

deeplab2

Runtime error

akhaliq3

spaces demo

506da10 almost 3 years ago

No virus

14.6 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.

	"""Visualizes and stores results of a panoptic-deeplab model."""
	import os.path
	from typing import Any, Dict, List, Text

	import numpy as np
	import tensorflow as tf

	from deeplab2 import common
	from deeplab2.data import coco_constants
	from deeplab2.data import dataset
	from deeplab2.trainer import vis_utils

	# The format of the labels.
	_IMAGE_FORMAT = '%06d_image'
	_CENTER_LABEL_FORMAT = '%06d_center_label'
	_OFFSET_LABEL_FORMAT = '%06d_offset_label'
	_PANOPTIC_LABEL_FORMAT = '%06d_panoptic_label'
	_SEMANTIC_LABEL_FORMAT = '%06d_semantic_label'

	# The format of the predictions.
	_INSTANCE_PREDICTION_FORMAT = '%06d_instance_prediction'
	_CENTER_HEATMAP_PREDICTION_FORMAT = '%06d_center_prediction'
	_OFFSET_PREDICTION_RGB_FORMAT = '%06d_offset_prediction_rgb'
	_PANOPTIC_PREDICTION_FORMAT = '%06d_panoptic_prediction'
	_SEMANTIC_PREDICTION_FORMAT = '%06d_semantic_prediction'

	# The format of others.
	_ANALYSIS_FORMAT = '%06d_semantic_error'

	# Conversion from train id to eval id.
	_CITYSCAPES_TRAIN_ID_TO_EVAL_ID = (
	7, 8, 11, 12, 13, 17, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 31, 32, 33, 0
	)
	_COCO_TRAIN_ID_TO_EVAL_ID = coco_constants.get_id_mapping_inverse()


	def _convert_train_id_to_eval_id(
	prediction: np.ndarray, dataset_name: str) -> np.ndarray:
	"""Converts the predicted label for evaluation.

	There are cases where the training labels are not equal to the evaluation
	labels. This function is used to perform the conversion so that we could
	evaluate the results on the evaluation server.

	Args:
	prediction: Semantic segmentation prediction.
	dataset_name: Dataset name.

	Returns:
	Semantic segmentation prediction whose labels have been changed.

	Raises:
	ValueError: If the dataset is not supported.
	"""
	if 'cityscapes' in dataset_name:
	train_id_to_eval_id = _CITYSCAPES_TRAIN_ID_TO_EVAL_ID
	elif 'coco' in dataset_name:
	train_id_to_eval_id = _COCO_TRAIN_ID_TO_EVAL_ID
	else:
	raise ValueError(
	'Unsupported dataset %s for converting semantic class IDs.' %
	dataset_name)
	length = np.maximum(256, len(train_id_to_eval_id))
	to_eval_id_map = np.zeros((length), dtype=prediction.dtype)
	dataset_ids = np.asarray(
	train_id_to_eval_id, dtype=prediction.dtype)
	to_eval_id_map[:len(train_id_to_eval_id)] = dataset_ids
	return to_eval_id_map[prediction]


	def _get_fg_mask(label_map: np.ndarray, thing_list: List[int]) -> np.ndarray:
	fg_mask = np.zeros_like(label_map, np.bool)
	for class_id in np.unique(label_map):
	if class_id in thing_list:
	fg_mask = np.logical_or(fg_mask, np.equal(label_map, class_id))
	fg_mask = np.expand_dims(fg_mask, axis=2)
	return fg_mask.astype(np.int)


	def store_raw_predictions(predictions: Dict[str, Any],
	image_filename: tf.Tensor,
	dataset_info: dataset.DatasetDescriptor,
	save_dir: Text,
	sequence: tf.Tensor,
	raw_panoptic_format='two_channel_png',
	convert_to_eval=True):
	"""Stores raw predictions to the specified path.

	Raw predictions are saved in the specified path with the specified
	`raw_panoptic_format`. For the `raw_panoptic_format`, we currently
	support `two_channel_png`, `three_channel_png` and `two_channel_numpy_array`.
	Note that `two_channel_png` and `three_channel_png` could not encode large
	values of semantic label and instance ID due to limited PNG channel size. In
	such a case, use `raw_panoptic_format` = `two_channel_numpy_array` to save
	the raw predictions as two channel numpy array (i.e., first channel encodes
	the semantic class and the second channel the instance ID).

	Args:
	predictions: A dctionary with string keys and any content. Tensors under
	common.PRED_SEMANTIC_KEY and common.PRED_PANOPTIC_KEY will be stored.
	image_filename: A tf.Tensor containing the image filename.
	dataset_info: A dataset.DatasetDescriptor specifying the dataset.
	save_dir: A path to the folder to write the output to.
	sequence: A tf.Tensor describing the sequence that the image belongs to.
	raw_panoptic_format: A string specifying what format the panoptic output
	should be stored. Supports:
	- 'two_channel_png': The popular format, also supported by the official
	COCO panoptic API (https://github.com/cocodataset/panopticapi), where
	the saved PNG image contains R-channel for semantic labels and
	G-channel for instance IDs.
	- 'three_channel_png': A simple extension of the 'two_channel_png' format,
	and is adopted in some video panoptic segmentation datasets (for
	example, KITTI-STEP and MOTChallenge-STEP), where the saved PNG image
	contains R-channel for semantic labels, G-channel for the values of
	(instance ID // 256), and B-channel for (instance ID % 256).
	- 'two_channel_numpy_array': A more flexible format (unconstrained by the
	PNG channel size), where the panoptic predictions are saved as a numpy
	array in the two channel format (i.e., first channel encodes the
	semantic class and the second channel the instance ID).
	convert_to_eval: A flag specyfing whether semantic class IDs should be
	converted to cityscapes eval IDs. This is usefulfor the official test
	sever evaluation.

	Raises:
	ValueError: An error occurs when semantic label or instance ID is larger
	than the values supported by the 'two_channel_png' or 'three_channel_png'
	format. Or, if the raw_panoptic_format is not supported.
	"""
	# Note: predictions[key] contains a tuple of length 1.
	predictions = {key: predictions[key][0] for key in predictions}
	predictions = vis_utils.squeeze_batch_dim_and_convert_to_numpy(predictions)
	image_filename = image_filename.numpy().decode('utf-8')
	image_filename = os.path.splitext(image_filename)[0]

	# Store raw semantic prediction.
	semantic_prediction = predictions[common.PRED_SEMANTIC_KEY]
	if convert_to_eval:
	semantic_prediction = _convert_train_id_to_eval_id(
	semantic_prediction, dataset_info.dataset_name)
	output_folder = os.path.join(save_dir, 'raw_semantic')
	if dataset_info.is_video_dataset:
	sequence = sequence.numpy().decode('utf-8')
	output_folder = os.path.join(output_folder, sequence)
	tf.io.gfile.makedirs(output_folder)
	vis_utils.save_annotation(
	semantic_prediction,
	output_folder,
	image_filename,
	add_colormap=False)

	if common.PRED_PANOPTIC_KEY in predictions:
	# Save the predicted panoptic annotations in two-channel format, where the
	# R-channel stores the semantic label while the G-channel stores the
	# instance label.
	panoptic_prediction = predictions[common.PRED_PANOPTIC_KEY]
	panoptic_outputs = np.zeros(
	(panoptic_prediction.shape[0], panoptic_prediction.shape[1], 3),
	dtype=panoptic_prediction.dtype)
	predicted_semantic_labels = (
	panoptic_prediction // dataset_info.panoptic_label_divisor)
	if convert_to_eval:
	predicted_semantic_labels = _convert_train_id_to_eval_id(
	predicted_semantic_labels, dataset_info.dataset_name)
	predicted_instance_labels = predictions[
	common.PRED_PANOPTIC_KEY] % dataset_info.panoptic_label_divisor

	output_folder = os.path.join(save_dir, 'raw_panoptic')
	if dataset_info.is_video_dataset:
	output_folder = os.path.join(output_folder, sequence)
	tf.io.gfile.makedirs(output_folder)
	if raw_panoptic_format == 'two_channel_png':
	if np.max(predicted_semantic_labels) > 255:
	raise ValueError('Overflow: Semantic IDs greater 255 are not supported '
	'for images of 8-bit. Please save output as numpy '
	'arrays instead.')
	if np.max(predicted_instance_labels) > 255:
	raise ValueError(
	'Overflow: Instance IDs greater 255 could not be encoded by '
	'G channel. Please save output as numpy arrays instead.')
	panoptic_outputs[:, :, 0] = predicted_semantic_labels
	panoptic_outputs[:, :, 1] = predicted_instance_labels
	vis_utils.save_annotation(panoptic_outputs,
	output_folder,
	image_filename,
	add_colormap=False)
	elif raw_panoptic_format == 'three_channel_png':
	if np.max(predicted_semantic_labels) > 255:
	raise ValueError('Overflow: Semantic IDs greater 255 are not supported '
	'for images of 8-bit. Please save output as numpy '
	'arrays instead.')
	if np.max(predicted_instance_labels) > 65535:
	raise ValueError(
	'Overflow: Instance IDs greater 65535 could not be encoded by '
	'G and B channels. Please save output as numpy arrays instead.')
	panoptic_outputs[:, :, 0] = predicted_semantic_labels
	panoptic_outputs[:, :, 1] = predicted_instance_labels // 256
	panoptic_outputs[:, :, 2] = predicted_instance_labels % 256
	vis_utils.save_annotation(panoptic_outputs,
	output_folder,
	image_filename,
	add_colormap=False)
	elif raw_panoptic_format == 'two_channel_numpy_array':
	panoptic_outputs[:, :, 0] = predicted_semantic_labels
	panoptic_outputs[:, :, 1] = predicted_instance_labels
	with tf.io.gfile.GFile(
	os.path.join(output_folder, image_filename + '.npy'), 'w') as f:
	np.save(f, panoptic_outputs)
	else:
	raise ValueError(
	'Unknown raw_panoptic_format %s.' % raw_panoptic_format)


	def store_predictions(predictions: Dict[str, Any], inputs: Dict[str, Any],
	image_id: int, dataset_info: dataset.DatasetDescriptor,
	save_dir: Text):
	"""Saves predictions and labels to the specified path."""
	predictions = {key: predictions[key][0] for key in predictions}
	predictions = vis_utils.squeeze_batch_dim_and_convert_to_numpy(predictions)
	inputs = {key: inputs[key][0] for key in inputs}
	del inputs[common.IMAGE_NAME]
	inputs = vis_utils.squeeze_batch_dim_and_convert_to_numpy(inputs)

	thing_list = dataset_info.class_has_instances_list
	label_divisor = dataset_info.panoptic_label_divisor
	colormap_name = dataset_info.colormap

	# 1. Save image.
	image = inputs[common.IMAGE]
	vis_utils.save_annotation(
	image,
	save_dir,
	_IMAGE_FORMAT % image_id,
	add_colormap=False)

	# 2. Save semantic predictions and semantic labels.
	vis_utils.save_annotation(
	predictions[common.PRED_SEMANTIC_KEY],
	save_dir,
	_SEMANTIC_PREDICTION_FORMAT % image_id,
	add_colormap=True,
	colormap_name=colormap_name)
	vis_utils.save_annotation(
	inputs[common.GT_SEMANTIC_RAW],
	save_dir,
	_SEMANTIC_LABEL_FORMAT % image_id,
	add_colormap=True,
	colormap_name=colormap_name)

	if common.PRED_CENTER_HEATMAP_KEY in predictions:
	# 3. Save center heatmap.
	heatmap_pred = predictions[common.PRED_CENTER_HEATMAP_KEY]
	heat_map_gt = inputs[common.GT_INSTANCE_CENTER_KEY]
	vis_utils.save_annotation(
	vis_utils.overlay_heatmap_on_image(
	heatmap_pred,
	image),
	save_dir,
	_CENTER_HEATMAP_PREDICTION_FORMAT % image_id,
	add_colormap=False)
	vis_utils.save_annotation(
	vis_utils.overlay_heatmap_on_image(
	heat_map_gt,
	image),
	save_dir,
	_CENTER_LABEL_FORMAT % image_id,
	add_colormap=False)

	if common.PRED_OFFSET_MAP_KEY in predictions:
	# 4. Save center offsets.
	center_offset_prediction = predictions[common.PRED_OFFSET_MAP_KEY]
	center_offset_prediction_rgb = vis_utils.flow_to_color(
	center_offset_prediction)
	semantic_prediction = predictions[common.PRED_SEMANTIC_KEY]
	pred_fg_mask = _get_fg_mask(semantic_prediction, thing_list)
	center_offset_prediction_rgb = (
	center_offset_prediction_rgb * pred_fg_mask)
	vis_utils.save_annotation(
	center_offset_prediction_rgb,
	save_dir,
	_OFFSET_PREDICTION_RGB_FORMAT % image_id,
	add_colormap=False)

	center_offset_label = inputs[common.GT_INSTANCE_REGRESSION_KEY]
	center_offset_label_rgb = vis_utils.flow_to_color(center_offset_label)
	gt_fg_mask = _get_fg_mask(inputs[common.GT_SEMANTIC_RAW], thing_list)
	center_offset_label_rgb = center_offset_label_rgb * gt_fg_mask

	vis_utils.save_annotation(
	center_offset_label_rgb,
	save_dir,
	_OFFSET_LABEL_FORMAT % image_id,
	add_colormap=False)

	if common.PRED_INSTANCE_KEY in predictions:
	# 5. Save instance map.
	vis_utils.save_annotation(
	vis_utils.create_rgb_from_instance_map(
	predictions[common.PRED_INSTANCE_KEY]),
	save_dir,
	_INSTANCE_PREDICTION_FORMAT % image_id,
	add_colormap=False)

	if common.PRED_PANOPTIC_KEY in predictions:
	# 6. Save panoptic segmentation.
	vis_utils.save_parsing_result(
	predictions[common.PRED_PANOPTIC_KEY],
	label_divisor=label_divisor,
	thing_list=thing_list,
	save_dir=save_dir,
	filename=_PANOPTIC_PREDICTION_FORMAT % image_id,
	colormap_name=colormap_name)
	vis_utils.save_parsing_result(
	parsing_result=inputs[common.GT_PANOPTIC_RAW],
	label_divisor=label_divisor,
	thing_list=thing_list,
	save_dir=save_dir,
	filename=_PANOPTIC_LABEL_FORMAT % image_id,
	colormap_name=colormap_name)

	# 7. Save error of semantic prediction.
	label = inputs[common.GT_SEMANTIC_RAW].astype(np.uint8)
	error_prediction = (
	(predictions[common.PRED_SEMANTIC_KEY] != label) &
	(label != dataset_info.ignore_label)).astype(np.uint8) * 255
	vis_utils.save_annotation(
	error_prediction,
	save_dir,
	_ANALYSIS_FORMAT % (image_id),
	add_colormap=False)