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 MaX-DeepLab results."""17 18import functools19from typing import List, Tuple, Dict, Text20 21import tensorflow as tf22 23from deeplab2 import common24from deeplab2 import config_pb225from deeplab2.data import dataset26from deeplab2.model import utils27 28 29def _get_transformer_class_prediction(30    transformer_class_probs: tf.Tensor,31    transformer_class_confidence_threshold: float32    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:33  """Computes the transformer class prediction and confidence score.34 35  Args:36    transformer_class_probs: A tf.Tensor of shape [num_mask_slots,37      num_thing_stuff_classes + 1]. It is a pixel level logit scores where the38      num_mask_slots is the number of mask slots (for both thing classes and39      stuff classes) in MaX-DeepLab. The last channel indicates a `void` class.40    transformer_class_confidence_threshold: A float for thresholding the41      confidence of the transformer_class_probs. The panoptic mask slots with42      class confidence less than the threshold are filtered and not used for43      panoptic prediction. Only masks whose confidence is larger than the44      threshold are counted in num_detections.45 46  Returns:47    A tuple of:48    - the detected mask class prediction as float32 tf.Tensor of shape49      [num_detections].50    - the detected mask indices as tf.Tensor of shape [num_detections].51    - the number of detections as tf.Tensor of shape [1].52  """53  transformer_class_pred = tf.cast(54      tf.argmax(transformer_class_probs, axis=-1), tf.float32)55  transformer_class_confidence = tf.reduce_max(56      transformer_class_probs, axis=-1, keepdims=False)57  # Filter mask IDs with class confidence less than the threshold.58  thresholded_mask = tf.cast(59      tf.greater_equal(transformer_class_confidence,60                       transformer_class_confidence_threshold), tf.float32)61  transformer_class_confidence = (transformer_class_confidence62                                  * thresholded_mask)63 64  detected_mask_indices = tf.where(tf.greater(thresholded_mask, 0.5))[:, 0]65  detected_mask_class_pred = tf.gather(66      transformer_class_pred, detected_mask_indices)67  num_detections = tf.shape(detected_mask_indices)[0]68  return detected_mask_class_pred, detected_mask_indices, num_detections69 70 71def _get_mask_id_and_semantic_maps(72    thing_class_ids: List[int],73    stuff_class_ids: List[int],74    pixel_space_mask_logits: tf.Tensor,75    transformer_class_probs: tf.Tensor,76    image_shape: List[int],77    pixel_confidence_threshold=0.4,78    transformer_class_confidence_threshold=0.7,79    pieces=1) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:80  """Computes the pixel-level mask ID map and semantic map per image.81 82  Args:83    thing_class_ids: A List of integers of shape [num_thing_classes] containing84      thing class indices.85    stuff_class_ids: A List of integers of shape [num_thing_classes] containing86      stuff class indices.87    pixel_space_mask_logits: A tf.Tensor of shape [height, width,88      num_mask_slots]. It is a pixel level logit scores where the89      num_mask_slots is the number of mask slots (for both thing classes90      and stuff classes) in MaX-DeepLab.91    transformer_class_probs: A tf.Tensor of shape [num_mask_slots,92      num_thing_stuff_classes + 1]. It is a pixel level logit scores where the93      num_mask_slots is the number of mask slots (for both thing classes and94      stuff classes) in MaX-DeepLab. The last channel indicates a `void` class.95    image_shape: A list of integers specifying the [height, width] of input96      image.97    pixel_confidence_threshold: A float indicating a threshold for the pixel98      level softmax probability confidence of transformer mask logits. If less99      than the threshold, the pixel locations have confidence `0` in100      `confident_regions` output, and represent `void` (ignore) regions.101    transformer_class_confidence_threshold: A float for thresholding the102      confidence of the transformer_class_probs. The panoptic mask slots with103      class confidence less than the threshold are filtered and not used for104      panoptic prediction.105    pieces: An integer indicating the number of pieces in the piece-wise106      operation. When computing panpotic prediction and confident regions, the107      mask logits are divided width-wise into multiple pieces and processed108      piece-wise due to the GPU memory limit. Then, the piece-wise outputs are109      concatenated along the width into the original mask shape. Defaults to 1.110 111  Returns:112    A tuple of:113    - the mask ID prediction as tf.Tensor with shape [height, width].114    - the semantic prediction as tf.Tensor with shape [height, width].115    - the thing region mask as tf.Tensor with shape [height, width].116    - the stuff region mask as tf.Tensor with shape [height, width].117 118  Raises:119    ValueError: When input image's `width - 1` is not divisible by `pieces`.120  """121  # The last channel indicates `void` class and thus is not included.122  transformer_class_probs = transformer_class_probs[..., :-1]123  # Generate mapping from mask IDs to dataset's thing and stuff semantic IDs.124  thing_stuff_class_ids = thing_class_ids + stuff_class_ids125 126  detected_mask_class_pred, detected_mask_indices, num_detections = (127      _get_transformer_class_prediction(transformer_class_probs,128                                        transformer_class_confidence_threshold))129  # If num_detections = 0, return empty result maps.130  def _return_empty_mask_id_and_semantic_maps():131    return (132        tf.ones([image_shape[0], image_shape[1]], dtype=tf.int32),133        tf.zeros([image_shape[0], image_shape[1]], dtype=tf.int32),134        tf.zeros([image_shape[0], image_shape[1]], dtype=tf.float32),135        tf.zeros([image_shape[0], image_shape[1]], dtype=tf.float32))136 137  # If num_detections > 0:138  def _generate_mask_id_and_semantic_maps():139    output_mask_id_map = []140    output_confident_region = []141    logits_width = pixel_space_mask_logits.get_shape().as_list()[1]142    output_width = image_shape[1]143 144    if (output_width - 1) % pieces > 0:145      raise ValueError('`output_width - 1` must be divisible by `pieces`.')146    # Use of input shape of a multiple of the feature stride, plus one, so that147    # it preserves left- and right-alignment.148    piece_output_width = (output_width - 1) // pieces + 1149 150    for piece_id in range(pieces):151      piece_begin = (logits_width - 1) // pieces * piece_id152      # Use of input shape of a multiple of the feature stride, plus one, so153      # that it preserves left- and right-alignment.154      piece_end = (logits_width - 1) // pieces * (piece_id + 1) + 1155      piece_pixel_mask_logits = (156          pixel_space_mask_logits[:, piece_begin:piece_end, :])157      piece_pixel_mask_logits = tf.compat.v1.image.resize_bilinear(158          tf.expand_dims(piece_pixel_mask_logits, 0),159          (image_shape[0], piece_output_width),160          align_corners=True)161      piece_pixel_mask_logits = tf.squeeze(piece_pixel_mask_logits, axis=0)162      piece_detected_pixel_mask_logits = tf.gather(163          piece_pixel_mask_logits, detected_mask_indices, axis=-1)164      # Filter the pixels which are assigned to a mask ID that does not survive.165      piece_max_logits = tf.reduce_max(piece_pixel_mask_logits, axis=-1)166      piece_detected_max_logits = tf.reduce_max(167          piece_detected_pixel_mask_logits, axis=-1)168      piece_detected_mask = tf.cast(tf.math.equal(169          piece_max_logits, piece_detected_max_logits), tf.float32)170      # Filter with pixel mask threshold.171      piece_pixel_confidence_map = tf.reduce_max(172          tf.nn.softmax(piece_detected_pixel_mask_logits, axis=-1), axis=-1)173      piece_confident_region = tf.cast(174          piece_pixel_confidence_map > pixel_confidence_threshold, tf.float32)175      piece_confident_region = piece_confident_region * piece_detected_mask176      piece_mask_id_map = tf.cast(177          tf.argmax(piece_detected_pixel_mask_logits, axis=-1), tf.int32)178      if piece_id == pieces - 1:179        output_mask_id_map.append(piece_mask_id_map)180        output_confident_region.append(piece_confident_region)181      else:182        output_mask_id_map.append(piece_mask_id_map[:, :-1])183        output_confident_region.append(piece_confident_region[:, :-1])184 185    mask_id_map = tf.concat(output_mask_id_map, axis=1)186    confident_region = tf.concat(output_confident_region, axis=1)187    mask_id_map_flat = tf.reshape(mask_id_map, [-1])188    mask_id_semantic_map_flat = tf.gather(189        detected_mask_class_pred, mask_id_map_flat)190    mask_id_semantic_map = tf.reshape(191        mask_id_semantic_map_flat, [image_shape[0], image_shape[1]])192    # Generate thing and stuff masks (with value 1/0 indicates the193    # presence/absence)194    thing_mask = tf.cast(mask_id_semantic_map < len(thing_class_ids),195                         tf.float32) * confident_region196    stuff_mask = tf.cast(mask_id_semantic_map >= len(thing_class_ids),197                         tf.float32) * confident_region198    # Generate semantic_map.199    semantic_map = tf.gather(200        tf.convert_to_tensor(thing_stuff_class_ids),201        tf.cast(tf.round(mask_id_semantic_map_flat), tf.int32))202    semantic_map = tf.reshape(semantic_map, [image_shape[0], image_shape[1]])203    # Add 1 because mask ID 0 is reserved for unconfident region.204    mask_id_map_plus_one = mask_id_map + 1205    semantic_map = tf.cast(tf.round(semantic_map), tf.int32)206    return (mask_id_map_plus_one, semantic_map, thing_mask, stuff_mask)207 208  mask_id_map_plus_one, semantic_map, thing_mask, stuff_mask = tf.cond(209      tf.cast(num_detections, tf.float32) < tf.cast(0.5, tf.float32),210      _return_empty_mask_id_and_semantic_maps,211      _generate_mask_id_and_semantic_maps)212 213  return (mask_id_map_plus_one, semantic_map, thing_mask, stuff_mask)214 215 216def _filter_by_count(input_index_map: tf.Tensor,217                     area_limit: int) -> Tuple[tf.Tensor, tf.Tensor]:218  """Filters input index map by area limit threshold per index.219 220  Args:221    input_index_map: A float32 tf.Tensor of shape [batch, height, width].222    area_limit: An integer specifying the number of pixels that each index223      regions need to have at least. If not over the limit, the index regions224      are masked (zeroed) out.225 226  Returns:227    masked input_index_map: A tf.Tensor with shape [batch, height, width],228      masked by the area_limit threshold.229    mask: A tf.Tensor with shape [batch, height, width]. It is a pixel-level230      mask with 1. indicating the regions over the area limit, and 0. otherwise.231  """232  batch_size = tf.shape(input_index_map)[0]233  index_map = tf.cast(tf.round(input_index_map), tf.int32)234  index_map_flat = tf.reshape(index_map, [batch_size, -1])235  counts = tf.math.bincount(index_map_flat, axis=-1)236  counts_map = tf.gather(counts, index_map_flat, batch_dims=1)237  counts_map = tf.reshape(counts_map, tf.shape(index_map))238 239  mask = tf.cast(240      tf.cast(counts_map, tf.float32) > tf.cast(area_limit - 0.5, tf.float32),241      input_index_map.dtype)242  return input_index_map * mask, mask243 244 245def _merge_mask_id_and_semantic_maps(246    mask_id_maps_plus_one: tf.Tensor,247    semantic_maps: tf.Tensor,248    thing_masks: tf.Tensor,249    stuff_masks: tf.Tensor,250    void_label: int,251    label_divisor: int,252    thing_area_limit: int,253    stuff_area_limit: int,) -> tf.Tensor:254  """Merges mask_id maps and semantic_maps to obtain panoptic segmentation.255 256  Args:257    mask_id_maps_plus_one: A tf.Tensor of shape [batch, height, width].258    semantic_maps: A tf.Tensor of shape [batch, height, width].259    thing_masks: A float32 tf.Tensor of shape [batch, height, width] containing260      masks with 1. at thing regions, 0. otherwise.261    stuff_masks: A float32 tf.Tensor of shape [batch, height, width] containing262      masks with 1. at thing regions, 0. otherwise.263    void_label: An integer specifying the void label.264    label_divisor: An integer specifying the label divisor of the dataset.265    thing_area_limit: An integer specifying the number of pixels that thing266      regions need to have at least. The thing region will be included in the267      panoptic prediction, only if its area is larger than the limit; otherwise,268      it will be re-assigned as void_label.269    stuff_area_limit: An integer specifying the number of pixels that stuff270      regions need to have at least. The stuff region will be included in the271      panoptic prediction, only if its area is larger than the limit; otherwise,272      it will be re-assigned as void_label.273 274  Returns:275    panoptic_maps: A tf.Tensor with shape [batch, height, width].276 277  """278  thing_mask_id_maps_plus_one = (tf.cast(mask_id_maps_plus_one, tf.float32)279                                 * thing_masks)280  # We increase semantic_maps by 1 before masking (zeroing) by thing_masks and281  # stuff_masks, to ensure all valid semantic IDs are greater than 0 and thus282  # not masked out.283  semantic_maps_plus_one = semantic_maps + 1284  tf.debugging.assert_less(285      tf.reduce_sum(thing_masks * stuff_masks), 0.5,286      message='thing_masks and stuff_masks must be mutually exclusive.')287 288  thing_semantic_maps = (tf.cast(semantic_maps_plus_one, tf.float32)289                         * thing_masks)290  stuff_semantic_maps = (tf.cast(semantic_maps_plus_one, tf.float32)291                         * stuff_masks)292 293  # Filter stuff_semantic_maps by stuff_area_limit.294  stuff_semantic_maps, _ = _filter_by_count(295      stuff_semantic_maps, stuff_area_limit)296  # Filter thing_mask_id_map and thing_semantic_map by thing_area_limit297  thing_mask_id_maps_plus_one, mask_id_count_filter_mask = _filter_by_count(298      thing_mask_id_maps_plus_one, thing_area_limit)299  thing_semantic_maps = thing_semantic_maps * mask_id_count_filter_mask300 301  # Filtered un-confident region will be replaced with `void_label`. The302  # "plus_one" will be reverted, the un-confident region (0) will be -1, and so303  # we add (void + 1)304  semantic_maps_new = thing_semantic_maps + stuff_semantic_maps - 1.0305  semantic_maps_new = (tf.cast(semantic_maps_new < -0.5, tf.float32)306                       * tf.cast(void_label + 1, tf.float32)307                       + semantic_maps_new)308  panoptic_maps = (semantic_maps_new * label_divisor309                   + thing_mask_id_maps_plus_one)310  panoptic_maps = tf.cast(tf.round(panoptic_maps), tf.int32)311  return panoptic_maps312 313 314def _get_panoptic_predictions(315    pixel_space_mask_logits: tf.Tensor,316    transformer_class_logits: tf.Tensor,317    thing_class_ids: List[int],318    void_label: int,319    label_divisor: int,320    thing_area_limit: int,321    stuff_area_limit: int,322    image_shape: List[int],323    pixel_confidence_threshold=0.4,324    transformer_class_confidence_threshold=0.7,325    pieces=1) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:326  """Computes the pixel-level panoptic, mask ID, and semantic maps.327 328  Args:329    pixel_space_mask_logits: A tf.Tensor of shape [batch, strided_height,330      strided_width, num_mask_slots]. It is a pixel level logit scores where the331      num_mask_slots is the number of mask slots (for both thing classes332      and stuff classes) in MaX-DeepLab.333    transformer_class_logits: A tf.Tensor of shape [batch, num_mask_slots,334      num_thing_stuff_classes + 1]. It is a pixel level logit scores where the335      num_mask_slots is the number of mask slots (for both thing classes and336      stuff classes) in MaX-DeepLab. The last channel indicates a `void` class.337    thing_class_ids: A List of integers of shape [num_thing_classes] containing338      thing class indices.339    void_label: An integer specifying the void label.340    label_divisor: An integer specifying the label divisor of the dataset.341    thing_area_limit: An integer specifying the number of pixels that thing342      regions need to have at least. The thing region will be included in the343      panoptic prediction, only if its area is larger than the limit; otherwise,344      it will be re-assigned as void_label.345    stuff_area_limit: An integer specifying the number of pixels that stuff346      regions need to have at least. The stuff region will be included in the347      panoptic prediction, only if its area is larger than the limit; otherwise,348      it will be re-assigned as void_label.349    image_shape: A list of integers specifying the [height, width] of input350      image.351    pixel_confidence_threshold: A float indicating a threshold for the pixel352      level softmax probability confidence of transformer mask logits. If less353      than the threshold, the pixel locations have confidence `0` in354      `confident_regions` output, and represent `void` (ignore) regions.355    transformer_class_confidence_threshold: A float for thresholding the356      confidence of the transformer_class_probs. The panoptic mask slots with357      class confidence less than the threshold are filtered and not used for358      panoptic prediction.359    pieces: An integer indicating the number of pieces in the piece-wise360      operation in `_get_mask_id_and_semantic_maps`. When computing panoptic361      prediction and confident regions, the mask logits are divided width-wise362      into multiple pieces and processed piece-wise due to the GPU memory limit.363      Then, the piece-wise outputs are concatenated along the width into the364      original mask shape. Defaults to 1.365 366  Returns:367    A tuple of:368    - the panoptic prediction as tf.Tensor with shape [batch, height, width].369    - the mask ID prediction as tf.Tensor with shape [batch, height, width].370    - the semantic prediction as tf.Tensor with shape [batch, height, width].371  """372  transformer_class_probs = tf.nn.softmax(transformer_class_logits, axis=-1)373  batch_size = tf.shape(transformer_class_logits)[0]374  # num_thing_stuff_classes does not include `void` class, so we decrease by 1.375  num_thing_stuff_classes = (376      transformer_class_logits.get_shape().as_list()[-1] - 1)377  # Generate thing and stuff class ids378  stuff_class_ids = utils.get_stuff_class_ids(379      num_thing_stuff_classes, thing_class_ids, void_label)380 381  mask_id_map_plus_one_lists = tf.TensorArray(382      tf.int32, size=batch_size, dynamic_size=False)383  semantic_map_lists = tf.TensorArray(384      tf.int32, size=batch_size, dynamic_size=False)385  thing_mask_lists = tf.TensorArray(386      tf.float32, size=batch_size, dynamic_size=False)387  stuff_mask_lists = tf.TensorArray(388      tf.float32, size=batch_size, dynamic_size=False)389  for i in tf.range(batch_size):390    mask_id_map_plus_one, semantic_map, thing_mask, stuff_mask = (391        _get_mask_id_and_semantic_maps(392            thing_class_ids, stuff_class_ids,393            pixel_space_mask_logits[i, ...], transformer_class_probs[i, ...],394            image_shape, pixel_confidence_threshold,395            transformer_class_confidence_threshold, pieces)396        )397    mask_id_map_plus_one_lists = mask_id_map_plus_one_lists.write(398        i, mask_id_map_plus_one)399    semantic_map_lists = semantic_map_lists.write(i, semantic_map)400    thing_mask_lists = thing_mask_lists.write(i, thing_mask)401    stuff_mask_lists = stuff_mask_lists.write(i, stuff_mask)402  # This does not work with unknown shapes.403  mask_id_maps_plus_one = mask_id_map_plus_one_lists.stack()404  semantic_maps = semantic_map_lists.stack()405  thing_masks = thing_mask_lists.stack()406  stuff_masks = stuff_mask_lists.stack()407 408  panoptic_maps = _merge_mask_id_and_semantic_maps(409      mask_id_maps_plus_one, semantic_maps, thing_masks, stuff_masks,410      void_label, label_divisor, thing_area_limit, stuff_area_limit)411  return panoptic_maps, mask_id_maps_plus_one, semantic_maps412 413 414class PostProcessor(tf.keras.layers.Layer):415  """This class contains code of a MaX-DeepLab post-processor."""416 417  def __init__(418      self,419      config: config_pb2.ExperimentOptions,420      dataset_descriptor: dataset.DatasetDescriptor):421    """Initializes a MaX-DeepLab post-processor.422 423    Args:424      config: A config_pb2.ExperimentOptions configuration.425      dataset_descriptor: A dataset.DatasetDescriptor.426    """427    super(PostProcessor, self).__init__(name='PostProcessor')428    self._post_processor = functools.partial(429        _get_panoptic_predictions,430        thing_class_ids=list(dataset_descriptor.class_has_instances_list),431        void_label=dataset_descriptor.ignore_label,432        label_divisor=dataset_descriptor.panoptic_label_divisor,433        thing_area_limit=config.evaluator_options.thing_area_limit,434        stuff_area_limit=config.evaluator_options.stuff_area_limit,435        image_shape=list(config.eval_dataset_options.crop_size),436        transformer_class_confidence_threshold=config.evaluator_options437        .transformer_class_confidence_threshold,438        pixel_confidence_threshold=config.evaluator_options439        .pixel_confidence_threshold,440        pieces=1)441 442  def call(self, result_dict: Dict[Text, tf.Tensor]) -> Dict[Text, tf.Tensor]:443    """Performs the post-processing given model predicted results.444 445    Args:446      result_dict: A dictionary of tf.Tensor containing model results. The dict447      has to contain448       - common.PRED_PIXEL_SPACE_MASK_LOGITS_KEY,449       - common.PRED_TRANSFORMER_CLASS_LOGITS_KEY,450 451    Returns:452      The post-processed dict of tf.Tensor, containing the following:453       - common.PRED_SEMANTIC_KEY,454       - common.PRED_INSTANCE_KEY,455       - common.PRED_PANOPTIC_KEY,456    """457    processed_dict = {}458    (processed_dict[common.PRED_PANOPTIC_KEY],459     processed_dict[common.PRED_INSTANCE_KEY],460     processed_dict[common.PRED_SEMANTIC_KEY]461    ) = self._post_processor(462        result_dict[common.PRED_PIXEL_SPACE_MASK_LOGITS_KEY],463        result_dict[common.PRED_TRANSFORMER_CLASS_LOGITS_KEY])464    return processed_dict465