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 Motion-DeepLab results."""17 18from typing import Tuple19 20import tensorflow as tf21 22 23def assign_instances_to_previous_tracks(24    prev_centers: tf.Tensor,25    current_centers: tf.Tensor,26    heatmap: tf.Tensor,27    offsets: tf.Tensor,28    panoptic_map: tf.Tensor,29    next_id: tf.Tensor,30    label_divisor: int,31    sigma=7) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:32  """Greedy assignment of current centers to previous centers.33 34  Current centers are selected in decreasing order of confidence (heatmap35  scores). These centers are transformed with the offsets and assigned to36  previous centers.37 38  Args:39    prev_centers: A tf.Tensor containing previous centers of shape [Np, 5]. This40      tensor contains:41      [0]: The x-coordinate.42      [1]: The y-coordinate.43      [2]: The panoptic ID.44      [3]: The geometric mean of width and height of the instance mask.45      [4]: The number of frames that no new masks got assigned to this center.46    current_centers: A tf.Tensor containing centers of current frame of shape47      [Nc, 5]. This tensor contains:48      [0]: The x-coordinate.49      [1]: The y-coordinate.50      [2]: The panoptic ID.51      [3]: The geometric mean of width and height of the instance mask.52      [4]: The number of frames that no new masks got assigned to this center.53    heatmap: A tf.Tensor of shape [batch, height, width] containing the center54      heatmap.55    offsets: A tf.Tensor of shape [batch, height, width, 2] containing the56      center offsets.57    panoptic_map: A tf.Tensor of shape [batch, height, width] containing the58      panoptic segmentation.59    next_id: A tf.Tensor of shape [1] containing the next ID.60    label_divisor: An integer specifying the label divisor for panoptic IDs.61    sigma: An optional integer specifying the number of frames that unmatched62      centers should be kept (default: 7).63 64  Returns:65    A tuple of three tf.Tensor:66      1. The updated panoptic segmentation map that contains track IDs.67      2. The updated tensor containing all current centers (including unmatched68        previous ones).69      3. The updated next ID that can be used for new tracks.70  """71  # Switch x and y coordinates for indexing.72  center_indices = tf.concat(73      [tf.zeros([tf.shape(current_centers)[0], 1], dtype=tf.int32),74       current_centers[:, 1:2], current_centers[:, 0:1]],75      axis=1)76  confidence_scores = tf.gather_nd(heatmap, center_indices)77 78  scores = tf.argsort(confidence_scores, direction='DESCENDING')79  cond = lambda i, *_: i < tf.shape(center_indices)[0]80 81  def body(i, current_centers_loop, prev_centers_loop, new_panoptic_map_loop,82           next_id_loop):83    row_index = scores[i]84    i = tf.add(i, 1)85    center_id = current_centers_loop[row_index, 2]86    center_location = current_centers_loop[row_index, :2]87    center_offset_yx = offsets[0, center_location[1], center_location[0], :]88    center_offset_xy = center_offset_yx[::-1]89    center_location = center_offset_xy + tf.cast(center_location, tf.float32)90    center_sem_id = center_id // label_divisor91    center_mask = tf.equal(panoptic_map, center_id)92    prev_centers_class = prev_centers_loop[:, 2] // label_divisor93    prev_centers_with_same_class = tf.squeeze(94        tf.cast(95            tf.gather(96                prev_centers_loop,97                tf.where(tf.equal(prev_centers_class, center_sem_id)),98                axis=0), tf.float32),99        axis=1)100 101    # Check if there are still unassigned previous centers of the same class.102    if tf.shape(prev_centers_with_same_class)[0] > 0:103      # For efficieny reasons, we do not take the sqrt when we compute the104      # minimal distances. See render_panoptic_map_as_heatmap as well.105      distances = tf.reduce_sum(106          tf.square(prev_centers_with_same_class[:, :2] - center_location),107          axis=1)108      prev_center_index = tf.math.argmin(109          distances, axis=0, output_type=tf.int32)110      min_dist = distances[prev_center_index]111 112      # If previous center is within a certain range, continue track.113      if min_dist < prev_centers_with_same_class[prev_center_index, 3]:114        new_center_id = tf.cast(115            prev_centers_with_same_class[prev_center_index, 2], dtype=tf.int32)116        shape = new_panoptic_map_loop.get_shape()117        new_panoptic_map_loop = tf.where(center_mask, new_center_id,118                                         new_panoptic_map_loop)119        new_panoptic_map_loop.set_shape(shape)120        current_centers_loop = tf.tensor_scatter_nd_update(121            current_centers_loop, tf.expand_dims([row_index, 2], 0),122            [new_center_id])123        # Remove previous center.124        prev_centers_loop = tf.squeeze(125            tf.gather(126                prev_centers_loop,127                tf.where(tf.not_equal(prev_centers_loop[:, 2], new_center_id)),128                axis=0),129            axis=1)130        return (i, current_centers_loop, prev_centers_loop,131                new_panoptic_map_loop, next_id_loop)132      else:133        # Assign new track ID134        new_center_id = center_sem_id * label_divisor + next_id_loop135        shape = new_panoptic_map_loop.get_shape()136        new_panoptic_map_loop = tf.where(center_mask, new_center_id,137                                         new_panoptic_map_loop)138        new_panoptic_map_loop.set_shape(shape)139        current_centers_loop = tf.tensor_scatter_nd_update(140            current_centers_loop, tf.expand_dims([row_index, 2], 0),141            [new_center_id])142        next_id_loop += 1143        return (i, current_centers_loop, prev_centers_loop,144                new_panoptic_map_loop, next_id_loop)145    else:146      # Assign new track ID147      new_center_id = center_sem_id * label_divisor + next_id_loop148      shape = new_panoptic_map_loop.get_shape()149      new_panoptic_map_loop = tf.where(center_mask, new_center_id,150                                       new_panoptic_map_loop)151      new_panoptic_map_loop.set_shape(shape)152      current_centers_loop = tf.tensor_scatter_nd_update(153          current_centers_loop, tf.expand_dims([row_index, 2], 0),154          [new_center_id])155      next_id_loop += 1156      return (i, current_centers_loop, prev_centers_loop, new_panoptic_map_loop,157              next_id_loop)158 159  loop_start_index = tf.constant(0)160  (_, current_centers,161   unmatched_centers, new_panoptic_map, next_id) = tf.while_loop(162       cond, body,163       (loop_start_index, current_centers, prev_centers, panoptic_map,164        next_id))165 166  # Keep unmatched centers for sigma frames.167  if tf.shape(unmatched_centers)[0] > 0:168    current_centers = tf.concat([current_centers, unmatched_centers], axis=0)169 170  number_centers = tf.shape(current_centers)[0]171  indices_row = tf.range(number_centers, dtype=tf.int32)172  indices_column = tf.repeat([4], number_centers, axis=0)173  indices = tf.stack([indices_row, indices_column], axis=1)174  current_centers = tf.tensor_scatter_nd_add(175      current_centers, indices,176      tf.repeat([1], number_centers, axis=0))177 178  # Remove centers after sigma frames.179  current_centers = tf.squeeze(180      tf.gather(181          current_centers,182          tf.where(tf.not_equal(current_centers[:, 4], sigma)),183          axis=0),184      axis=1)185 186  return new_panoptic_map, current_centers, next_id187 188 189def render_panoptic_map_as_heatmap(190    panoptic_map: tf.Tensor, sigma: int, label_divisor: int,191    void_label: int) -> Tuple[tf.Tensor, tf.Tensor]:192  """Extracts centers from panoptic map and renders as heatmap."""193  gaussian_size = 6 * sigma + 3194  x = tf.range(gaussian_size, dtype=tf.float32)195  y = tf.expand_dims(x, axis=1)196  x0, y0 = 3 * sigma + 1, 3 * sigma + 1197  gaussian = tf.math.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma**2))198  gaussian = tf.cast(tf.reshape(gaussian, [-1]), tf.float32)199 200  height = tf.shape(panoptic_map)[1]201  width = tf.shape(panoptic_map)[2]202  # Pad center to make boundary handling easier.203  center_pad_begin = int(round(3 * sigma + 1))204  center_pad_end = int(round(3 * sigma + 2))205  center_pad = center_pad_begin + center_pad_end206 207  center = tf.zeros((height + center_pad, width + center_pad))208  unique_ids, _ = tf.unique(tf.reshape(panoptic_map, [-1]))209  centers_and_ids = tf.TensorArray(210      tf.int32, size=0, dynamic_size=True, clear_after_read=False)211  counter = tf.zeros([], dtype=tf.int32)212 213  for panoptic_id in unique_ids:214    semantic_id = panoptic_id // label_divisor215    # Filter out IDs that should be ignored, are stuff classes or crowd.216    # Stuff classes and crowd regions both have IDs of the form panoptic_id =217    # semantic_id * label_divisor218    if semantic_id == void_label or panoptic_id % label_divisor == 0:219      continue220 221    # Convert [[0, y0, x0], ...] to [[0, ...], [y0, ...], [x0, ...]].222    mask_index = tf.cast(223        tf.transpose(tf.where(panoptic_map == panoptic_id)), tf.float32)224    mask_size = (225        tf.reduce_max(mask_index, axis=1) - tf.reduce_min(mask_index, axis=1))226    # The radius is defined as the geometric mean of width and height.227    # For efficieny reasons, we do not take the sqrt when we compute the minimal228    # distances. See assign_instances_to_previous_tracks as well.229    mask_radius = tf.cast(tf.round(mask_size[1] * mask_size[2]), tf.int32)230    centers = tf.reduce_mean(mask_index, axis=1)231 232    center_x = tf.cast(tf.round(centers[2]), tf.int32)233    center_y = tf.cast(tf.round(centers[1]), tf.int32)234    centers_and_ids = centers_and_ids.write(235        counter,236        [center_x, center_y, tf.cast(panoptic_id, tf.int32), mask_radius, 0])237    counter += 1238 239    # Due to the padding with center_pad_begin in center, the computed center240    # becomes the upper left corner in the center tensor.241    upper_left = center_x, center_y242    bottom_right = (upper_left[0] + gaussian_size,243                    upper_left[1] + gaussian_size)244 245    indices_x, indices_y = tf.meshgrid(246        tf.range(upper_left[0], bottom_right[0]),247        tf.range(upper_left[1], bottom_right[1]))248    indices = tf.transpose(249        tf.stack([tf.reshape(indices_y, [-1]),250                  tf.reshape(indices_x, [-1])]))251 252    center = tf.tensor_scatter_nd_max(253        center, indices, gaussian, name='center_scatter')254 255  center = center[center_pad_begin:(center_pad_begin + height),256                  center_pad_begin:(center_pad_begin + width)]257  return tf.expand_dims(center, axis=0), centers_and_ids.stack()258