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"""AutoAugment utility file.17 18Please cite or refer to the following papers:19- Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le.20"Autoaugment: Learning augmentation policies from data." In CVPR, 2019.21 22- Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le.23"Randaugment: Practical automated data augmentation with a reduced search24space." In CVPR, 2020.25"""26 27import inspect28 29import tensorflow as tf30 31from deeplab2.data.preprocessing import autoaugment_policy32 33 34# This signifies the max integer that the controller RNN could predict for the35# augmentation scheme.36_MAX_LEVEL = 10.37 38 39def blend(image1, image2, factor):40  """Blends image1 and image2 using 'factor'.41 42  Factor can be above 0.0.  A value of 0.0 means only image1 is used.43  A value of 1.0 means only image2 is used.  A value between 0.0 and44  1.0 means we linearly interpolate the pixel values between the two45  images.  A value greater than 1.0 "extrapolates" the difference46  between the two pixel values, and we clip the results to values47  between 0 and 255.48 49  Args:50    image1: An image Tensor of type uint8.51    image2: An image Tensor of type uint8.52    factor: A floating point value above 0.0.53 54  Returns:55    A blended image Tensor of type uint8.56  """57  if factor == 0.0:58    return tf.convert_to_tensor(image1)59  if factor == 1.0:60    return tf.convert_to_tensor(image2)61 62  image1 = tf.cast(image1, tf.float32)63  image2 = tf.cast(image2, tf.float32)64 65  difference = image2 - image166  scaled = factor * difference67 68  # Do addition in float.69  temp = tf.cast(image1, tf.float32) + scaled70 71  # Interpolate72  if factor > 0.0 and factor < 1.0:73    # Interpolation means we always stay within 0 and 255.74    return tf.cast(temp, tf.uint8)75 76  # Extrapolate:77  #78  # We need to clip and then cast.79  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)80 81 82def solarize(image, threshold=128):83  # For each pixel in the image, select the pixel84  # if the value is less than the threshold.85  # Otherwise, subtract 255 from the pixel.86  return tf.where(image < threshold, image, 255 - image)87 88 89def invert(image):90  """Inverts the image pixels."""91  image = tf.convert_to_tensor(image)92  return 255 - image93 94 95def color(image, factor):96  """Equivalent of PIL Color."""97  degenerate = tf.image.grayscale_to_rgb(tf.image.rgb_to_grayscale(image))98  return blend(degenerate, image, factor)99 100 101def contrast(image, factor):102  """Equivalent of PIL Contrast."""103  degenerate = tf.image.rgb_to_grayscale(image)104  # Cast before calling tf.histogram.105  degenerate = tf.cast(degenerate, tf.int32)106 107  # Compute the grayscale histogram, then compute the mean pixel value,108  # and create a constant image size of that value.  Use that as the109  # blending degenerate target of the original image.110  hist = tf.histogram_fixed_width(degenerate, [0, 255], nbins=256)111  mean = tf.reduce_sum(tf.cast(hist, tf.float32)) / 256.0112  degenerate = tf.ones_like(degenerate, dtype=tf.float32) * mean113  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)114  degenerate = tf.image.grayscale_to_rgb(tf.cast(degenerate, tf.uint8))115  return blend(degenerate, image, factor)116 117 118def brightness(image, factor):119  """Equivalent of PIL Brightness."""120  degenerate = tf.zeros_like(image)121  return blend(degenerate, image, factor)122 123 124def posterize(image, bits):125  """Equivalent of PIL Posterize."""126  shift = 8 - bits127  return tf.bitwise.left_shift(tf.bitwise.right_shift(image, shift), shift)128 129 130def autocontrast(image):131  """Implements Autocontrast function from PIL using TF ops.132 133  Args:134    image: A 3D uint8 tensor.135 136  Returns:137    The image after it has had autocontrast applied to it and will be of type138    uint8.139  """140 141  def scale_channel(image):142    """Scale the 2D image using the autocontrast rule."""143    # A possibly cheaper version can be done using cumsum/unique_with_counts144    # over the histogram values, rather than iterating over the entire image.145    # to compute mins and maxes.146    lo = tf.cast(tf.reduce_min(image), tf.float32)147    hi = tf.cast(tf.reduce_max(image), tf.float32)148 149    # Scale the image, making the lowest value 0 and the highest value 255.150    def scale_values(im):151      scale = 255.0 / (hi - lo)152      offset = -lo * scale153      im = tf.cast(im, tf.float32) * scale + offset154      im = tf.clip_by_value(im, 0.0, 255.0)155      return tf.cast(im, tf.uint8)156 157    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)158    return result159 160  # Assumes RGB for now.  Scales each channel independently161  # and then stacks the result.162  s1 = scale_channel(image[:, :, 0])163  s2 = scale_channel(image[:, :, 1])164  s3 = scale_channel(image[:, :, 2])165  image = tf.stack([s1, s2, s3], 2)166  return image167 168 169def sharpness(image, factor):170  """Implements Sharpness function from PIL using TF ops."""171  orig_image = image172  image = tf.cast(image, tf.float32)173  # Make image 4D for conv operation.174  image = tf.expand_dims(image, 0)175  # SMOOTH PIL Kernel.176  kernel = tf.constant(177      [[1, 1, 1], [1, 5, 1], [1, 1, 1]], dtype=tf.float32,178      shape=[3, 3, 1, 1]) / 13.179  # Tile across channel dimension.180  kernel = tf.tile(kernel, [1, 1, 3, 1])181  strides = [1, 1, 1, 1]182  degenerate = tf.nn.depthwise_conv2d(183      image, kernel, strides, padding='VALID', dilations=[1, 1])184  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)185  degenerate = tf.squeeze(tf.cast(degenerate, tf.uint8), [0])186 187  # For the borders of the resulting image, fill in the values of the188  # original image.189  mask = tf.ones_like(degenerate)190  padded_mask = tf.pad(mask, [[1, 1], [1, 1], [0, 0]])191  padded_degenerate = tf.pad(degenerate, [[1, 1], [1, 1], [0, 0]])192  result = tf.where(tf.equal(padded_mask, 1), padded_degenerate, orig_image)193 194  # Blend the final result.195  return blend(result, orig_image, factor)196 197 198def equalize(image):199  """Implements Equalize function from PIL using TF ops."""200  def scale_channel(im, c):201    """Scale the data in the channel to implement equalize."""202    im = tf.cast(im[:, :, c], tf.int32)203    # Compute the histogram of the image channel.204    histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)205 206    # For the purposes of computing the step, filter out the nonzeros.207    nonzero = tf.where(tf.not_equal(histo, 0))208    nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])209    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255210 211    def build_lut(histo, step):212      # Compute the cumulative sum, shifting by step // 2213      # and then normalization by step.214      lut = (tf.cumsum(histo) + (step // 2)) // step215      # Shift lut, prepending with 0.216      lut = tf.concat([[0], lut[:-1]], 0)217      # Clip the counts to be in range.  This is done218      # in the C code for image.point.219      return tf.clip_by_value(lut, 0, 255)220 221    # If step is zero, return the original image.  Otherwise, build222    # lut from the full histogram and step and then index from it.223    result = tf.cond(tf.equal(step, 0),224                     lambda: im,225                     lambda: tf.gather(build_lut(histo, step), im))226 227    return tf.cast(result, tf.uint8)228 229  # Assumes RGB for now.  Scales each channel independently230  # and then stacks the result.231  s1 = scale_channel(image, 0)232  s2 = scale_channel(image, 1)233  s3 = scale_channel(image, 2)234  image = tf.stack([s1, s2, s3], 2)235  return image236 237 238NAME_TO_FUNC = {239    'AutoContrast': autocontrast,240    'Equalize': equalize,241    'Invert': invert,242    'Posterize': posterize,243    'Solarize': solarize,244    'Color': color,245    'Contrast': contrast,246    'Brightness': brightness,247    'Sharpness': sharpness,248}249 250 251def _enhance_level_to_arg(level):252  return ((level/_MAX_LEVEL) * 1.8 + 0.1,)253 254 255def level_to_arg():256  return {257      'AutoContrast':258          lambda level: (),259      'Equalize':260          lambda level: (),261      'Invert':262          lambda level: (),263      'Posterize': lambda level: (int((level/_MAX_LEVEL) * 4),),264      'Solarize': lambda level: (int((level/_MAX_LEVEL) * 256),),265      'Color':266          _enhance_level_to_arg,267      'Contrast':268          _enhance_level_to_arg,269      'Brightness':270          _enhance_level_to_arg,271      'Sharpness':272          _enhance_level_to_arg,273  }274 275 276def label_wrapper(func):277  """Adds a label function argument to func and returns unchanged label."""278  def wrapper(images, label, *args, **kwargs):279    return func(images, *args, **kwargs), label280  return wrapper281 282 283def _parse_policy_info(name, prob, level, replace_value, ignore_label):284  """Returns the function corresponding to `name` and update `level` param."""285  func = NAME_TO_FUNC[name]286  args = level_to_arg()[name](level)287 288  if 'prob' in inspect.getfullargspec(func)[0]:289    args = tuple([prob] + list(args))290 291  # Add in replace arg if it is required for the function that is being called.292  if 'replace' in inspect.getfullargspec(func)[0]:293    # Make sure ignore_label is also in the argument.294    assert 'ignore_label' in inspect.getfullargspec(func)[0]295    # Make sure replace is the second from last argument296    assert 'replace' == inspect.getfullargspec(func)[0][-2]297    # Make sure ignore_label is the final argument298    assert 'ignore_label' == inspect.getfullargspec(func)[0][-1]299    args = tuple(list(args) + [replace_value, ignore_label])300 301  # Add label as the second positional argument for the function if it does302  # not already exist.303  if 'label' not in inspect.getfullargspec(func)[0]:304    func = label_wrapper(func)305  return (func, prob, args)306 307 308def _apply_func_with_prob(func, image, args, prob, label):309  """Apply `func` to image w/ `args` as input with probability `prob`."""310  assert isinstance(args, tuple)311  assert 'label' == inspect.getfullargspec(func)[0][1]312 313  # If prob is a function argument, then this randomness is being handled314  # inside the function, so make sure it is always called.315  if 'prob' in inspect.getfullargspec(func)[0]:316    prob = 1.0317 318  # Apply the function with probability `prob`.319  should_apply_op = tf.cast(320      tf.floor(tf.random.uniform([], dtype=tf.float32) + prob), tf.bool)321  augmented_image, augmented_label = tf.cond(322      should_apply_op,323      lambda: func(image, label, *args),324      lambda: (image, label))325  return augmented_image, augmented_label326 327 328def select_and_apply_random_policy(policies, image, label):329  """Select a random policy from `policies` and apply it to `image`."""330  policy_to_select = tf.random.uniform([], maxval=len(policies), dtype=tf.int32)331  # Note that using tf.case instead of tf.conds would result in significantly332  # larger graphs and would even break export for some larger policies.333  for (i, policy) in enumerate(policies):334    image, label = tf.cond(335        tf.equal(i, policy_to_select),336        lambda selected_policy=policy: selected_policy(image, label),337        lambda: (image, label))338  return (image, label)339 340 341def build_and_apply_autoaugment_policy(policies, image, label, ignore_label):342  """Builds a policy from the given policies passed in and applies to image.343 344  Args:345    policies: list of lists of tuples in the form `(func, prob, level)`, `func`346      is a string name of the augmentation function, `prob` is the probability347      of applying the `func` operation, `level` is the input argument for348      `func`.349    image: tf.Tensor that the resulting policy will be applied to.350    label: tf.Tensor that the resulting policy will be applied to.351    ignore_label: The label value which will be ignored for training and352      evaluation.353 354  Returns:355    A version of image that now has data augmentation applied to it based on356    the `policies` pass into the function. Additionally, returns bboxes if357    a value for them is passed in that is not None358  """359  replace_value = [128, 128, 128]360 361  # func is the string name of the augmentation function, prob is the362  # probability of applying the operation and level is the parameter associated363  # with the tf op.364 365  # tf_policies are functions that take in an image and return an augmented366  # image.367  tf_policies = []368  for policy in policies:369    tf_policy = []370    # Link string name to the correct python function and make sure the correct371    # argument is passed into that function.372    for policy_info in policy:373      policy_info = (374          list(policy_info) + [replace_value, ignore_label])375 376      tf_policy.append(_parse_policy_info(*policy_info))377    # Now build the tf policy that will apply the augmentation procedue378    # on image.379    def make_final_policy(tf_policy_):380      def final_policy(image_, label_):381        for func, prob, args in tf_policy_:382          image_, label_ = _apply_func_with_prob(383              func, image_, args, prob, label_)384        return image_, label_385      return final_policy386    tf_policies.append(make_final_policy(tf_policy))387 388  augmented_images, augmented_label = select_and_apply_random_policy(389      tf_policies, image, label)390  # If no bounding boxes were specified, then just return the images.391  return (augmented_images, augmented_label)392 393 394def distort_image_with_autoaugment(image,395                                   label,396                                   ignore_label,397                                   augmentation_name=None):398  """Applies the AutoAugment policy to `image` and `label`.399 400  Args:401    image: `Tensor` of shape [height, width, 3] representing an image.402    label: `Tensor` of shape [height, width, 1] representing a label.403    ignore_label: The label value which will be ignored for training and404      evaluation.405    augmentation_name: The name of the AutoAugment policy to use. See406      autoaugment_policy.py for available_policies.407 408  Returns:409    A tuple containing the augmented versions of `image` and `label`.410 411  Raises:412    ValueError: If the augmentation_name is not in available_policies.413  """414  if augmentation_name:415    available_policies = autoaugment_policy.available_policies416    if augmentation_name not in available_policies:417      raise ValueError(418          'Invalid augmentation_name: {}'.format(augmentation_name))419    policy = available_policies[augmentation_name]420    return build_and_apply_autoaugment_policy(421        policy, image, label, ignore_label)422  return image, label423