karolmajek/maxdeeplab

1# coding=utf-82# Copyright 2021 The Deeplab2 Authors.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16"""This file contains functions to post-process Panoptic-DeepLab results."""17 18import functools19from typing import Tuple, Dict, Text20 21import tensorflow as tf22 23from deeplab2 import common24from deeplab2 import config_pb225from deeplab2.data import dataset26from deeplab2.model import utils27from deeplab2.tensorflow_ops.python.ops import merge_semantic_and_instance_maps_op as merge_ops28 29 30def _get_semantic_predictions(semantic_logits: tf.Tensor) -> tf.Tensor:31  """Computes the semantic classes from the predictions.32 33  Args:34    semantic_logits: A tf.tensor of shape [batch, height, width, classes].35 36  Returns:37    A tf.Tensor containing the semantic class prediction of shape38      [batch, height, width].39  """40  return tf.argmax(semantic_logits, axis=-1, output_type=tf.int32)41 42 43def _get_instance_centers_from_heatmap(44    center_heatmap: tf.Tensor, center_threshold: float, nms_kernel_size: int,45    keep_k_centers: int) -> Tuple[tf.Tensor, tf.Tensor]:46  """Computes a list of instance centers.47 48  Args:49    center_heatmap: A tf.Tensor of shape [height, width, 1].50    center_threshold: A float setting the threshold for the center heatmap.51    nms_kernel_size: An integer specifying the nms kernel size.52    keep_k_centers: An integer specifying the number of centers to keep (K).53      Non-positive values will keep all centers.54 55  Returns:56    A tuple of57    - tf.Tensor of shape [N, 2] containing N center coordinates (after58      non-maximum suppression) in (y, x) order.59    - tf.Tensor of shape [height, width] containing the center heatmap after60      non-maximum suppression.61  """62  # Threshold center map.63  center_heatmap = tf.where(64      tf.greater(center_heatmap, center_threshold), center_heatmap, 0.0)65 66  # Non-maximum suppression.67  padded_map = utils.add_zero_padding(center_heatmap, nms_kernel_size, rank=3)68  pooled_center_heatmap = tf.keras.backend.pool2d(69      tf.expand_dims(padded_map, 0),70      pool_size=(nms_kernel_size, nms_kernel_size),71      strides=(1, 1),72      padding='valid',73      pool_mode='max')74  center_heatmap = tf.where(75      tf.equal(pooled_center_heatmap, center_heatmap), center_heatmap, 0.0)76  center_heatmap = tf.squeeze(center_heatmap, axis=[0, 3])77 78  # `centers` is of shape (N, 2) with (y, x) order of the second dimension.79  centers = tf.where(tf.greater(center_heatmap, 0.0))80 81  if keep_k_centers > 0 and tf.shape(centers)[0] > keep_k_centers:82    topk_scores, _ = tf.math.top_k(83        tf.reshape(center_heatmap, [-1]), keep_k_centers, sorted=False)84    centers = tf.where(tf.greater(center_heatmap, topk_scores[-1]))85 86  return centers, center_heatmap87 88 89def _find_closest_center_per_pixel(centers: tf.Tensor,90                                   center_offsets: tf.Tensor) -> tf.Tensor:91  """Assigns all pixels to their closest center.92 93  Args:94    centers: A tf.Tensor of shape [N, 2] containing N centers with coordinate95      order (y, x).96    center_offsets: A tf.Tensor of shape [height, width, 2].97 98  Returns:99    A tf.Tensor of shape [height, width] containing the index of the closest100      center, per pixel.101  """102  height = tf.shape(center_offsets)[0]103  width = tf.shape(center_offsets)[1]104 105  x_coord, y_coord = tf.meshgrid(tf.range(width), tf.range(height))106  coord = tf.stack([y_coord, x_coord], axis=-1)107 108  center_per_pixel = tf.cast(coord, tf.float32) + center_offsets109 110  # centers: [N, 2] -> [N, 1, 2].111  # center_per_pixel: [H, W, 2] -> [1, H*W, 2].112  centers = tf.cast(tf.expand_dims(centers, 1), tf.float32)113  center_per_pixel = tf.reshape(center_per_pixel, [height*width, 2])114  center_per_pixel = tf.expand_dims(center_per_pixel, 0)115 116  # distances: [N, H*W].117  distances = tf.norm(centers - center_per_pixel, axis=-1)118 119  return tf.reshape(tf.argmin(distances, axis=0), [height, width])120 121 122def _get_instances_from_heatmap_and_offset(123    semantic_segmentation: tf.Tensor, center_heatmap: tf.Tensor,124    center_offsets: tf.Tensor, center_threshold: float,125    thing_class_ids: tf.Tensor, nms_kernel_size: int,126    keep_k_centers: int) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:127  """Computes the instance assignment per pixel.128 129  Args:130    semantic_segmentation: A tf.Tensor containing the semantic labels of shape131      [height, width].132    center_heatmap: A tf.Tensor of shape [height, width, 1].133    center_offsets: A tf.Tensor of shape [height, width, 2].134    center_threshold: A float setting the threshold for the center heatmap.135    thing_class_ids: A tf.Tensor of shape [N] containing N thing indices.136    nms_kernel_size: An integer specifying the nms kernel size.137    keep_k_centers: An integer specifying the number of centers to keep.138      Negative values will keep all centers.139 140  Returns:141    A tuple of:142    - tf.Tensor containing the instance segmentation (filtered with the `thing`143      segmentation from the semantic segmentation output) with shape144      [height, width].145    - tf.Tensor containing the processed centermap with shape [height, width].146    - tf.Tensor containing instance scores (where higher "score" is a reasonable147      signal of a higher confidence detection.) Will be of shape [height, width]148      with the score for a pixel being the score of the instance it belongs to.149      The scores will be zero for pixels in background/"stuff" regions.150  """151  thing_segmentation = tf.zeros_like(semantic_segmentation)152  for thing_id in thing_class_ids:153    thing_segmentation = tf.where(tf.equal(semantic_segmentation, thing_id),154                                  1,155                                  thing_segmentation)156 157  centers, processed_center_heatmap = _get_instance_centers_from_heatmap(158      center_heatmap, center_threshold, nms_kernel_size, keep_k_centers)159  if tf.shape(centers)[0] == 0:160    return (tf.zeros_like(semantic_segmentation), processed_center_heatmap,161            tf.zeros_like(processed_center_heatmap))162 163  instance_center_index = _find_closest_center_per_pixel(164      centers, center_offsets)165  # Instance IDs should start with 1. So we use the index into the centers, but166  # shifted by 1.167  instance_segmentation = tf.cast(instance_center_index, tf.int32) + 1168 169  # The value of the heatmap at an instance's center is used as the score170  # for that instance.171  instance_scores = tf.gather_nd(processed_center_heatmap, centers)172  tf.debugging.assert_shapes([173      (centers, ('N', 2)),174      (instance_scores, ('N',)),175  ])176  # This will map the instance scores back to the image space: where each pixel177  # has a value equal to the score of its instance.178  flat_center_index = tf.reshape(instance_center_index, [-1])179  instance_score_map = tf.gather(instance_scores, flat_center_index)180  instance_score_map = tf.reshape(instance_score_map,181                                  tf.shape(instance_segmentation))182  instance_score_map *= tf.cast(thing_segmentation, tf.float32)183 184  return (thing_segmentation * instance_segmentation, processed_center_heatmap,185          instance_score_map)186 187 188@tf.function189def _get_panoptic_predictions(190    semantic_logits: tf.Tensor, center_heatmap: tf.Tensor,191    center_offsets: tf.Tensor, center_threshold: float,192    thing_class_ids: tf.Tensor, label_divisor: int, stuff_area_limit: int,193    void_label: int, nms_kernel_size: int, keep_k_centers: int,194    merge_semantic_and_instance_with_tf_op: bool195) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:196  """Computes the semantic class and instance ID per pixel.197 198  Args:199    semantic_logits: A tf.Tensor of shape [batch, height, width, classes].200    center_heatmap: A tf.Tensor of shape [batch, height, width, 1].201    center_offsets: A tf.Tensor of shape [batch, height, width, 2].202    center_threshold: A float setting the threshold for the center heatmap.203    thing_class_ids: A tf.Tensor of shape [N] containing N thing indices.204    label_divisor: An integer specifying the label divisor of the dataset.205    stuff_area_limit: An integer specifying the number of pixels that stuff206      regions need to have at least. The stuff region will be included in the207      panoptic prediction, only if its area is larger than the limit; otherwise,208      it will be re-assigned as void_label.209    void_label: An integer specifying the void label.210    nms_kernel_size: An integer specifying the nms kernel size.211    keep_k_centers: An integer specifying the number of centers to keep.212      Negative values will keep all centers.213    merge_semantic_and_instance_with_tf_op: Boolean, specifying the merging214      operation uses TensorFlow (CUDA kernel) implementation (True) or215      tf.py_function implementation (False). Note the tf.py_function216      implementation is simply used as a backup solution when you could not217      successfully compile the provided TensorFlow implementation. To reproduce218      our results, please use the provided TensorFlow implementation `merge_ops`219      (i.e., set to True).220 221  Returns:222    A tuple of:223    - the panoptic prediction as tf.Tensor with shape [batch, height, width].224    - the semantic prediction as tf.Tensor with shape [batch, height, width].225    - the instance prediction as tf.Tensor with shape [batch, height, width].226    - the centermap prediction as tf.Tensor with shape [batch, height, width].227    - the instance score maps as tf.Tensor with shape [batch, height, width].228  """229  semantic_prediction = _get_semantic_predictions(semantic_logits)230  batch_size = tf.shape(semantic_logits)[0]231 232  instance_map_lists = tf.TensorArray(233      tf.int32, size=batch_size, dynamic_size=False)234  center_map_lists = tf.TensorArray(235      tf.float32, size=batch_size, dynamic_size=False)236  instance_score_map_lists = tf.TensorArray(237      tf.float32, size=batch_size, dynamic_size=False)238 239  for i in tf.range(batch_size):240    (instance_map, center_map,241     instance_score_map) = _get_instances_from_heatmap_and_offset(242         semantic_prediction[i, ...], center_heatmap[i, ...],243         center_offsets[i, ...], center_threshold, thing_class_ids,244         nms_kernel_size, keep_k_centers)245    instance_map_lists = instance_map_lists.write(i, instance_map)246    center_map_lists = center_map_lists.write(i, center_map)247    instance_score_map_lists = instance_score_map_lists.write(248        i, instance_score_map)249 250  # This does not work with unknown shapes.251  instance_maps = instance_map_lists.stack()252  center_maps = center_map_lists.stack()253  instance_score_maps = instance_score_map_lists.stack()254 255  if merge_semantic_and_instance_with_tf_op:256    panoptic_prediction = merge_ops.merge_semantic_and_instance_maps(257        semantic_prediction, instance_maps, thing_class_ids, label_divisor,258        stuff_area_limit, void_label)259  else:260    panoptic_prediction = _merge_semantic_and_instance_maps(261        semantic_prediction, instance_maps, thing_class_ids, label_divisor,262        stuff_area_limit, void_label)263  return (panoptic_prediction, semantic_prediction, instance_maps, center_maps,264          instance_score_maps)265 266 267@tf.function268def _merge_semantic_and_instance_maps(269    semantic_prediction: tf.Tensor,270    instance_maps: tf.Tensor,271    thing_class_ids: tf.Tensor,272    label_divisor: int,273    stuff_area_limit: int,274    void_label: int) -> tf.Tensor:275  """Merges semantic and instance maps to obtain panoptic segmentation.276 277  This function merges the semantic segmentation and class-agnostic278  instance segmentation to form the panoptic segmentation. In particular,279  the class label of each instance mask is inferred from the majority280  votes from the corresponding pixels in the semantic segmentation. This281  operation is first poposed in the DeeperLab paper and adopted by the282  Panoptic-DeepLab.283 284  - DeeperLab: Single-Shot Image Parser, T-J Yang, et al. arXiv:1902.05093.285  - Panoptic-DeepLab, B. Cheng, et al. In CVPR, 2020.286 287  Note that this function only supports batch = 1 for simplicity. Additionally,288  this function has a slightly different implementation from the provided289  TensorFlow implementation `merge_ops` but with a similar performance. This290  function is mainly used as a backup solution when you could not successfully291  compile the provided TensorFlow implementation. To reproduce our results,292  please use the provided TensorFlow implementation (i.e., not use this293  function, but the `merge_ops.merge_semantic_and_instance_maps`).294 295  Args:296    semantic_prediction: A tf.Tensor of shape [batch, height, width].297    instance_maps: A tf.Tensor of shape [batch, height, width].298    thing_class_ids: A tf.Tensor of shape [N] containing N thing indices.299    label_divisor: An integer specifying the label divisor of the dataset.300    stuff_area_limit: An integer specifying the number of pixels that stuff301      regions need to have at least. The stuff region will be included in the302      panoptic prediction, only if its area is larger than the limit; otherwise,303      it will be re-assigned as void_label.304    void_label: An integer specifying the void label.305 306  Returns:307    panoptic_prediction: A tf.Tensor with shape [batch, height, width].308  """309  prediction_shape = semantic_prediction.get_shape().as_list()310  # This implementation only supports batch size of 1. Since model construction311  # might lose batch size information (and leave it to None), override it here.312  prediction_shape[0] = 1313  semantic_prediction = tf.ensure_shape(semantic_prediction, prediction_shape)314  instance_maps = tf.ensure_shape(instance_maps, prediction_shape)315 316  # Default panoptic_prediction to have semantic label = void_label.317  panoptic_prediction = tf.ones_like(318      semantic_prediction) * void_label * label_divisor319 320  # Start to paste predicted `thing` regions to panoptic_prediction.321  # Infer `thing` segmentation regions from semantic prediction.322  semantic_thing_segmentation = tf.zeros_like(semantic_prediction,323                                              dtype=tf.bool)324  for thing_class in thing_class_ids:325    semantic_thing_segmentation = tf.math.logical_or(326        semantic_thing_segmentation,327        semantic_prediction == thing_class)328  # Keep track of how many instances for each semantic label.329  num_instance_per_semantic_label = tf.TensorArray(330      tf.int32, size=0, dynamic_size=True, clear_after_read=False)331  instance_ids, _ = tf.unique(tf.reshape(instance_maps, [-1]))332  for instance_id in instance_ids:333    # Instance ID 0 is reserved for crowd region.334    if instance_id == 0:335      continue336    thing_mask = tf.math.logical_and(instance_maps == instance_id,337                                     semantic_thing_segmentation)338    if tf.reduce_sum(tf.cast(thing_mask, tf.int32)) == 0:339      continue340    semantic_bin_counts = tf.math.bincount(341        tf.boolean_mask(semantic_prediction, thing_mask))342    semantic_majority = tf.cast(343        tf.math.argmax(semantic_bin_counts), tf.int32)344 345    while num_instance_per_semantic_label.size() <= semantic_majority:346      num_instance_per_semantic_label = num_instance_per_semantic_label.write(347          num_instance_per_semantic_label.size(), 0)348 349    new_instance_id = (350        num_instance_per_semantic_label.read(semantic_majority) + 1)351    num_instance_per_semantic_label = num_instance_per_semantic_label.write(352        semantic_majority, new_instance_id)353    panoptic_prediction = tf.where(354        thing_mask,355        tf.ones_like(panoptic_prediction) * semantic_majority * label_divisor356        + new_instance_id,357        panoptic_prediction)358 359  # Done with `num_instance_per_semantic_label` tensor array.360  num_instance_per_semantic_label.close()361 362  # Start to paste predicted `stuff` regions to panoptic prediction.363  instance_stuff_regions = instance_maps == 0364  semantic_ids, _ = tf.unique(tf.reshape(semantic_prediction, [-1]))365  for semantic_id in semantic_ids:366    if tf.reduce_sum(tf.cast(thing_class_ids == semantic_id, tf.int32)) > 0:367      continue368    # Check stuff area.369    stuff_mask = tf.math.logical_and(semantic_prediction == semantic_id,370                                     instance_stuff_regions)371    stuff_area = tf.reduce_sum(tf.cast(stuff_mask, tf.int32))372    if stuff_area >= stuff_area_limit:373      panoptic_prediction = tf.where(374          stuff_mask,375          tf.ones_like(panoptic_prediction) * semantic_id * label_divisor,376          panoptic_prediction)377 378  return panoptic_prediction379 380 381class SemanticOnlyPostProcessor(tf.keras.layers.Layer):382  """This class contains code of a semantic only post-processor."""383 384  def __init__(self):385    """Initializes a semantic only post-processor."""386    super(SemanticOnlyPostProcessor, self).__init__(387        name='SemanticOnlyPostProcessor')388 389  def call(self, result_dict: Dict[Text, tf.Tensor]) -> Dict[Text, tf.Tensor]:390    """Performs the post-processing given model predicted results.391 392    Args:393      result_dict: A dictionary of tf.Tensor containing model results. The dict394      has to contain395       - common.PRED_SEMANTIC_PROBS_KEY,396 397    Returns:398      The post-processed dict of tf.Tensor, containing the following:399       - common.PRED_SEMANTIC_KEY,400    """401    processed_dict = {}402    processed_dict[common.PRED_SEMANTIC_KEY] = _get_semantic_predictions(403        result_dict[common.PRED_SEMANTIC_PROBS_KEY])404    return processed_dict405 406 407class PostProcessor(tf.keras.layers.Layer):408  """This class contains code of a Panoptic-Deeplab post-processor."""409 410  def __init__(411      self,412      config: config_pb2.ExperimentOptions,413      dataset_descriptor: dataset.DatasetDescriptor):414    """Initializes a Panoptic-Deeplab post-processor.415 416    Args:417      config: A config_pb2.ExperimentOptions configuration.418      dataset_descriptor: A dataset.DatasetDescriptor.419    """420    super(PostProcessor, self).__init__(name='PostProcessor')421    self._post_processor = functools.partial(422        _get_panoptic_predictions,423        center_threshold=config.evaluator_options.center_score_threshold,424        thing_class_ids=tf.convert_to_tensor(425            dataset_descriptor.class_has_instances_list),426        label_divisor=dataset_descriptor.panoptic_label_divisor,427        stuff_area_limit=config.evaluator_options.stuff_area_limit,428        void_label=dataset_descriptor.ignore_label,429        nms_kernel_size=config.evaluator_options.nms_kernel,430        keep_k_centers=config.evaluator_options.keep_k_centers,431        merge_semantic_and_instance_with_tf_op=(432            config.evaluator_options.merge_semantic_and_instance_with_tf_op),433        )434 435  def call(self, result_dict: Dict[Text, tf.Tensor]) -> Dict[Text, tf.Tensor]:436    """Performs the post-processing given model predicted results.437 438    Args:439      result_dict: A dictionary of tf.Tensor containing model results. The dict440      has to contain441        - common.PRED_SEMANTIC_PROBS_KEY,442        - common.PRED_CENTER_HEATMAP_KEY,443        - common.PRED_OFFSET_MAP_KEY,444 445    Returns:446      The post-processed dict of tf.Tensor, containing the following:447        - common.PRED_SEMANTIC_KEY,448        - common.PRED_INSTANCE_KEY,449        - common.PRED_PANOPTIC_KEY,450        - common.PRED_INSTANCE_CENTER_KEY,451        - common.PRED_INSTANCE_SCORES_KEY,452    """453    processed_dict = {}454    (processed_dict[common.PRED_PANOPTIC_KEY],455     processed_dict[common.PRED_SEMANTIC_KEY],456     processed_dict[common.PRED_INSTANCE_KEY],457     processed_dict[common.PRED_INSTANCE_CENTER_KEY],458     processed_dict[common.PRED_INSTANCE_SCORES_KEY]459    ) = self._post_processor(460        result_dict[common.PRED_SEMANTIC_PROBS_KEY],461        result_dict[common.PRED_CENTER_HEATMAP_KEY],462        result_dict[common.PRED_OFFSET_MAP_KEY])463    return processed_dict464