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"""This file contains functions to preprocess images and labels."""17 18import tensorflow as tf19 20from deeplab2.data.preprocessing import autoaugment_utils21from deeplab2.data.preprocessing import preprocess_utils22 23# The probability of flipping the images and labels24# left-right during training25_PROB_OF_FLIP = 0.526 27_MEAN_PIXEL = [127.5, 127.5, 127.5]28 29 30def _pad_image_and_label(image, label, offset_height, offset_width,31                         target_height, target_width, ignore_label=None):32  """Pads the image and the label to the given size.33 34  Args:35    image: A tf.Tensor of shape [height, width, channels].36    label: A tf.Tensor of shape [height, width, 1] or None.37    offset_height: The number of rows of zeros to add on top of the image and38      label.39    offset_width: The number of columns of zeros to add on the left of the image40      and label.41    target_height: The total height after padding.42    target_width: The total width after padding.43    ignore_label: The ignore_label for the label. Must only be set when label is44      given.45 46  Returns:47    The padded image and label as a tuple (padded_image, padded_label).48 49  Raises:50    tf.errors.InvalidArgumentError: An error occurs if the padding configuration51      is invalid.52    ValueError: An error occurs if label is given without an ignore_label.53  """54  height = tf.shape(image)[0]55  width = tf.shape(image)[1]56  original_dtype = image.dtype57  if original_dtype not in (tf.float32, tf.float64):58    image = tf.cast(image, tf.float32)59 60  bottom_padding = target_height - offset_height - height61  right_padding = target_width - offset_width - width62 63  assert_bottom_padding = tf.assert_greater(64      bottom_padding, -1,65      'The padding configuration is not valid. Please either increase the '66      'target size or reduce the padding offset.')67  assert_right_padding = tf.assert_greater(68      right_padding, -1, 'The padding configuration is not valid. Please either'69      ' increase the target size or reduce the padding offset.')70  with tf.control_dependencies([assert_bottom_padding, assert_right_padding]):71    paddings = [[offset_height, bottom_padding], [offset_width, right_padding],72                [0, 0]]73 74    image = image - _MEAN_PIXEL75    image = tf.pad(image, paddings)76    image = image + _MEAN_PIXEL77    image = tf.cast(image, original_dtype)78 79    if label is not None:80      if ignore_label is None:81        raise ValueError(82            'If a label is given, the ignore label must be set too.')83      label = tf.pad(label, paddings, constant_values=ignore_label)84 85    return image, label86 87 88def _update_max_resize_value(max_resize_value, crop_size, is_inference=False):89  """Checks and may update max_resize_value.90 91  Args:92    max_resize_value: A 2-tuple of (height, width), maximum allowed value93      after resize. If a single element is given, then height and width94      share the same value. None, empty or having 0 indicates no maximum value95      will be used.96    crop_size: A 2-tuple of (height, width), crop size used.97    is_inference: Boolean, whether the model is performing inference or not.98 99  Returns:100    Updated max_resize_value.101  """102  max_resize_value = preprocess_utils.process_resize_value(max_resize_value)103  if max_resize_value is None and is_inference:104    # During inference, default max_resize_value to crop size to allow105    # model taking input images with larger sizes.106    max_resize_value = crop_size107 108  if max_resize_value is None:109    return None110 111  if max_resize_value[0] > crop_size[0] or max_resize_value[1] > crop_size[1]:112    raise ValueError(113        'Maximum resize value provided (%s) exceeds model crop size (%s)' %114        (max_resize_value, crop_size))115  return max_resize_value116 117 118def preprocess_image_and_label(image,119                               label,120                               crop_height,121                               crop_width,122                               prev_image=None,123                               prev_label=None,124                               min_resize_value=None,125                               max_resize_value=None,126                               resize_factor=None,127                               min_scale_factor=1.,128                               max_scale_factor=1.,129                               scale_factor_step_size=0,130                               ignore_label=None,131                               is_training=True,132                               autoaugment_policy_name=None):133  """Preprocesses the image and label.134 135  Args:136    image: A tf.Tensor containing the image with shape [height, width, 3].137    label: A tf.Tensor containing the label with shape [height, width, 1] or138      None.139    crop_height: The height value used to crop the image and label.140    crop_width: The width value used to crop the image and label.141    prev_image: An optional tensor of shape [image_height, image_width, 3].142    prev_label: An optional tensor of shape [label_height, label_width, 1].143    min_resize_value: A 2-tuple of (height, width), desired minimum value144      after resize. If a single element is given, then height and width share145      the same value. None, empty or having 0 indicates no minimum value will146      be used.147    max_resize_value: A 2-tuple of (height, width), maximum allowed value148      after resize. If a single element is given, then height and width149      share the same value. None, empty or having 0 indicates no maximum value150      will be used.151    resize_factor: Resized dimensions are multiple of factor plus one.152    min_scale_factor: Minimum scale factor for random scale augmentation.153    max_scale_factor: Maximum scale factor for random scale augmentation.154    scale_factor_step_size: The step size from min scale factor to max scale155      factor. The input is randomly scaled based on the value of156      (min_scale_factor, max_scale_factor, scale_factor_step_size).157    ignore_label: The label value which will be ignored for training and158      evaluation.159    is_training: If the preprocessing is used for training or not.160    autoaugment_policy_name: String, autoaugment policy name. See161        autoaugment_policy.py for available policies.162 163  Returns:164    resized_image: The resized input image without other augmentations as a165      tf.Tensor.166    processed_image: The preprocessed image as a tf.Tensor.167    label: The preprocessed groundtruth segmentation label as a tf.Tensor.168 169  Raises:170    ValueError: Ground truth label not provided during training.171  """172  if is_training and label is None:173    raise ValueError('During training, label must be provided.')174 175  image.get_shape().assert_is_compatible_with(tf.TensorShape([None, None, 3]))176 177  # Keep reference to original image.178  resized_image = image179  if prev_image is not None:180    image = tf.concat([image, prev_image], axis=2)181  processed_image = tf.cast(image, tf.float32)182  processed_prev_image = None183 184  if label is not None:185    label.get_shape().assert_is_compatible_with(tf.TensorShape([None, None, 1]))186    if prev_label is not None:187      label = tf.concat([label, prev_label], axis=2)188    label = tf.cast(label, tf.int32)189 190  # Resize image and label to the desired range.191  if any([min_resize_value, max_resize_value, not is_training]):192    max_resize_value = _update_max_resize_value(193        max_resize_value,194        crop_size=(crop_height, crop_width),195        is_inference=not is_training)196 197    processed_image, label = (198        preprocess_utils.resize_to_range(199            image=processed_image,200            label=label,201            min_size=min_resize_value,202            max_size=max_resize_value,203            factor=resize_factor,204            align_corners=True))205    if prev_image is None:206      resized_image = tf.identity(processed_image)207    else:208      resized_image, _ = tf.split(processed_image, 2, axis=2)209 210  if prev_image is not None:211    processed_image, processed_prev_image = tf.split(processed_image, 2, axis=2)212 213  if prev_label is not None:214    label, prev_label = tf.split(label, 2, axis=2)215 216  if not is_training:217    image_height = tf.shape(processed_image)[0]218    image_width = tf.shape(processed_image)[1]219 220    offset_height = 0221    offset_width = 0222    processed_image, label = _pad_image_and_label(processed_image, label,223                                                  offset_height, offset_width,224                                                  crop_height, crop_width,225                                                  ignore_label)226    processed_image.set_shape([crop_height, crop_width, 3])227    if label is not None:228      label.set_shape([crop_height, crop_width, 1])229    if prev_image is not None:230      processed_prev_image, prev_label = _pad_image_and_label(231          processed_prev_image, prev_label, offset_height, offset_width,232          crop_height, crop_width, ignore_label)233      processed_prev_image.set_shape([crop_height, crop_width, 3])234      if prev_label is not None:235        prev_label.set_shape([crop_height, crop_width, 1])236    return (resized_image, processed_image, label, processed_prev_image,237            prev_label)238 239  # Data augmentation by randomly scaling the inputs.240  scale = preprocess_utils.get_random_scale(241      min_scale_factor, max_scale_factor, scale_factor_step_size)242  processed_image, label = preprocess_utils.randomly_scale_image_and_label(243      processed_image, label, scale)244  if processed_prev_image is not None:245    (processed_prev_image,246     prev_label) = preprocess_utils.randomly_scale_image_and_label(247         processed_prev_image, prev_label, scale)248 249  # Apply autoaugment if any.250  if autoaugment_policy_name:251    processed_image, label = _autoaugment_helper(252        processed_image, label, ignore_label, autoaugment_policy_name)253    if processed_prev_image is not None:254      processed_prev_image, prev_label = _autoaugment_helper(255          processed_prev_image, prev_label, ignore_label,256          autoaugment_policy_name)257 258  # Pad image and label to have dimensions >= [crop_height, crop_width].259  image_height = tf.shape(processed_image)[0]260  image_width = tf.shape(processed_image)[1]261  target_height = image_height + tf.maximum(crop_height - image_height, 0)262  target_width = image_width + tf.maximum(crop_width - image_width, 0)263 264  # Randomly crop the image and label.265  def _uniform_offset(margin):266    return tf.random.uniform(267        [], minval=0, maxval=tf.maximum(margin, 1), dtype=tf.int32)268 269  offset_height = _uniform_offset(crop_height - image_height)270  offset_width = _uniform_offset(crop_width - image_width)271  processed_image, label = _pad_image_and_label(processed_image, label,272                                                offset_height, offset_width,273                                                target_height, target_width,274                                                ignore_label)275  if processed_prev_image is not None:276    processed_prev_image, prev_label = _pad_image_and_label(277        processed_prev_image, prev_label, offset_height, offset_width,278        target_height, target_width, ignore_label)279 280  if processed_prev_image is not None:281    (processed_image, label, processed_prev_image,282     prev_label) = preprocess_utils.random_crop(283         [processed_image, label, processed_prev_image, prev_label],284         crop_height, crop_width)285    # Randomly left-right flip the image and label.286    (processed_image, label, processed_prev_image, prev_label,287     _) = preprocess_utils.flip_dim(288         [processed_image, label, processed_prev_image, prev_label],289         _PROB_OF_FLIP,290         dim=1)291  else:292    processed_image, label = preprocess_utils.random_crop(293        [processed_image, label], crop_height, crop_width)294    # Randomly left-right flip the image and label.295    processed_image, label, _ = preprocess_utils.flip_dim(296        [processed_image, label], _PROB_OF_FLIP, dim=1)297 298  return resized_image, processed_image, label, processed_prev_image, prev_label299 300 301def _autoaugment_helper(image, label, ignore_label, policy_name):302  image = tf.cast(image, tf.uint8)303  label = tf.cast(label, tf.int32)304  image, label = autoaugment_utils.distort_image_with_autoaugment(305      image, label, ignore_label, policy_name)306  image = tf.cast(image, tf.float32)307  return image, label308