karolmajek/Axial-DeepLab-SWideRNet

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"""Utility functions related to preprocessing inputs."""17 18import numpy as np19import tensorflow as tf20 21 22def flip_dim(tensor_list, prob=0.5, dim=1):23  """Randomly flips a dimension of the given tensor.24 25  The decision to randomly flip the `Tensors` is made together. In other words,26  all or none of the images pass in are flipped.27 28  Note that tf.random_flip_left_right and tf.random_flip_up_down isn't used so29  that we can control for the probability as well as ensure the same decision30  is applied across the images.31 32  Args:33    tensor_list: A list of `Tensors` with the same number of dimensions.34    prob: The probability of a left-right flip.35    dim: The dimension to flip, 0, 1, ..36 37  Returns:38    outputs: A list of the possibly flipped `Tensors` as well as an indicator39    `Tensor` at the end whose value is `True` if the inputs were flipped and40    `False` otherwise.41 42  Raises:43    ValueError: If dim is negative or greater than the dimension of a `Tensor`.44  """45  random_value = tf.random.uniform([])46 47  def flip():48    flipped = []49    for tensor in tensor_list:50      if dim < 0 or dim >= len(tensor.get_shape().as_list()):51        raise ValueError('dim must represent a valid dimension.')52      flipped.append(tf.reverse(tensor, [dim]))53    return flipped54 55  is_flipped = tf.less_equal(random_value, prob)56  outputs = tf.cond(is_flipped, flip, lambda: tensor_list)57  if not isinstance(outputs, (list, tuple)):58    outputs = [outputs]59  outputs.append(is_flipped)60 61  return outputs62 63 64def get_label_resize_method(label):65  """Returns the resize method of labels depending on label dtype.66 67  Args:68    label: Groundtruth label tensor.69 70  Returns:71    tf.image.ResizeMethod.BILINEAR, if label dtype is floating.72    tf.image.ResizeMethod.NEAREST_NEIGHBOR, if label dtype is integer.73 74  Raises:75    ValueError: If label is neither floating nor integer.76  """77  if label.dtype.is_floating:78    return tf.image.ResizeMethod.BILINEAR79  elif label.dtype.is_integer:80    return tf.image.ResizeMethod.NEAREST_NEIGHBOR81  else:82    raise ValueError('Label type must be either floating or integer.')83 84 85def _crop(image, offset_height, offset_width, crop_height, crop_width):86  """Crops the given image using the provided offsets and sizes.87 88  Note that the method doesn't assume we know the input image size but it does89  assume we know the input image rank.90 91  Args:92    image: an image of shape [height, width, channels].93    offset_height: a scalar tensor indicating the height offset.94    offset_width: a scalar tensor indicating the width offset.95    crop_height: the height of the cropped image.96    crop_width: the width of the cropped image.97 98  Returns:99    The cropped (and resized) image.100 101  Raises:102    ValueError: if `image` doesn't have rank of 3.103    InvalidArgumentError: if the rank is not 3 or if the image dimensions are104      less than the crop size.105  """106  original_shape = tf.shape(image)107 108  if len(image.get_shape().as_list()) != 3:109    raise ValueError('input must have rank of 3')110  original_channels = image.get_shape().as_list()[2]111 112  rank_assertion = tf.Assert(113      tf.equal(tf.rank(image), 3),114      ['Rank of image must be equal to 3.'])115  with tf.control_dependencies([rank_assertion]):116    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])117 118  size_assertion = tf.Assert(119      tf.logical_and(120          tf.greater_equal(original_shape[0], crop_height),121          tf.greater_equal(original_shape[1], crop_width)),122      ['Crop size greater than the image size.'])123 124  offsets = tf.cast(tf.stack([offset_height, offset_width, 0]), tf.int32)125 126  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to127  # define the crop size.128  with tf.control_dependencies([size_assertion]):129    image = tf.slice(image, offsets, cropped_shape)130  image = tf.reshape(image, cropped_shape)131  image.set_shape([crop_height, crop_width, original_channels])132  return image133 134 135def random_crop(image_list, crop_height, crop_width):136  """Crops the given list of images.137 138  The function applies the same crop to each image in the list. This can be139  effectively applied when there are multiple image inputs of the same140  dimension such as:141 142    image, depths, normals = random_crop([image, depths, normals], 120, 150)143 144  Args:145    image_list: a list of image tensors of the same dimension but possibly146      varying channel.147    crop_height: the new height.148    crop_width: the new width.149 150  Returns:151    the image_list with cropped images.152 153  Raises:154    ValueError: if there are multiple image inputs provided with different size155      or the images are smaller than the crop dimensions.156  """157  if not image_list:158    raise ValueError('Empty image_list.')159 160  # Compute the rank assertions.161  rank_assertions = []162  for i in range(len(image_list)):163    image_rank = tf.rank(image_list[i])164    rank_assert = tf.Assert(165        tf.equal(image_rank, 3), [166            'Wrong rank for tensor %d in image_list [expected] [actual]', i, 3,167            image_rank168        ])169    rank_assertions.append(rank_assert)170 171  with tf.control_dependencies([rank_assertions[0]]):172    image_shape = tf.shape(image_list[0])173  image_height = image_shape[0]174  image_width = image_shape[1]175  crop_size_assert = tf.Assert(176      tf.logical_and(177          tf.greater_equal(image_height, crop_height),178          tf.greater_equal(image_width, crop_width)),179      ['Crop size greater than the image size.'])180 181  asserts = [rank_assertions[0], crop_size_assert]182 183  for i in range(1, len(image_list)):184    image = image_list[i]185    asserts.append(rank_assertions[i])186    with tf.control_dependencies([rank_assertions[i]]):187      shape = tf.shape(image)188    height = shape[0]189    width = shape[1]190 191    height_assert = tf.Assert(192        tf.equal(height, image_height), [193            'Wrong height for tensor %d in image_list [expected][actual]', i,194            height, image_height195        ])196    width_assert = tf.Assert(197        tf.equal(width, image_width), [198            'Wrong width for tensor %d in image_list [expected][actual]', i,199            width, image_width200        ])201    asserts.extend([height_assert, width_assert])202 203  # Create a random bounding box.204  #205  # Use tf.random.uniform and not numpy.random.rand as doing the former would206  # generate random numbers at graph eval time, unlike the latter which207  # generates random numbers at graph definition time.208  with tf.control_dependencies(asserts):209    max_offset_height = tf.reshape(image_height - crop_height + 1, [])210    max_offset_width = tf.reshape(image_width - crop_width + 1, [])211  offset_height = tf.random.uniform([],212                                    maxval=max_offset_height,213                                    dtype=tf.int32)214  offset_width = tf.random.uniform([], maxval=max_offset_width, dtype=tf.int32)215 216  return [_crop(image, offset_height, offset_width,217                crop_height, crop_width) for image in image_list]218 219 220def get_random_scale(min_scale_factor, max_scale_factor, step_size):221  """Gets a random scale value.222 223  Args:224    min_scale_factor: Minimum scale value.225    max_scale_factor: Maximum scale value.226    step_size: The step size from minimum to maximum value.227 228  Returns:229    A tensor with random scale value selected between minimum and maximum value.230    If `min_scale_factor` and `max_scale_factor` are the same, a number is231    returned instead.232 233  Raises:234    ValueError: min_scale_factor has unexpected value.235  """236  if min_scale_factor < 0 or min_scale_factor > max_scale_factor:237    raise ValueError('Unexpected value of min_scale_factor.')238 239  if min_scale_factor == max_scale_factor:240    return np.float32(min_scale_factor)241 242  # When step_size = 0, we sample the value uniformly from [min, max).243  if step_size == 0:244    return tf.random.uniform([1],245                             minval=min_scale_factor,246                             maxval=max_scale_factor)247 248  # When step_size != 0, we randomly select one discrete value from [min, max].249  num_steps = int((max_scale_factor - min_scale_factor) / step_size + 1)250  scale_factors = tf.linspace(min_scale_factor, max_scale_factor, num_steps)251  shuffled_scale_factors = tf.random.shuffle(scale_factors)252  return shuffled_scale_factors[0]253 254 255def randomly_scale_image_and_label(image, label=None, scale=1.0):256  """Randomly scales image and label.257 258  Args:259    image: Image with shape [height, width, 3].260    label: Label with shape [height, width, 1].261    scale: The value to scale image and label.262 263  Returns:264    Scaled image and label.265  """266  # No random scaling if scale == 1.267  if scale == 1.0:268    return image, label269  image_shape = tf.shape(image)270  new_dim = tf.cast(271      tf.cast([image_shape[0], image_shape[1]], tf.float32) * scale,272      tf.int32)273 274  # Need squeeze and expand_dims because image interpolation takes275  # 4D tensors as input.276  image = tf.squeeze(277      tf.compat.v1.image.resize_bilinear(278          tf.expand_dims(image, 0), new_dim, align_corners=True), [0])279  if label is not None:280    label = tf.compat.v1.image.resize(281        label,282        new_dim,283        method=get_label_resize_method(label),284        align_corners=True)285 286  return image, label287 288 289def resolve_shape(tensor, rank=None):290  """Fully resolves the shape of a Tensor.291 292  Use as much as possible the shape components already known during graph293  creation and resolve the remaining ones during runtime.294 295  Args:296    tensor: Input tensor whose shape we query.297    rank: The rank of the tensor, provided that we know it.298 299  Returns:300    shape: The full shape of the tensor.301  """302  if rank is not None:303    shape = tensor.get_shape().with_rank(rank).as_list()304  else:305    shape = tensor.get_shape().as_list()306 307  if None in shape:308    dynamic_shape = tf.shape(tensor)309    for i in range(len(shape)):310      if shape[i] is None:311        shape[i] = dynamic_shape[i]312 313  return shape314 315 316def _scale_dim(original_size, factor):317  """Helper method to scale one input dimension by the given factor."""318  original_size = tf.cast(original_size, tf.float32)319  factor = tf.cast(factor, tf.float32)320  return tf.cast(tf.floor(original_size * factor), tf.int32)321 322 323def process_resize_value(resize_spec):324  """Helper method to process input resize spec.325 326  Args:327    resize_spec: Either None, a python scalar, or a sequence with length <=2.328      Each value in the sequence should be a python integer.329 330  Returns:331    None if input size is not valid, or 2-tuple of (height, width), derived332      from input resize_spec.333  """334  if not resize_spec:335    return None336 337  if isinstance(resize_spec, int):338    # For conveniences and also backward compatibility.339    resize_spec = (resize_spec,)340 341  resize_spec = tuple(resize_spec)342 343  if len(resize_spec) == 1:344    resize_spec = (resize_spec[0], resize_spec[0])345 346  if len(resize_spec) != 2:347    raise ValueError('Unable to process input resize_spec: %s' % resize_spec)348 349  if resize_spec[0] <= 0 or resize_spec[1] <= 0:350    return None351 352  return resize_spec353 354 355def _resize_to_match_min_size(input_shape, min_size):356  """Returns the resized shape so that both sides match minimum size.357 358  Note: the input image will still be scaled if input height and width359  are already greater than minimum size.360 361  Args:362    input_shape: A 2-tuple, (height, width) of the input image. Each value can363      be either a python integer or a integer scalar tensor.364    min_size: A tuple of (minimum height, minimum width) to specify the365      minimum shape after resize. The input shape would be scaled so that both366      height and width will be greater than or equal to their minimum value.367 368  Returns:369    A 2-tuple, (height, width), resized input shape which preserves input370      aspect ratio.371  """372  input_height, input_width = input_shape373  min_height, min_width = min_size374 375  scale_factor = tf.maximum(min_height / input_height, min_width / input_width)376  return (_scale_dim(input_height, scale_factor),377          _scale_dim(input_width, scale_factor))378 379 380def _resize_to_fit_max_size(input_shape, max_size):381  """Returns the resized shape so that both sides fit within max size.382 383  Note: if input shape is already smaller or equal to maximum size, no resize384    operation would be performed.385 386  Args:387    input_shape: A 2-tuple, (height, width) of the input image. Each value can388      be either a python integer or a integer scalar tensor.389    max_size: A tuple of (minimum height, minimum width) to specify390      the maximum allowed shape after resize.391 392  Returns:393    A 2-tuple, (height, width), resized input shape which preserves input394      aspect ratio.395  """396  input_height, input_width = input_shape397  max_height, max_width = max_size398  scale_factor = tf.minimum(max_height / input_height, max_width / input_width)399 400  scale_factor = tf.minimum(tf.cast(scale_factor, tf.float32),401                            tf.cast(1.0, tf.float32))402  return (_scale_dim(input_height, scale_factor),403          _scale_dim(input_width, scale_factor))404 405 406def resize_to_range_helper(input_shape, min_size, max_size=None, factor=None):407  """Determines output size in specified range.408 409  The output size (height and/or width) can be described by two cases:410  1. If current side can be rescaled so its minimum size is equal to min_size411     without the other side exceeding its max_size, then do so.412  2. Otherwise, resize so at least one side is reaching its max_size.413 414  An integer in `range(factor)` is added to the computed sides so that the415  final dimensions are multiples of `factor` plus one.416 417  Args:418    input_shape: A 2-tuple, (height, width) of the input image. Each value can419      be either a python integer or a integer scalar tensor.420    min_size: A 2-tuple of (height, width), desired minimum value after resize.421      If a single element is given, then height and width share the same422      min_size. None, empty or having 0 indicates no minimum value will be used.423    max_size: A 2-tuple of (height, width), maximum allowed value after resize.424      If a single element is given, then height and width share the same425      max_size. None, empty or having 0 indicates no maximum value will be used.426      Note that the output dimension is no larger than max_size and may be427      slightly smaller than max_size when factor is not None.428    factor: None or integer, make output size multiple of factor plus one.429 430  Returns:431    A 1-D tensor containing the [new_height, new_width].432  """433  output_shape = input_shape434 435  min_size = process_resize_value(min_size)436  if min_size:437    output_shape = _resize_to_match_min_size(input_shape, min_size)438 439  max_size = process_resize_value(max_size)440  if max_size:441    if factor:442      # Update max_size to be a multiple of factor plus 1 and make sure the443      # max dimension after resizing is no larger than max_size.444      max_size = (max_size[0] - (max_size[0] - 1) % factor,445                  max_size[1] - (max_size[1] - 1) % factor)446 447    output_shape = _resize_to_fit_max_size(output_shape, max_size)448 449  output_shape = tf.stack(output_shape)450  # Ensure that both output sides are multiples of factor plus one.451  if factor:452    output_shape += (factor - (output_shape - 1) % factor) % factor453 454  return output_shape455 456 457def resize_to_range(image,458                    label=None,459                    min_size=None,460                    max_size=None,461                    factor=None,462                    align_corners=True,463                    method=tf.image.ResizeMethod.BILINEAR):464  """Resizes image or label so their sides are within the provided range.465 466  The output size (height and/or width) can be described by two cases:467  1. If current side can be rescaled so its minimum size is equal to min_size468     without the other side exceeding its max_size, then do so.469  2. Otherwise, resize so at least one side is reaching its max_size.470 471  An integer in `range(factor)` is added to the computed sides so that the472  final dimensions are multiples of `factor` plus one.473 474  Args:475    image: A 3D tensor of shape [height, width, channels].476    label: (optional) A 3D tensor of shape [height, width, channels].477    min_size: A 2-tuple of (height, width), desired minimum value after resize.478      If a single element is given, then height and width share the same479      min_size. None, empty or having 0 indicates no minimum value will be used.480    max_size: A 2-tuple of (height, width), maximum allowed value after resize.481      If a single element is given, then height and width share the same482      max_size. None, empty or having 0 indicates no maximum value will be used.483      Note that the output dimension is no larger than max_size and may be484      slightly smaller than max_size when factor is not None.485    factor: Make output size multiple of factor plus one.486    align_corners: If True, exactly align all 4 corners of input and output.487    method: Image resize method. Defaults to tf.image.ResizeMethod.BILINEAR.488 489  Returns:490    resized_image: A 3-D tensor of shape [new_height, new_width, channels],491      where the image has been resized with the specified method.492    resized_label: Either None (if input label is None) or a 3-D tensor,493      where the input label has been resized accordingly.494 495  Raises:496    ValueError: If the image is not a 3D tensor.497  """498  orig_height, orig_width, _ = resolve_shape(image, rank=3)499  new_size = resize_to_range_helper(input_shape=(orig_height, orig_width),500                                    min_size=min_size,501                                    max_size=max_size,502                                    factor=factor)503 504  resized_image = tf.compat.v1.image.resize(505      image, new_size, method=method, align_corners=align_corners)506 507  if label is None:508    return resized_image, None509 510  resized_label = tf.compat.v1.image.resize(511      label,512      new_size,513      method=get_label_resize_method(label),514      align_corners=align_corners)515 516  return resized_image, resized_label517