utils/panoptic_instances.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.

"""Tensorflow code for working with object instances in segmentation."""

from typing import Iterable
from typing import Optional
from typing import Tuple
from typing import Union

import tensorflow as tf


def instances_without_ignore_categories(panoptic_labels: tf.Tensor,
                                        ignore_categories: Union[tf.Tensor,
                                                                 Iterable[int]],
                                        panoptic_divisor: Union[tf.Tensor,
                                                                int] = 256):
  """Determines which instances to keep after ignoring a set of categories.

  Args:
    panoptic_labels: An integer tensor of panoptic labels of shape `[height,
      width]`. Each element will be `category * panoptic_divisor + instance` for
      a pixel.
    ignore_categories: An iterable or tensor of integer category labels.
      Instances where the category portion of the label in `panoptic_labels` are
      in the ignore set will not be included in the results.
    panoptic_divisor: The divisor used to multiply the category label when
      constructing panoptic labels, as in integer or integer scalar tensor.

  Returns:
    A boolean tensor masking which of the input `panoptic_labels` corresponds
    to an instance that will be kept, or equivalently *not* ignored.
  """
  ignore_categories = tf.convert_to_tensor(
      ignore_categories, dtype=panoptic_labels.dtype)
  panoptic_divisor = tf.convert_to_tensor(
      panoptic_divisor, dtype=panoptic_labels.dtype)

  instance_category = tf.math.floordiv(panoptic_labels, panoptic_divisor)
  instance_is_ignored = tf.math.reduce_any(
      tf.equal(
          tf.expand_dims(instance_category, 1),
          tf.expand_dims(ignore_categories, 0)),
      axis=1)
  instance_is_kept = tf.math.logical_not(instance_is_ignored)

  return instance_is_kept


def _broadcast_over_instances(t, num_instances):
  out_shape = tf.concat([tf.shape(t), [num_instances]], axis=0)
  return tf.broadcast_to(tf.expand_dims(t, -1), out_shape)


def instance_boxes_from_masks(
    panoptic_labels: tf.Tensor,
    ignore_categories: Optional[Union[tf.Tensor, Iterable[int]]] = None,
    panoptic_divisor: Union[tf.Tensor, int] = 256):
  """Finds the bounding boxes around instances, given a panoptic label map.

  Args:
    panoptic_labels: An integer tensor of panoptic labelsof shape `[height,
      width]`. Each element will be `category * panoptic_divisor + instance` for
      a pixel.
    ignore_categories: An iterable or tensor of integer category labels.
      Instances where the category portion of the label in `panoptic_labels` are
      in the ignore set will not be included in the results.
    panoptic_divisor: The divisor used to multiply the category label when
      constructing panoptic labels, as in integer or integer scalar tensor.

  Returns:
    A tuple of arrays (unique_labels, box_coords).
    unique_labels: An tensor of each possible non-ignored label value in
      `panoptic_labels`, in the same order as the boxes.
    box_coords: An tensor of shape `[num_labels, 4]`. Each row is one box as
      `[ymin, xmin, ymax, xmax]`.
  """
  label_shape = tf.shape(panoptic_labels)
  height = label_shape[0]
  width = label_shape[1]
  x_coord, y_coord = tf.meshgrid(
      tf.range(width, dtype=tf.float32), tf.range(height, dtype=tf.float32))

  unique_labels, flat_instance_index = tf.unique(
      tf.reshape(panoptic_labels, [height * width]))
  num_instances = tf.size(unique_labels)
  instance_index = tf.reshape(flat_instance_index, [height, width])

  y_coord_repeated = _broadcast_over_instances(y_coord, num_instances)
  x_coord_repeated = _broadcast_over_instances(x_coord, num_instances)
  instance_index_repeated = _broadcast_over_instances(instance_index,
                                                      num_instances)

  instance_index_matches = tf.math.equal(
      instance_index_repeated,
      tf.reshape(tf.range(num_instances), [1, 1, num_instances]))

  # In these tensors, each slice in the 3rd dimension corresponds to an
  # instance. We replace the pixels that do _not_ belong to that instance with
  # a +/- infinity in order that it not be included in the reduce_min/max below.
  inf3d = tf.broadcast_to([[[float('Inf')]]], tf.shape(x_coord_repeated))
  y_or_inf = tf.where(instance_index_matches, y_coord_repeated, inf3d)
  y_or_neg_inf = tf.where(instance_index_matches, y_coord_repeated, -inf3d)
  x_or_inf = tf.where(instance_index_matches, x_coord_repeated, inf3d)
  x_or_neg_inf = tf.where(instance_index_matches, x_coord_repeated, -inf3d)

  y_min = tf.reduce_min(y_or_inf, axis=[0, 1])
  x_min = tf.reduce_min(x_or_inf, axis=[0, 1])
  y_max = tf.reduce_max(y_or_neg_inf, axis=[0, 1]) + 1
  x_max = tf.reduce_max(x_or_neg_inf, axis=[0, 1]) + 1

  box_coords = tf.stack([y_min, x_min, y_max, x_max], axis=1)

  if ignore_categories is not None:
    # Filter out the boxes that correspond to instances in the "ignore"
    # categories.
    instance_is_kept = instances_without_ignore_categories(
        unique_labels, ignore_categories, panoptic_divisor)

    unique_labels = tf.boolean_mask(unique_labels, instance_is_kept)
    box_coords = tf.boolean_mask(box_coords, instance_is_kept)

  return unique_labels, box_coords


def per_instance_masks(panoptic_labels: tf.Tensor,
                       instance_panoptic_labels: tf.Tensor,
                       out_dtype: tf.dtypes.DType = tf.bool) -> tf.Tensor:
  """3D tensor where each slice in 3rd dimensions is an instance's mask."""
  num_instances = tf.size(instance_panoptic_labels)
  matches = tf.equal(
      tf.expand_dims(panoptic_labels, 0),
      tf.reshape(instance_panoptic_labels, [num_instances, 1, 1]))
  return tf.cast(matches, out_dtype)


def _average_per_instance(map_tensor: tf.Tensor, panoptic_labels: tf.Tensor,
                          instance_panoptic_labels: tf.Tensor,
                          instance_area: tf.Tensor) -> tf.Tensor:
  """Finds the average of the values in map_tensor over each instance."""

  # For each instance (in the 3rd dim), generate a map that has, per-pixel:
  # - The input value if that pixel belongs to the instance.
  # - Zero otherwise.
  pixel_in_instance = per_instance_masks(panoptic_labels,
                                         instance_panoptic_labels)

  map_dtype = map_tensor.dtype
  num_instances = tf.size(instance_panoptic_labels)
  map_or_zero = tf.where(pixel_in_instance, tf.expand_dims(map_tensor, 0),
                         tf.zeros([num_instances, 1, 1], dtype=map_dtype))

  # Average the semantic probabilities over each instance.
  instance_total_prob = tf.math.reduce_sum(map_or_zero, axis=[1, 2])
  instance_avg_prob = tf.divide(instance_total_prob,
                                tf.cast(instance_area, map_dtype))

  return instance_avg_prob


# pyformat: disable
def per_instance_semantic_probabilities(
    panoptic_labels: tf.Tensor,
    instance_panoptic_labels: tf.Tensor,
    instance_area: tf.Tensor,
    semantic_probability: tf.Tensor,
    panoptic_divisor: Union[tf.Tensor, int],
    ignore_label: Union[tf.Tensor, int]) -> tf.Tensor:
  """Mean probability for the semantic label of each unique instance."""
  # pyformat: enable
  # Get the probability associated with the semantic label encoded in the
  # panoptic_labels at each pixel.
  panoptic_divisor = tf.convert_to_tensor(panoptic_divisor, dtype=tf.int32)
  ignore_label = tf.convert_to_tensor(ignore_label, dtype=tf.int32)
  semantic_label_map = tf.math.floordiv(panoptic_labels, panoptic_divisor)

  map_shape = tf.shape(semantic_label_map)
  height = map_shape[0]
  width = map_shape[1]
  num_pixels = height * width

  semantic_index = tf.reshape(semantic_label_map, [num_pixels])
  # Use 0 as the index for a pixel with the "ignore" label, since that semantic
  # label may not be a valid index into the class axis of the
  # semantic_probability tensor.
  semantic_index = tf.where(semantic_index == ignore_label, 0, semantic_index)

  x, y = tf.meshgrid(tf.range(width), tf.range(height))
  probability_index = tf.stack([
      tf.reshape(y, [num_pixels]),
      tf.reshape(x, [num_pixels]),
      semantic_index,
  ],
                               axis=1)

  pixel_semantic_probability = tf.reshape(
      tf.gather_nd(semantic_probability, probability_index), [height, width])
  # Set the probability for the "ignore" pixels to 0.
  pixel_semantic_probability = tf.where(semantic_label_map == ignore_label, 0.0,
                                        pixel_semantic_probability)

  instance_avg_prob = _average_per_instance(pixel_semantic_probability,
                                            panoptic_labels,
                                            instance_panoptic_labels,
                                            instance_area)

  return instance_avg_prob


def combined_instance_scores(
    panoptic_labels: tf.Tensor, semantic_probability: tf.Tensor,
    instance_score_map: tf.Tensor, panoptic_divisor: Union[tf.Tensor, int],
    ignore_label: Union[tf.Tensor, int]) -> Tuple[tf.Tensor, tf.Tensor]:
  """Combines (with a product) predicted semantic and instance probabilities.

  Args:
    panoptic_labels: A 2D integer tensor of panoptic format labels (each pixel
      entry is `semantic_label * panoptic_divisor + instance_label`).
    semantic_probability: A 3D float tensor, where the 3rd dimension is over
      semantic labels, and each spatial location will have the discrete
      distribution of the probabilities of the semantic classes.
    instance_score_map: A 2D float tensor, where the pixels for an instance will
      have the probability of that being an instance.
    panoptic_divisor: Integer scalar divisor/multiplier used to construct the
      panoptic labels.
    ignore_label: Integer scalar, for the "ignore" semantic label in the
      panoptic labels.

  Returns:
    A tuple of instance labels and the combined scores for those instances, each
    as a 1D tensor.
  """
  panoptic_divisor = tf.convert_to_tensor(panoptic_divisor, dtype=tf.int32)
  ignore_label = tf.convert_to_tensor(ignore_label, dtype=tf.int32)

  num_pixels = tf.size(panoptic_labels)
  instance_panoptic_labels, _, instance_area = tf.unique_with_counts(
      tf.reshape(panoptic_labels, [num_pixels]))

  instance_semantic_labels = tf.math.floordiv(instance_panoptic_labels,
                                              panoptic_divisor)
  instance_panoptic_labels = tf.boolean_mask(
      instance_panoptic_labels, instance_semantic_labels != ignore_label)
  instance_area = tf.boolean_mask(instance_area,
                                  instance_semantic_labels != ignore_label)

  instance_semantic_probabilities = per_instance_semantic_probabilities(
      panoptic_labels, instance_panoptic_labels, instance_area,
      semantic_probability, panoptic_divisor, ignore_label)

  instance_scores = _average_per_instance(instance_score_map, panoptic_labels,
                                          instance_panoptic_labels,
                                          instance_area)

  combined_scores = instance_semantic_probabilities * instance_scores
  return instance_panoptic_labels, combined_scores


def per_instance_is_crowd(is_crowd_map: tf.Tensor, id_map: tf.Tensor,
                          output_ids: tf.Tensor) -> tf.Tensor:
  """Determines the per-instance is_crowd value from a boolian is_crowd map.

  Args:
    is_crowd_map: A 2D boolean tensor. Where it is True, the instance in that
      region is a "crowd" instance. It is assumed that all pixels in an instance
      will have the same value in this map.
    id_map: A 2D integer tensor, with the instance id label at each pixel.
    output_ids: A 1D integer vector tensor, the per-instance ids for which to
      output the is_crowd values.

  Returns:
    A 1D boolean vector tensor, with the per-instance is_crowd value. The ith
    element of the return value will be the is_crowd result for the segment
    with the ith element of the output_ids argument.
  """
  flat_is_crowd_map = tf.reshape(is_crowd_map, [-1])
  flat_id_map = tf.reshape(id_map, [-1])

  # Get an is_crowd value from the map for each id.
  # Only need an arbtitrary value due to assumption that the is_crowd map does
  # not vary over an instance.
  unique_ids, unique_index = tf.unique(flat_id_map)
  unique_is_crowd = tf.scatter_nd(
      tf.expand_dims(unique_index, 1), flat_is_crowd_map, tf.shape(unique_ids))

  # Map from the order/set in unique_ids to that in output_ids
  matching_ids = tf.math.equal(
      tf.expand_dims(output_ids, 1), tf.expand_dims(unique_ids, 0))
  matching_index = tf.where(matching_ids)[:, 1]
  return tf.gather(unique_is_crowd, matching_index)