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 utility functions for the model code."""17 18from typing import Any, List, MutableMapping, MutableSequence, Optional, Set19 20import tensorflow as tf21 22from deeplab2 import common23from deeplab2 import config_pb224 25layers = tf.keras.layers26 27_PREDICTION_WITH_NEAREST_UPSAMPLING = (28    common.PRED_INSTANCE_KEY,29    common.PRED_INSTANCE_CENTER_KEY,30    common.PRED_INSTANCE_SCORES_KEY,31    common.PRED_PANOPTIC_KEY,32    common.PRED_SEMANTIC_KEY,33    common.PRED_NEXT_PANOPTIC_KEY,34    common.PRED_CONCAT_NEXT_PANOPTIC_KEY,35    common.PRED_CENTER_HEATMAP_KEY,36)37 38_PREDICTION_WITH_BILINEAR_UPSAMPLING = (39    common.PRED_SEMANTIC_PROBS_KEY,40    common.PRED_OFFSET_MAP_KEY,41)42 43_INPUT_WITH_NEAREST_UPSAMPLING = (44    common.GT_INSTANCE_CENTER_KEY,45)46 47_INPUT_WITH_BILINEAR_UPSAMPLING = (48    common.IMAGE,49    common.GT_INSTANCE_REGRESSION_KEY50)51 52 53def _scale_helper(value, scale):54  if isinstance(value, tf.Tensor):55    return tf.cast(56        (tf.cast(value, dtype=tf.float32) - 1.0) * scale + 1.0,57        dtype=tf.int32)58  else:59    return int((float(value) - 1.0) * scale + 1.0)60 61 62def scale_mutable_sequence(input_sequence: MutableSequence[int],63                           scale: float) -> MutableSequence[int]:64  return [_scale_helper(x, scale) for x in input_sequence]65 66 67def scale_int_list(int_list, scale):68  return [int(x * scale) for x in int_list]69 70 71def undo_image_preprocessing(image_in: tf.Tensor, method: str,72                             perform_crop: bool,73                             regions_to_crop: List[int],74                             output_shape: List[int]) -> tf.Tensor:75  """Undoes the image preprocessing.76 77  In particular, this function slices out the valid regions (determined by78  `regions_to_crop`) in the input when perform_crop is True. After79  that, we resize the results to the desired `output_shape`.80 81  Args:82    image_in: Input image Tensor with shape [batch, height, width, n_channels].83    method: Image resize method.84    perform_crop: Boolean, performing crop or not.85    regions_to_crop: The regions to crop [height, width]. Will only apply86      cropping at the bottom right.87    output_shape: Desired shape after resizing [height, width].88 89  Returns:90    Outputs after cropping (if perform_crop = True) and resizing.91  """92  if perform_crop:93    image_out = image_in[94        :, :regions_to_crop[0], :regions_to_crop[1], :]95  else:96    image_out = image_in97  return resize_align_corners(image_out, output_shape, method=method)98 99 100def undo_preprocessing(input_or_prediction_dict: MutableMapping[str, Any],101                       regions_to_crop: List[int],102                       output_shape: List[int]) -> MutableMapping[str, Any]:103  """Undoes preprocessing for predictions.104 105  Args:106    input_or_prediction_dict: A dictionary storing different types of inputs or107      predictions.108    regions_to_crop: The regions to crop [height, width]. Will only apply109      cropping at the bottom right.110    output_shape: Desired shape after resizing [height, width].111 112  Returns:113    inputs or predictions after cropping (if perform_crop = True) and resizing.114  """115  for key in input_or_prediction_dict.keys():116    if key in _PREDICTION_WITH_NEAREST_UPSAMPLING or key in _INPUT_WITH_NEAREST_UPSAMPLING:117      input_or_prediction_dict[key] = tf.squeeze(118          undo_image_preprocessing(119              tf.expand_dims(input_or_prediction_dict[key], 3),120              'nearest',121              perform_crop=True,122              regions_to_crop=regions_to_crop,123              output_shape=output_shape),124          axis=3)125    elif key in _PREDICTION_WITH_BILINEAR_UPSAMPLING or key in _INPUT_WITH_BILINEAR_UPSAMPLING:126      input_or_prediction_dict[key] = undo_image_preprocessing(127          input_or_prediction_dict[key],128          'bilinear',129          perform_crop=True,130          regions_to_crop=regions_to_crop,131          output_shape=output_shape)132    else:133      # We only undo preprocessing for those defined in134      # _{PREDICTION,INPUT}_WITH_{NEAREST,BILINEAR}_UPSAMPLING.135      # Other intermediate results are skipped.136      continue137  return input_or_prediction_dict138 139 140def add_zero_padding(input_tensor: tf.Tensor, kernel_size: int,141                     rank: int) -> tf.Tensor:142  """Adds zero-padding to the input_tensor."""143  pad_total = kernel_size - 1144  pad_begin = pad_total // 2145  pad_end = pad_total - pad_begin146  if rank == 3:147    return tf.pad(148        input_tensor,149        paddings=[[pad_begin, pad_end], [pad_begin, pad_end], [0, 0]])150  else:151    return tf.pad(152        input_tensor,153        paddings=[[0, 0], [pad_begin, pad_end], [pad_begin, pad_end], [0, 0]])154 155 156def resize_and_rescale_offsets(input_tensor: tf.Tensor, target_size):157  """Bilinearly resizes and rescales the offsets.158 159  Args:160    input_tensor: A tf.Tensor of shape [batch, height, width, 2].161    target_size: A list or tuple or 1D tf.Tensor that specifies the height and162      width after resizing.163 164  Returns:165    The input_tensor resized to shape `[batch, target_height, target_width, 2]`.166      Moreover, the offsets along the y-axis are rescaled by a factor equal to167      (target_height - 1) / (reference_height - 1) and the offsets along the168      x-axis are rescaled by a factor equal to169      (target_width - 1) / (reference_width - 1).170  """171  input_size_y = tf.shape(input_tensor)[1]172  input_size_x = tf.shape(input_tensor)[2]173 174  scale_y = tf.cast(target_size[0] - 1, tf.float32) / tf.cast(175      input_size_y - 1, tf.float32)176  scale_x = tf.cast(target_size[1] - 1, tf.float32) / tf.cast(177      input_size_x - 1, tf.float32)178 179  target_y, target_x = tf.split(180      value=input_tensor, num_or_size_splits=2, axis=3)181  target_y *= scale_y182  target_x *= scale_x183  target = tf.concat([target_y, target_x], 3)184  return resize_bilinear(target, target_size)185 186 187def resize_align_corners(input_tensor, target_size, method='bilinear'):188  """Resizes the input_tensor to target_size.189 190  This returns the same output as tf.compat.v1.image.resize(input_tensor,191  target_size, align_corners=True).192 193  Args:194    input_tensor: A tf.Tensor of shape [batch, height, width, channels].195    target_size: A list or tuple or 1D tf.Tensor that specifies the height and196      width after resizing.197    method: An optional string specifying the method used for resizing.198      Supported options are 'nearest' and 'bilinear'.199 200  Returns:201    The resized tensor.202 203  Raises:204    ValueError: An error occurs if 1) the input tensor's rank is not 4 or 2) the205      resizing method is not supported.206  """207  if method == 'bilinear':208    tf_method = tf.compat.v1.image.ResizeMethod.BILINEAR209  elif method == 'nearest':210    tf_method = tf.compat.v1.image.ResizeMethod.NEAREST_NEIGHBOR211  else:212    raise ValueError('The given method %s is not supported. Please use bilinear'213                     ' or nearest.' % method)214 215  tf.debugging.assert_rank(216      input_tensor, 4,217      message='Input tensor to resize method should have rank of 4.')218 219  return tf.compat.v1.image.resize(220      input_tensor,221      target_size,222      method=tf_method,223      align_corners=True,224      name='resize_align_corners')225 226 227def resize_bilinear(images,228                    size,229                    align_corners=True,230                    name=None):231  """TPU memory efficient version of tf.compat.v1.image.resize_bilinear.232 233  ResizeBilinear on TPU requires padded batch and channel dimensions. On a234  TPUv3, the worst case could lead to 256x memory consumption, if the235  input is, for example, [1, 257, 513, 1]. In this function, we replace the236  default resize_bilinear by two resize_bilinear operations, which put one image237  axis on the channel axis. This reduces TPU padding when batch * channel is238  small and height * width is large.239 240  Args:241    images: Input image of shape [B, H, W, C].242    size: A list of two elements: [height, width]. The new size for the images.243    align_corners: Whether to align corners of the image.244    name: Name of the operation.245 246  Returns:247    Resized image.248  """249  _, height, width, channel = images.get_shape().as_list()250  if height == size[0] and width == size[1]:251    return images252  dtype = images.dtype253  images = tf.cast(images, tf.float32)254  # We check the channel axis only since the batch size is similar (usually 1 or255  # 2). In this way, this if-else easily supports dynamic batch size without256  # using tf.cond().257  if channel > 32 or not align_corners:258    images = tf.compat.v1.image.resize_bilinear(259        images, size,260        align_corners=align_corners,261        name=name)262  else:263    images = tf.transpose(images, [0, 3, 1, 2])264    images = tf.compat.v1.image.resize_bilinear(265        images, [channel, size[0]],266        align_corners=align_corners,267        name=name + '_height' if name else None)268    images = tf.transpose(images, [0, 1, 3, 2])269    images = tf.compat.v1.image.resize_bilinear(270        images, [channel, size[1]],271        align_corners=align_corners,272        name=name + '_width' if name else None)273    images = tf.transpose(images, [0, 3, 2, 1])274  return tf.cast(images, dtype)275 276 277def make_divisible(value: float,278                   divisor: int,279                   min_value: Optional[float] = None) -> int:280  """Ensures all layers have channels that are divisible by the divisor.281 282  Args:283    value: A `float` of original value.284    divisor: An `int` of the divisor that needs to be checked upon.285    min_value: A `float` of  minimum value threshold.286 287  Returns:288    The adjusted value in `int` that is divisible by divisor.289 290  Raises:291    ValueError: Minimual value should be divisible by divisor.292  """293  if min_value is None:294    min_value = divisor295  elif min_value % divisor != 0:296    raise ValueError('Minimual value should be divisible by divisor.')297 298  new_value = max(min_value, int(value + divisor / 2) // divisor * divisor)299  # Make sure that round down does not go down by more than 10%.300  if new_value < 0.9 * value:301    new_value += divisor302  return int(new_value)303 304 305def transpose_and_reshape_for_attention_operation(inputs):306  """Sequentially transposes and reshapes the tensor.307 308  Args:309    inputs: An input [batch, num_heads, length, channel] tensor.310 311  Returns:312    output: An output [batch, length, num_heads * channel] tensor.313  """314  _, num_heads, length, channel = inputs.get_shape().as_list()315  transposed_inputs = tf.transpose(inputs, [0, 2, 1, 3])316  return tf.reshape(transposed_inputs, [-1, length, num_heads * channel])317 318 319def reshape_and_transpose_for_attention_operation(inputs, num_heads):320  """Sequentially reshapes and transposes the tensor.321 322  Args:323    inputs: An input [batch, length, num_heads * channel] tensor.324    num_heads: An integer, the number of attention heads.325 326  Returns:327    output: An output [batch, num_heads, length, channel] tensor.328  """329  _, length, channels = inputs.get_shape().as_list()330  inputs = tf.reshape(inputs, [-1, length, num_heads, channels // num_heads])331  return tf.transpose(inputs, [0, 2, 1, 3])332 333 334def get_layer_name(private_attribute_name):335  if private_attribute_name[0] != '_':336    raise ValueError('Private attribute name should start with a \'_\'.')337  return private_attribute_name[1:]338 339 340def get_stem_current_name(index):341  return '_basic_block{}'.format(index + 1)342 343 344def get_low_level_conv_fusion_conv_current_names(index):345  return ('_low_level_conv{}'.format(index + 1),346          '_fusion_conv{}'.format(index + 1))347 348 349def get_conv_bn_act_current_name(index, use_bn, activation):350  name = '_conv{}'.format(index + 1)351  if use_bn:352    name += '_bn'353  if (activation is not None and354      activation.lower() != 'none' and355      activation.lower() != 'linear'):356    name += '_act'357  return name358 359 360def safe_setattr(obj, name, value):361  """A conflict-safe version of setattr().362 363  Different from setattr(), this function raises ValueError if the object364  already has an attribute with the same name.365 366  Args:367    obj: An object whose attribute has to be set.368    name: A string, the name of the attribute.369    value: Any type, the value given to the attribute.370 371  Raises:372    ValueError: If the object already has an attribute with the same name.373  """374  if hasattr(obj, name):375    raise ValueError('The object already has an attribute with the same name.')376  setattr(obj, name, value)377 378 379def pad_sequence_with_none(sequence, target_length):380  return list(sequence) + [None] * (target_length - len(sequence))381 382 383def strided_downsample(input_tensor, target_size):384  """Strided downsamples a tensor to the target size.385 386  The stride_height and stride_width is computed by (height - 1) //387  (target_height - 1) and (width - 1) // (target_width - 1). We raise an error388  if stride_height != stride_width, since this is not intended in our current389  use cases. But this check can be removed if different strides are desired.390  This function supports static shape only.391 392  Args:393    input_tensor: A [batch, height, width] tf.Tensor to be downsampled.394    target_size: A list of two integers, [target_height, target_width], the395      target size after downsampling.396 397  Returns:398    output_tensor: A [batch, target_height, target_width] tf.Tensor, the399      downsampled result.400 401  Raises:402    ValueError: If the input cannot be downsampled with integer stride, i.e.,403      (height - 1) % (target_height - 1) != 0, or (width - 1) % (target_width -404      1) != 0.405    ValueError: If the height axis stride does not equal to the width axis406      stride.407  """408  input_height, input_width = input_tensor.get_shape().as_list()[1:3]409  target_height, target_width = target_size410 411  if ((input_height - 1) % (target_height - 1) or412      (input_width - 1) % (target_width - 1)):413    raise ValueError('The input cannot be downsampled with integer striding. '414                     'Please ensure (height - 1) % (target_height - 1) == 0 '415                     'and (width - 1) % (target_width - 1) == 0.')416  stride_height = (input_height - 1) // (target_height - 1)417  stride_width = (input_width - 1) // (target_width - 1)418  if stride_height != stride_width:419    raise ValueError('The height axis stride does not equal to the width axis '420                     'stride.')421  if stride_height > 1 or stride_width > 1:422    return input_tensor[:, ::stride_height, ::stride_width]423  return input_tensor424 425 426def get_stuff_class_ids(num_thing_stuff_classes: int,427                        thing_class_ids: List[int],428                        void_label: int) -> List[int]:429  """Computes stuff_class_ids.430 431  The stuff_class_ids are computed from the num_thing_stuff_classes, the432  thing_class_ids and the void_label.433 434  Args:435    num_thing_stuff_classes: An integer specifying the number of stuff and thing436      classes, not including `void` class.437    thing_class_ids: A List of integers of length [num_thing_classes] containing438      thing class indices.439    void_label: An integer specifying the void label.440 441  Returns:442    stuff_class_ids: A sorted List of integers of shape [num_stuff_classes]443      containing stuff class indices.444  """445  if void_label >= num_thing_stuff_classes:446    thing_stuff_class_ids = list(range(num_thing_stuff_classes))447  else:448    thing_stuff_class_ids = [_ for _ in range(num_thing_stuff_classes + 1)449                             if _ is not void_label]450  return sorted(set(thing_stuff_class_ids) - set(thing_class_ids))451 452 453def get_supported_tasks(454    config: config_pb2.ExperimentOptions) -> Set[str]:455  """Gets currently supported tasks for each meta_architecture.456 457  Args:458    config: A config_pb2.ExperimentOptions configuration.459 460  Returns:461    supported_tasks: A set of strings (see common.py), optionally462     - common.TASK_PANOPTIC_SEGMENTATION,463     - common.TASK_INSTANCE_SEGMENTATION,464     - common.TASK_VIDEO_PANOPTIC_SEGMENTATION,465  """466  supported_tasks = set()467  meta_architecture = config.model_options.WhichOneof('meta_architecture')468  is_max_deeplab = meta_architecture == 'max_deeplab'469  is_motion_deeplab = meta_architecture == 'motion_deeplab'470  is_panoptic_deeplab = meta_architecture == 'panoptic_deeplab'471  is_vip_deeplab = meta_architecture == 'vip_deeplab'472  is_panoptic = (473      (config.model_options.panoptic_deeplab.instance.enable and474       is_panoptic_deeplab) or475      is_motion_deeplab or is_max_deeplab or is_vip_deeplab)476  if is_panoptic:477    supported_tasks.add(common.TASK_PANOPTIC_SEGMENTATION)478    # MaX-DeepLab does not support evaluating instance segmentation mask AP yet.479    if not is_max_deeplab:480      supported_tasks.add(common.TASK_INSTANCE_SEGMENTATION)481  if is_motion_deeplab or is_vip_deeplab:482    supported_tasks.add(common.TASK_VIDEO_PANOPTIC_SEGMENTATION)483  if is_vip_deeplab:484    supported_tasks.add(common.TASK_DEPTH_AWARE_VIDEO_PANOPTIC_SEGMENTATION)485  return supported_tasks486