trainer/vis.py

# 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)