paddlex 2.0.0rc4__py3-none-any.whl → 3.0.0b2__py3-none-any.whl

paddlex/ppdet/data/transform/autoaugment_utils.py

@@ -1,1593 +0,0 @@
-# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# Reference:
-#   https://github.com/tensorflow/tpu/blob/master/models/official/detection/utils/autoaugment_utils.py
-"""AutoAugment util file."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import inspect
-import math
-from PIL import Image, ImageEnhance
-import numpy as np
-import os
-import sys
-import cv2
-from copy import deepcopy
-
-# This signifies the max integer that the controller RNN could predict for the
-# augmentation scheme.
-_MAX_LEVEL = 10.
-
-# Represents an invalid bounding box that is used for checking for padding
-# lists of bounding box coordinates for a few augmentation operations
-_INVALID_BOX = [[-1.0, -1.0, -1.0, -1.0]]
-
-
-def policy_v0():
-    """Autoaugment policy that was used in AutoAugment Detection Paper."""
-    # Each tuple is an augmentation operation of the form
-    # (operation, probability, magnitude). Each element in policy is a
-    # sub-policy that will be applied sequentially on the image.
-    policy = [
-        [('TranslateX_BBox', 0.6, 4), ('Equalize', 0.8, 10)],
-        [('TranslateY_Only_BBoxes', 0.2, 2), ('Cutout', 0.8, 8)],
-        [('Sharpness', 0.0, 8), ('ShearX_BBox', 0.4, 0)],
-        [('ShearY_BBox', 1.0, 2), ('TranslateY_Only_BBoxes', 0.6, 6)],
-        [('Rotate_BBox', 0.6, 10), ('Color', 1.0, 6)],
-    ]
-    return policy
-
-
-def policy_v1():
-    """Autoaugment policy that was used in AutoAugment Detection Paper."""
-    # Each tuple is an augmentation operation of the form
-    # (operation, probability, magnitude). Each element in policy is a
-    # sub-policy that will be applied sequentially on the image.
-    policy = [
-        [('TranslateX_BBox', 0.6, 4), ('Equalize', 0.8, 10)],
-        [('TranslateY_Only_BBoxes', 0.2, 2), ('Cutout', 0.8, 8)],
-        [('Sharpness', 0.0, 8), ('ShearX_BBox', 0.4, 0)],
-        [('ShearY_BBox', 1.0, 2), ('TranslateY_Only_BBoxes', 0.6, 6)],
-        [('Rotate_BBox', 0.6, 10), ('Color', 1.0, 6)],
-        [('Color', 0.0, 0), ('ShearX_Only_BBoxes', 0.8, 4)],
-        [('ShearY_Only_BBoxes', 0.8, 2), ('Flip_Only_BBoxes', 0.0, 10)],
-        [('Equalize', 0.6, 10), ('TranslateX_BBox', 0.2, 2)],
-        [('Color', 1.0, 10), ('TranslateY_Only_BBoxes', 0.4, 6)],
-        [('Rotate_BBox', 0.8, 10), ('Contrast', 0.0, 10)],  # ,
-        [('Cutout', 0.2, 2), ('Brightness', 0.8, 10)],
-        [('Color', 1.0, 6), ('Equalize', 1.0, 2)],
-        [('Cutout_Only_BBoxes', 0.4, 6), ('TranslateY_Only_BBoxes', 0.8, 2)],
-        [('Color', 0.2, 8), ('Rotate_BBox', 0.8, 10)],
-        [('Sharpness', 0.4, 4), ('TranslateY_Only_BBoxes', 0.0, 4)],
-        [('Sharpness', 1.0, 4), ('SolarizeAdd', 0.4, 4)],
-        [('Rotate_BBox', 1.0, 8), ('Sharpness', 0.2, 8)],
-        [('ShearY_BBox', 0.6, 10), ('Equalize_Only_BBoxes', 0.6, 8)],
-        [('ShearX_BBox', 0.2, 6), ('TranslateY_Only_BBoxes', 0.2, 10)],
-        [('SolarizeAdd', 0.6, 8), ('Brightness', 0.8, 10)],
-    ]
-    return policy
-
-
-def policy_vtest():
-    """Autoaugment test policy for debugging."""
-    # Each tuple is an augmentation operation of the form
-    # (operation, probability, magnitude). Each element in policy is a
-    # sub-policy that will be applied sequentially on the image.
-    policy = [[('TranslateX_BBox', 1.0, 4), ('Equalize', 1.0, 10)], ]
-    return policy
-
-
-def policy_v2():
-    """Additional policy that performs well on object detection."""
-    # Each tuple is an augmentation operation of the form
-    # (operation, probability, magnitude). Each element in policy is a
-    # sub-policy that will be applied sequentially on the image.
-    policy = [
-        [('Color', 0.0, 6), ('Cutout', 0.6, 8), ('Sharpness', 0.4, 8)],
-        [('Rotate_BBox', 0.4, 8), ('Sharpness', 0.4, 2),
-         ('Rotate_BBox', 0.8, 10)],
-        [('TranslateY_BBox', 1.0, 8), ('AutoContrast', 0.8, 2)],
-        [('AutoContrast', 0.4, 6), ('ShearX_BBox', 0.8, 8),
-         ('Brightness', 0.0, 10)],
-        [('SolarizeAdd', 0.2, 6), ('Contrast', 0.0, 10),
-         ('AutoContrast', 0.6, 0)],
-        [('Cutout', 0.2, 0), ('Solarize', 0.8, 8), ('Color', 1.0, 4)],
-        [('TranslateY_BBox', 0.0, 4), ('Equalize', 0.6, 8),
-         ('Solarize', 0.0, 10)],
-        [('TranslateY_BBox', 0.2, 2), ('ShearY_BBox', 0.8, 8),
-         ('Rotate_BBox', 0.8, 8)],
-        [('Cutout', 0.8, 8), ('Brightness', 0.8, 8), ('Cutout', 0.2, 2)],
-        [('Color', 0.8, 4), ('TranslateY_BBox', 1.0, 6),
-         ('Rotate_BBox', 0.6, 6)],
-        [('Rotate_BBox', 0.6, 10), ('BBox_Cutout', 1.0, 4), ('Cutout', 0.2, 8)
-         ],
-        [('Rotate_BBox', 0.0, 0), ('Equalize', 0.6, 6),
-         ('ShearY_BBox', 0.6, 8)],
-        [('Brightness', 0.8, 8), ('AutoContrast', 0.4, 2),
-         ('Brightness', 0.2, 2)],
-        [('TranslateY_BBox', 0.4, 8), ('Solarize', 0.4, 6),
-         ('SolarizeAdd', 0.2, 10)],
-        [('Contrast', 1.0, 10), ('SolarizeAdd', 0.2, 8), ('Equalize', 0.2, 4)],
-    ]
-    return policy
-
-
-def policy_v3():
-    """"Additional policy that performs well on object detection."""
-    # Each tuple is an augmentation operation of the form
-    # (operation, probability, magnitude). Each element in policy is a
-    # sub-policy that will be applied sequentially on the image.
-    policy = [
-        [('Posterize', 0.8, 2), ('TranslateX_BBox', 1.0, 8)],
-        [('BBox_Cutout', 0.2, 10), ('Sharpness', 1.0, 8)],
-        [('Rotate_BBox', 0.6, 8), ('Rotate_BBox', 0.8, 10)],
-        [('Equalize', 0.8, 10), ('AutoContrast', 0.2, 10)],
-        [('SolarizeAdd', 0.2, 2), ('TranslateY_BBox', 0.2, 8)],
-        [('Sharpness', 0.0, 2), ('Color', 0.4, 8)],
-        [('Equalize', 1.0, 8), ('TranslateY_BBox', 1.0, 8)],
-        [('Posterize', 0.6, 2), ('Rotate_BBox', 0.0, 10)],
-        [('AutoContrast', 0.6, 0), ('Rotate_BBox', 1.0, 6)],
-        [('Equalize', 0.0, 4), ('Cutout', 0.8, 10)],
-        [('Brightness', 1.0, 2), ('TranslateY_BBox', 1.0, 6)],
-        [('Contrast', 0.0, 2), ('ShearY_BBox', 0.8, 0)],
-        [('AutoContrast', 0.8, 10), ('Contrast', 0.2, 10)],
-        [('Rotate_BBox', 1.0, 10), ('Cutout', 1.0, 10)],
-        [('SolarizeAdd', 0.8, 6), ('Equalize', 0.8, 8)],
-    ]
-    return policy
-
-
-def _equal(val1, val2, eps=1e-8):
-    return abs(val1 - val2) <= eps
-
-
-def blend(image1, image2, factor):
-    """Blend image1 and image2 using 'factor'.
-
-    Factor can be above 0.0.    A value of 0.0 means only image1 is used.
-    A value of 1.0 means only image2 is used.    A value between 0.0 and
-    1.0 means we linearly interpolate the pixel values between the two
-    images.    A value greater than 1.0 "extrapolates" the difference
-    between the two pixel values, and we clip the results to values
-    between 0 and 255.
-
-    Args:
-        image1: An image Tensor of type uint8.
-        image2: An image Tensor of type uint8.
-        factor: A floating point value above 0.0.
-
-    Returns:
-        A blended image Tensor of type uint8.
-    """
-    if factor == 0.0:
-        return image1
-    if factor == 1.0:
-        return image2
-
-    image1 = image1.astype(np.float32)
-    image2 = image2.astype(np.float32)
-
-    difference = image2 - image1
-    scaled = factor * difference
-
-    # Do addition in float.
-    temp = image1 + scaled
-
-    # Interpolate
-    if factor > 0.0 and factor < 1.0:
-        # Interpolation means we always stay within 0 and 255.
-        return temp.astype(np.uint8)
-
-    # Extrapolate:
-    #
-    # We need to clip and then cast.
-    return np.clip(temp, a_min=0, a_max=255).astype(np.uint8)
-
-
-def cutout(image, pad_size, replace=0):
-    """Apply cutout (https://arxiv.org/abs/1708.04552) to image.
-
-    This operation applies a (2*pad_size x 2*pad_size) mask of zeros to
-    a random location within `img`. The pixel values filled in will be of the
-    value `replace`. The located where the mask will be applied is randomly
-    chosen uniformly over the whole image.
-
-    Args:
-        image: An image Tensor of type uint8.
-        pad_size: Specifies how big the zero mask that will be generated is that
-            is applied to the image. The mask will be of size
-            (2*pad_size x 2*pad_size).
-        replace: What pixel value to fill in the image in the area that has
-            the cutout mask applied to it.
-
-    Returns:
-        An image Tensor that is of type uint8.
-    Example:
-        img = cv2.imread( "/home/vis/gry/train/img_data/test.jpg", cv2.COLOR_BGR2RGB )
-        new_img = cutout(img, pad_size=50, replace=0)
-    """
-    image_height, image_width = image.shape[0], image.shape[1]
-
-    cutout_center_height = np.random.randint(low=0, high=image_height)
-    cutout_center_width = np.random.randint(low=0, high=image_width)
-
-    lower_pad = np.maximum(0, cutout_center_height - pad_size)
-    upper_pad = np.maximum(0, image_height - cutout_center_height - pad_size)
-    left_pad = np.maximum(0, cutout_center_width - pad_size)
-    right_pad = np.maximum(0, image_width - cutout_center_width - pad_size)
-
-    cutout_shape = [
-        image_height - (lower_pad + upper_pad),
-        image_width - (left_pad + right_pad)
-    ]
-    padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]
-    mask = np.pad(np.zeros(
-        cutout_shape, dtype=image.dtype),
-                  padding_dims,
-                  'constant',
-                  constant_values=1)
-    mask = np.expand_dims(mask, -1)
-    mask = np.tile(mask, [1, 1, 3])
-    image = np.where(
-        np.equal(mask, 0),
-        np.ones_like(
-            image, dtype=image.dtype) * replace,
-        image)
-    return image.astype(np.uint8)
-
-
-def solarize(image, threshold=128):
-    # For each pixel in the image, select the pixel
-    # if the value is less than the threshold.
-    # Otherwise, subtract 255 from the pixel.
-    return np.where(image < threshold, image, 255 - image)
-
-
-def solarize_add(image, addition=0, threshold=128):
-    # For each pixel in the image less than threshold
-    # we add 'addition' amount to it and then clip the
-    # pixel value to be between 0 and 255. The value
-    # of 'addition' is between -128 and 128.
-    added_image = image.astype(np.int64) + addition
-    added_image = np.clip(added_image, a_min=0, a_max=255).astype(np.uint8)
-    return np.where(image < threshold, added_image, image)
-
-
-def color(image, factor):
-    """use cv2 to deal"""
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    degenerate = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
-    return blend(degenerate, image, factor)
-
-
-# refer to https://github.com/4uiiurz1/pytorch-auto-augment/blob/024b2eac4140c38df8342f09998e307234cafc80/auto_augment.py#L197
-def contrast(img, factor):
-    img = ImageEnhance.Contrast(Image.fromarray(img)).enhance(factor)
-    return np.array(img)
-
-
-def brightness(image, factor):
-    """Equivalent of PIL Brightness."""
-    degenerate = np.zeros_like(image)
-    return blend(degenerate, image, factor)
-
-
-def posterize(image, bits):
-    """Equivalent of PIL Posterize."""
-    shift = 8 - bits
-    return np.left_shift(np.right_shift(image, shift), shift)
-
-
-def rotate(image, degrees, replace):
-    """Rotates the image by degrees either clockwise or counterclockwise.
-
-    Args:
-        image: An image Tensor of type uint8.
-        degrees: Float, a scalar angle in degrees to rotate all images by. If
-            degrees is positive the image will be rotated clockwise otherwise it will
-            be rotated counterclockwise.
-        replace: A one or three value 1D tensor to fill empty pixels caused by
-            the rotate operation.
-
-    Returns:
-        The rotated version of image.
-    """
-    image = wrap(image)
-    image = Image.fromarray(image)
-    image = image.rotate(degrees)
-    image = np.array(image, dtype=np.uint8)
-    return unwrap(image, replace)
-
-
-def random_shift_bbox(image,
-                      bbox,
-                      pixel_scaling,
-                      replace,
-                      new_min_bbox_coords=None):
-    """Move the bbox and the image content to a slightly new random location.
-
-    Args:
-        image: 3D uint8 Tensor.
-        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
-            of type float that represents the normalized coordinates between 0 and 1.
-            The potential values for the new min corner of the bbox will be between
-            [old_min - pixel_scaling * bbox_height/2,
-             old_min - pixel_scaling * bbox_height/2].
-        pixel_scaling: A float between 0 and 1 that specifies the pixel range
-            that the new bbox location will be sampled from.
-        replace: A one or three value 1D tensor to fill empty pixels.
-        new_min_bbox_coords: If not None, then this is a tuple that specifies the
-            (min_y, min_x) coordinates of the new bbox. Normally this is randomly
-            specified, but this allows it to be manually set. The coordinates are
-            the absolute coordinates between 0 and image height/width and are int32.
-
-    Returns:
-        The new image that will have the shifted bbox location in it along with
-        the new bbox that contains the new coordinates.
-    """
-    # Obtains image height and width and create helper clip functions.
-    image_height, image_width = image.shape[0], image.shape[1]
-    image_height = float(image_height)
-    image_width = float(image_width)
-
-    def clip_y(val):
-        return np.clip(val, a_min=0, a_max=image_height - 1).astype(np.int32)
-
-    def clip_x(val):
-        return np.clip(val, a_min=0, a_max=image_width - 1).astype(np.int32)
-
-    # Convert bbox to pixel coordinates.
-    min_y = int(image_height * bbox[0])
-    min_x = int(image_width * bbox[1])
-    max_y = clip_y(image_height * bbox[2])
-    max_x = clip_x(image_width * bbox[3])
-
-    bbox_height, bbox_width = (max_y - min_y + 1, max_x - min_x + 1)
-    image_height = int(image_height)
-    image_width = int(image_width)
-
-    # Select the new min/max bbox ranges that are used for sampling the
-    # new min x/y coordinates of the shifted bbox.
-    minval_y = clip_y(min_y - np.int32(pixel_scaling * float(bbox_height) /
-                                       2.0))
-    maxval_y = clip_y(min_y + np.int32(pixel_scaling * float(bbox_height) /
-                                       2.0))
-    minval_x = clip_x(min_x - np.int32(pixel_scaling * float(bbox_width) /
-                                       2.0))
-    maxval_x = clip_x(min_x + np.int32(pixel_scaling * float(bbox_width) /
-                                       2.0))
-
-    # Sample and calculate the new unclipped min/max coordinates of the new bbox.
-    if new_min_bbox_coords is None:
-        unclipped_new_min_y = np.random.randint(
-            low=minval_y, high=maxval_y, dtype=np.int32)
-        unclipped_new_min_x = np.random.randint(
-            low=minval_x, high=maxval_x, dtype=np.int32)
-    else:
-        unclipped_new_min_y, unclipped_new_min_x = (
-            clip_y(new_min_bbox_coords[0]), clip_x(new_min_bbox_coords[1]))
-    unclipped_new_max_y = unclipped_new_min_y + bbox_height - 1
-    unclipped_new_max_x = unclipped_new_min_x + bbox_width - 1
-
-    # Determine if any of the new bbox was shifted outside the current image.
-    # This is used for determining if any of the original bbox content should be
-    # discarded.
-    new_min_y, new_min_x, new_max_y, new_max_x = (
-        clip_y(unclipped_new_min_y), clip_x(unclipped_new_min_x),
-        clip_y(unclipped_new_max_y), clip_x(unclipped_new_max_x))
-    shifted_min_y = (new_min_y - unclipped_new_min_y) + min_y
-    shifted_max_y = max_y - (unclipped_new_max_y - new_max_y)
-    shifted_min_x = (new_min_x - unclipped_new_min_x) + min_x
-    shifted_max_x = max_x - (unclipped_new_max_x - new_max_x)
-
-    # Create the new bbox tensor by converting pixel integer values to floats.
-    new_bbox = np.stack([
-        float(new_min_y) / float(image_height), float(new_min_x) /
-        float(image_width), float(new_max_y) / float(image_height),
-        float(new_max_x) / float(image_width)
-    ])
-
-    # Copy the contents in the bbox and fill the old bbox location
-    # with gray (128).
-    bbox_content = image[shifted_min_y:shifted_max_y + 1, shifted_min_x:
-                         shifted_max_x + 1, :]
-
-    def mask_and_add_image(min_y_, min_x_, max_y_, max_x_, mask,
-                           content_tensor, image_):
-        """Applies mask to bbox region in image then adds content_tensor to it."""
-        mask = np.pad(mask, [[min_y_, (image_height - 1) - max_y_],
-                             [min_x_, (image_width - 1) - max_x_], [0, 0]],
-                      'constant',
-                      constant_values=1)
-
-        content_tensor = np.pad(content_tensor,
-                                [[min_y_, (image_height - 1) - max_y_],
-                                 [min_x_, (image_width - 1) - max_x_], [0, 0]],
-                                'constant',
-                                constant_values=0)
-        return image_ * mask + content_tensor
-
-    # Zero out original bbox location.
-    mask = np.zeros_like(image)[min_y:max_y + 1, min_x:max_x + 1, :]
-    grey_tensor = np.zeros_like(mask) + replace[0]
-    image = mask_and_add_image(min_y, min_x, max_y, max_x, mask, grey_tensor,
-                               image)
-
-    # Fill in bbox content to new bbox location.
-    mask = np.zeros_like(bbox_content)
-    image = mask_and_add_image(new_min_y, new_min_x, new_max_y, new_max_x,
-                               mask, bbox_content, image)
-
-    return image.astype(np.uint8), new_bbox
-
-
-def _clip_bbox(min_y, min_x, max_y, max_x):
-    """Clip bounding box coordinates between 0 and 1.
-
-    Args:
-        min_y: Normalized bbox coordinate of type float between 0 and 1.
-        min_x: Normalized bbox coordinate of type float between 0 and 1.
-        max_y: Normalized bbox coordinate of type float between 0 and 1.
-        max_x: Normalized bbox coordinate of type float between 0 and 1.
-
-    Returns:
-        Clipped coordinate values between 0 and 1.
-    """
-    min_y = np.clip(min_y, a_min=0, a_max=1.0)
-    min_x = np.clip(min_x, a_min=0, a_max=1.0)
-    max_y = np.clip(max_y, a_min=0, a_max=1.0)
-    max_x = np.clip(max_x, a_min=0, a_max=1.0)
-    return min_y, min_x, max_y, max_x
-
-
-def _check_bbox_area(min_y, min_x, max_y, max_x, delta=0.05):
-    """Adjusts bbox coordinates to make sure the area is > 0.
-
-    Args:
-        min_y: Normalized bbox coordinate of type float between 0 and 1.
-        min_x: Normalized bbox coordinate of type float between 0 and 1.
-        max_y: Normalized bbox coordinate of type float between 0 and 1.
-        max_x: Normalized bbox coordinate of type float between 0 and 1.
-        delta: Float, this is used to create a gap of size 2 * delta between
-            bbox min/max coordinates that are the same on the boundary.
-            This prevents the bbox from having an area of zero.
-
-    Returns:
-        Tuple of new bbox coordinates between 0 and 1 that will now have a
-        guaranteed area > 0.
-    """
-    height = max_y - min_y
-    width = max_x - min_x
-
-    def _adjust_bbox_boundaries(min_coord, max_coord):
-        # Make sure max is never 0 and min is never 1.
-        max_coord = np.maximum(max_coord, 0.0 + delta)
-        min_coord = np.minimum(min_coord, 1.0 - delta)
-        return min_coord, max_coord
-
-    if _equal(height, 0):
-        min_y, max_y = _adjust_bbox_boundaries(min_y, max_y)
-
-    if _equal(width, 0):
-        min_x, max_x = _adjust_bbox_boundaries(min_x, max_x)
-
-    return min_y, min_x, max_y, max_x
-
-
-def _scale_bbox_only_op_probability(prob):
-    """Reduce the probability of the bbox-only operation.
-
-    Probability is reduced so that we do not distort the content of too many
-    bounding boxes that are close to each other. The value of 3.0 was a chosen
-    hyper parameter when designing the autoaugment algorithm that we found
-    empirically to work well.
-
-    Args:
-        prob: Float that is the probability of applying the bbox-only operation.
-
-    Returns:
-        Reduced probability.
-    """
-    return prob / 3.0
-
-
-def _apply_bbox_augmentation(image, bbox, augmentation_func, *args):
-    """Applies augmentation_func to the subsection of image indicated by bbox.
-
-    Args:
-        image: 3D uint8 Tensor.
-        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
-            of type float that represents the normalized coordinates between 0 and 1.
-        augmentation_func: Augmentation function that will be applied to the
-            subsection of image.
-        *args: Additional parameters that will be passed into augmentation_func
-            when it is called.
-
-    Returns:
-        A modified version of image, where the bbox location in the image will
-        have `ugmentation_func applied to it.
-    """
-    image_height = image.shape[0]
-    image_width = image.shape[1]
-
-    min_y = int(image_height * bbox[0])
-    min_x = int(image_width * bbox[1])
-    max_y = int(image_height * bbox[2])
-    max_x = int(image_width * bbox[3])
-
-    # Clip to be sure the max values do not fall out of range.
-    max_y = np.minimum(max_y, image_height - 1)
-    max_x = np.minimum(max_x, image_width - 1)
-
-    # Get the sub-tensor that is the image within the bounding box region.
-    bbox_content = image[min_y:max_y + 1, min_x:max_x + 1, :]
-
-    # Apply the augmentation function to the bbox portion of the image.
-    augmented_bbox_content = augmentation_func(bbox_content, *args)
-
-    # Pad the augmented_bbox_content and the mask to match the shape of original
-    # image.
-    augmented_bbox_content = np.pad(
-        augmented_bbox_content, [[min_y, (image_height - 1) - max_y],
-                                 [min_x, (image_width - 1) - max_x], [0, 0]],
-        'constant',
-        constant_values=1)
-
-    # Create a mask that will be used to zero out a part of the original image.
-    mask_tensor = np.zeros_like(bbox_content)
-
-    mask_tensor = np.pad(mask_tensor,
-                         [[min_y, (image_height - 1) - max_y],
-                          [min_x, (image_width - 1) - max_x], [0, 0]],
-                         'constant',
-                         constant_values=1)
-    # Replace the old bbox content with the new augmented content.
-    image = image * mask_tensor + augmented_bbox_content
-    return image.astype(np.uint8)
-
-
-def _concat_bbox(bbox, bboxes):
-    """Helper function that concates bbox to bboxes along the first dimension."""
-
-    # Note if all elements in bboxes are -1 (_INVALID_BOX), then this means
-    # we discard bboxes and start the bboxes Tensor with the current bbox.
-    bboxes_sum_check = np.sum(bboxes)
-    bbox = np.expand_dims(bbox, 0)
-    # This check will be true when it is an _INVALID_BOX
-    if _equal(bboxes_sum_check, -4):
-        bboxes = bbox
-    else:
-        bboxes = np.concatenate([bboxes, bbox], 0)
-    return bboxes
-
-
-def _apply_bbox_augmentation_wrapper(image, bbox, new_bboxes, prob,
-                                     augmentation_func, func_changes_bbox,
-                                     *args):
-    """Applies _apply_bbox_augmentation with probability prob.
-
-    Args:
-        image: 3D uint8 Tensor.
-        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
-            of type float that represents the normalized coordinates between 0 and 1.
-        new_bboxes: 2D Tensor that is a list of the bboxes in the image after they
-            have been altered by aug_func. These will only be changed when
-            func_changes_bbox is set to true. Each bbox has 4 elements
-            (min_y, min_x, max_y, max_x) of type float that are the normalized
-            bbox coordinates between 0 and 1.
-        prob: Float that is the probability of applying _apply_bbox_augmentation.
-        augmentation_func: Augmentation function that will be applied to the
-            subsection of image.
-        func_changes_bbox: Boolean. Does augmentation_func return bbox in addition
-            to image.
-        *args: Additional parameters that will be passed into augmentation_func
-            when it is called.
-
-    Returns:
-        A tuple. Fist element is a modified version of image, where the bbox
-        location in the image will have augmentation_func applied to it if it is
-        chosen to be called with probability `prob`. The second element is a
-        Tensor of Tensors of length 4 that will contain the altered bbox after
-        applying augmentation_func.
-    """
-    should_apply_op = (np.random.rand() + prob >= 1)
-    if func_changes_bbox:
-        if should_apply_op:
-            augmented_image, bbox = augmentation_func(image, bbox, *args)
-        else:
-            augmented_image, bbox = (image, bbox)
-    else:
-        if should_apply_op:
-            augmented_image = _apply_bbox_augmentation(
-                image, bbox, augmentation_func, *args)
-        else:
-            augmented_image = image
-    new_bboxes = _concat_bbox(bbox, new_bboxes)
-    return augmented_image.astype(np.uint8), new_bboxes
-
-
-def _apply_multi_bbox_augmentation(image, bboxes, prob, aug_func,
-                                   func_changes_bbox, *args):
-    """Applies aug_func to the image for each bbox in bboxes.
-
-    Args:
-        image: 3D uint8 Tensor.
-        bboxes: 2D Tensor that is a list of the bboxes in the image. Each bbox
-            has 4 elements (min_y, min_x, max_y, max_x) of type float.
-        prob: Float that is the probability of applying aug_func to a specific
-            bounding box within the image.
-        aug_func: Augmentation function that will be applied to the
-            subsections of image indicated by the bbox values in bboxes.
-        func_changes_bbox: Boolean. Does augmentation_func return bbox in addition
-            to image.
-        *args: Additional parameters that will be passed into augmentation_func
-            when it is called.
-
-    Returns:
-        A modified version of image, where each bbox location in the image will
-        have augmentation_func applied to it if it is chosen to be called with
-        probability prob independently across all bboxes. Also the final
-        bboxes are returned that will be unchanged if func_changes_bbox is set to
-        false and if true, the new altered ones will be returned.
-    """
-    # Will keep track of the new altered bboxes after aug_func is repeatedly
-    # applied. The -1 values are a dummy value and this first Tensor will be
-    # removed upon appending the first real bbox.
-    new_bboxes = np.array(_INVALID_BOX)
-
-    # If the bboxes are empty, then just give it _INVALID_BOX. The result
-    # will be thrown away.
-    bboxes = np.array((_INVALID_BOX)) if bboxes.size == 0 else bboxes
-
-    assert bboxes.shape[1] == 4, "bboxes.shape[1] must be 4!!!!"
-
-    # pylint:disable=g-long-lambda
-    # pylint:disable=line-too-long
-    wrapped_aug_func = lambda _image, bbox, _new_bboxes: _apply_bbox_augmentation_wrapper(_image, bbox, _new_bboxes, prob, aug_func, func_changes_bbox, *args)
-    # pylint:enable=g-long-lambda
-    # pylint:enable=line-too-long
-
-    # Setup the while_loop.
-    num_bboxes = bboxes.shape[0]  # We loop until we go over all bboxes.
-    idx = 0  # Counter for the while loop.
-
-    # Conditional function when to end the loop once we go over all bboxes
-    # images_and_bboxes contain (_image, _new_bboxes)
-    def cond(_idx, _images_and_bboxes):
-        return _idx < num_bboxes
-
-    # Shuffle the bboxes so that the augmentation order is not deterministic if
-    # we are not changing the bboxes with aug_func.
-    # if not func_changes_bbox:
-    #     print(bboxes)
-    #     loop_bboxes = np.take(bboxes,np.random.permutation(bboxes.shape[0]),axis=0)
-    #     print(loop_bboxes)
-    # else:
-    #     loop_bboxes = bboxes
-    # we can not shuffle the bbox because it does not contain class information here
-    loop_bboxes = deepcopy(bboxes)
-
-    # Main function of while_loop where we repeatedly apply augmentation on the
-    # bboxes in the image.
-    # pylint:disable=g-long-lambda
-    body = lambda _idx, _images_and_bboxes: [
-            _idx + 1, wrapped_aug_func(_images_and_bboxes[0],
-                                         loop_bboxes[_idx],
-                                         _images_and_bboxes[1])]
-    while (cond(idx, (image, new_bboxes))):
-        idx, (image, new_bboxes) = body(idx, (image, new_bboxes))
-
-    # Either return the altered bboxes or the original ones depending on if
-    # we altered them in anyway.
-    if func_changes_bbox:
-        final_bboxes = new_bboxes
-    else:
-        final_bboxes = bboxes
-    return image, final_bboxes
-
-
-def _apply_multi_bbox_augmentation_wrapper(image, bboxes, prob, aug_func,
-                                           func_changes_bbox, *args):
-    """Checks to be sure num bboxes > 0 before calling inner function."""
-    num_bboxes = len(bboxes)
-    new_image = deepcopy(image)
-    new_bboxes = deepcopy(bboxes)
-    if num_bboxes != 0:
-        new_image, new_bboxes = _apply_multi_bbox_augmentation(
-            new_image, new_bboxes, prob, aug_func, func_changes_bbox, *args)
-    return new_image, new_bboxes
-
-
-def rotate_only_bboxes(image, bboxes, prob, degrees, replace):
-    """Apply rotate to each bbox in the image with probability prob."""
-    func_changes_bbox = False
-    prob = _scale_bbox_only_op_probability(prob)
-    return _apply_multi_bbox_augmentation_wrapper(
-        image, bboxes, prob, rotate, func_changes_bbox, degrees, replace)
-
-
-def shear_x_only_bboxes(image, bboxes, prob, level, replace):
-    """Apply shear_x to each bbox in the image with probability prob."""
-    func_changes_bbox = False
-    prob = _scale_bbox_only_op_probability(prob)
-    return _apply_multi_bbox_augmentation_wrapper(
-        image, bboxes, prob, shear_x, func_changes_bbox, level, replace)
-
-
-def shear_y_only_bboxes(image, bboxes, prob, level, replace):
-    """Apply shear_y to each bbox in the image with probability prob."""
-    func_changes_bbox = False
-    prob = _scale_bbox_only_op_probability(prob)
-    return _apply_multi_bbox_augmentation_wrapper(
-        image, bboxes, prob, shear_y, func_changes_bbox, level, replace)
-
-
-def translate_x_only_bboxes(image, bboxes, prob, pixels, replace):
-    """Apply translate_x to each bbox in the image with probability prob."""
-    func_changes_bbox = False
-    prob = _scale_bbox_only_op_probability(prob)
-    return _apply_multi_bbox_augmentation_wrapper(
-        image, bboxes, prob, translate_x, func_changes_bbox, pixels, replace)
-
-
-def translate_y_only_bboxes(image, bboxes, prob, pixels, replace):
-    """Apply translate_y to each bbox in the image with probability prob."""
-    func_changes_bbox = False
-    prob = _scale_bbox_only_op_probability(prob)
-    return _apply_multi_bbox_augmentation_wrapper(
-        image, bboxes, prob, translate_y, func_changes_bbox, pixels, replace)
-
-
-def flip_only_bboxes(image, bboxes, prob):
-    """Apply flip_lr to each bbox in the image with probability prob."""
-    func_changes_bbox = False
-    prob = _scale_bbox_only_op_probability(prob)
-    return _apply_multi_bbox_augmentation_wrapper(image, bboxes, prob,
-                                                  np.fliplr, func_changes_bbox)
-
-
-def solarize_only_bboxes(image, bboxes, prob, threshold):
-    """Apply solarize to each bbox in the image with probability prob."""
-    func_changes_bbox = False
-    prob = _scale_bbox_only_op_probability(prob)
-    return _apply_multi_bbox_augmentation_wrapper(
-        image, bboxes, prob, solarize, func_changes_bbox, threshold)
-
-
-def equalize_only_bboxes(image, bboxes, prob):
-    """Apply equalize to each bbox in the image with probability prob."""
-    func_changes_bbox = False
-    prob = _scale_bbox_only_op_probability(prob)
-    return _apply_multi_bbox_augmentation_wrapper(image, bboxes, prob,
-                                                  equalize, func_changes_bbox)
-
-
-def cutout_only_bboxes(image, bboxes, prob, pad_size, replace):
-    """Apply cutout to each bbox in the image with probability prob."""
-    func_changes_bbox = False
-    prob = _scale_bbox_only_op_probability(prob)
-    return _apply_multi_bbox_augmentation_wrapper(
-        image, bboxes, prob, cutout, func_changes_bbox, pad_size, replace)
-
-
-def _rotate_bbox(bbox, image_height, image_width, degrees):
-    """Rotates the bbox coordinated by degrees.
-
-    Args:
-        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
-            of type float that represents the normalized coordinates between 0 and 1.
-        image_height: Int, height of the image.
-        image_width: Int, height of the image.
-        degrees: Float, a scalar angle in degrees to rotate all images by. If
-            degrees is positive the image will be rotated clockwise otherwise it will
-            be rotated counterclockwise.
-
-    Returns:
-        A tensor of the same shape as bbox, but now with the rotated coordinates.
-    """
-    image_height, image_width = (float(image_height), float(image_width))
-
-    # Convert from degrees to radians.
-    degrees_to_radians = math.pi / 180.0
-    radians = degrees * degrees_to_radians
-
-    # Translate the bbox to the center of the image and turn the normalized 0-1
-    # coordinates to absolute pixel locations.
-    # Y coordinates are made negative as the y axis of images goes down with
-    # increasing pixel values, so we negate to make sure x axis and y axis points
-    # are in the traditionally positive direction.
-    min_y = -int(image_height * (bbox[0] - 0.5))
-    min_x = int(image_width * (bbox[1] - 0.5))
-    max_y = -int(image_height * (bbox[2] - 0.5))
-    max_x = int(image_width * (bbox[3] - 0.5))
-    coordinates = np.stack([[min_y, min_x], [min_y, max_x], [max_y, min_x],
-                            [max_y, max_x]]).astype(np.float32)
-    # Rotate the coordinates according to the rotation matrix clockwise if
-    # radians is positive, else negative
-    rotation_matrix = np.stack([[math.cos(radians), math.sin(radians)],
-                                [-math.sin(radians), math.cos(radians)]])
-    new_coords = np.matmul(rotation_matrix,
-                           np.transpose(coordinates)).astype(np.int32)
-
-    # Find min/max values and convert them back to normalized 0-1 floats.
-    min_y = -(float(np.max(new_coords[0, :])) / image_height - 0.5)
-    min_x = float(np.min(new_coords[1, :])) / image_width + 0.5
-    max_y = -(float(np.min(new_coords[0, :])) / image_height - 0.5)
-    max_x = float(np.max(new_coords[1, :])) / image_width + 0.5
-
-    # Clip the bboxes to be sure the fall between [0, 1].
-    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
-    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
-    return np.stack([min_y, min_x, max_y, max_x])
-
-
-def rotate_with_bboxes(image, bboxes, degrees, replace):
-    # Rotate the image.
-    image = rotate(image, degrees, replace)
-
-    # Convert bbox coordinates to pixel values.
-    image_height, image_width = image.shape[:2]
-    # pylint:disable=g-long-lambda
-    wrapped_rotate_bbox = lambda bbox: _rotate_bbox(bbox, image_height, image_width, degrees)
-    # pylint:enable=g-long-lambda
-    new_bboxes = np.zeros_like(bboxes)
-    for idx in range(len(bboxes)):
-        new_bboxes[idx] = wrapped_rotate_bbox(bboxes[idx])
-    return image, new_bboxes
-
-
-def translate_x(image, pixels, replace):
-    """Equivalent of PIL Translate in X dimension."""
-    image = Image.fromarray(wrap(image))
-    image = image.transform(image.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0))
-    return unwrap(np.array(image), replace)
-
-
-def translate_y(image, pixels, replace):
-    """Equivalent of PIL Translate in Y dimension."""
-    image = Image.fromarray(wrap(image))
-    image = image.transform(image.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels))
-    return unwrap(np.array(image), replace)
-
-
-def _shift_bbox(bbox, image_height, image_width, pixels, shift_horizontal):
-    """Shifts the bbox coordinates by pixels.
-
-    Args:
-        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
-            of type float that represents the normalized coordinates between 0 and 1.
-        image_height: Int, height of the image.
-        image_width: Int, width of the image.
-        pixels: An int. How many pixels to shift the bbox.
-        shift_horizontal: Boolean. If true then shift in X dimension else shift in
-            Y dimension.
-
-    Returns:
-        A tensor of the same shape as bbox, but now with the shifted coordinates.
-    """
-    pixels = int(pixels)
-    # Convert bbox to integer pixel locations.
-    min_y = int(float(image_height) * bbox[0])
-    min_x = int(float(image_width) * bbox[1])
-    max_y = int(float(image_height) * bbox[2])
-    max_x = int(float(image_width) * bbox[3])
-
-    if shift_horizontal:
-        min_x = np.maximum(0, min_x - pixels)
-        max_x = np.minimum(image_width, max_x - pixels)
-    else:
-        min_y = np.maximum(0, min_y - pixels)
-        max_y = np.minimum(image_height, max_y - pixels)
-
-    # Convert bbox back to floats.
-    min_y = float(min_y) / float(image_height)
-    min_x = float(min_x) / float(image_width)
-    max_y = float(max_y) / float(image_height)
-    max_x = float(max_x) / float(image_width)
-
-    # Clip the bboxes to be sure the fall between [0, 1].
-    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
-    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
-    return np.stack([min_y, min_x, max_y, max_x])
-
-
-def translate_bbox(image, bboxes, pixels, replace, shift_horizontal):
-    """Equivalent of PIL Translate in X/Y dimension that shifts image and bbox.
-
-    Args:
-        image: 3D uint8 Tensor.
-        bboxes: 2D Tensor that is a list of the bboxes in the image. Each bbox
-            has 4 elements (min_y, min_x, max_y, max_x) of type float with values
-            between [0, 1].
-        pixels: An int. How many pixels to shift the image and bboxes
-        replace: A one or three value 1D tensor to fill empty pixels.
-        shift_horizontal: Boolean. If true then shift in X dimension else shift in
-            Y dimension.
-
-    Returns:
-        A tuple containing a 3D uint8 Tensor that will be the result of translating
-        image by pixels. The second element of the tuple is bboxes, where now
-        the coordinates will be shifted to reflect the shifted image.
-    """
-    if shift_horizontal:
-        image = translate_x(image, pixels, replace)
-    else:
-        image = translate_y(image, pixels, replace)
-
-    # Convert bbox coordinates to pixel values.
-    image_height, image_width = image.shape[0], image.shape[1]
-    # pylint:disable=g-long-lambda
-    wrapped_shift_bbox = lambda bbox: _shift_bbox(bbox, image_height, image_width, pixels, shift_horizontal)
-    # pylint:enable=g-long-lambda
-    new_bboxes = deepcopy(bboxes)
-    num_bboxes = len(bboxes)
-    for idx in range(num_bboxes):
-        new_bboxes[idx] = wrapped_shift_bbox(bboxes[idx])
-    return image.astype(np.uint8), new_bboxes
-
-
-def shear_x(image, level, replace):
-    """Equivalent of PIL Shearing in X dimension."""
-    # Shear parallel to x axis is a projective transform
-    # with a matrix form of:
-    # [1    level
-    #    0    1].
-    image = Image.fromarray(wrap(image))
-    image = image.transform(image.size, Image.AFFINE, (1, level, 0, 0, 1, 0))
-    return unwrap(np.array(image), replace)
-
-
-def shear_y(image, level, replace):
-    """Equivalent of PIL Shearing in Y dimension."""
-    # Shear parallel to y axis is a projective transform
-    # with a matrix form of:
-    # [1    0
-    #    level    1].
-    image = Image.fromarray(wrap(image))
-    image = image.transform(image.size, Image.AFFINE, (1, 0, 0, level, 1, 0))
-    return unwrap(np.array(image), replace)
-
-
-def _shear_bbox(bbox, image_height, image_width, level, shear_horizontal):
-    """Shifts the bbox according to how the image was sheared.
-
-    Args:
-        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
-            of type float that represents the normalized coordinates between 0 and 1.
-        image_height: Int, height of the image.
-        image_width: Int, height of the image.
-        level: Float. How much to shear the image.
-        shear_horizontal: If true then shear in X dimension else shear in
-            the Y dimension.
-
-    Returns:
-        A tensor of the same shape as bbox, but now with the shifted coordinates.
-    """
-    image_height, image_width = (float(image_height), float(image_width))
-
-    # Change bbox coordinates to be pixels.
-    min_y = int(image_height * bbox[0])
-    min_x = int(image_width * bbox[1])
-    max_y = int(image_height * bbox[2])
-    max_x = int(image_width * bbox[3])
-    coordinates = np.stack(
-        [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
-    coordinates = coordinates.astype(np.float32)
-
-    # Shear the coordinates according to the translation matrix.
-    if shear_horizontal:
-        translation_matrix = np.stack([[1, 0], [-level, 1]])
-    else:
-        translation_matrix = np.stack([[1, -level], [0, 1]])
-    translation_matrix = translation_matrix.astype(np.float32)
-    new_coords = np.matmul(translation_matrix,
-                           np.transpose(coordinates)).astype(np.int32)
-
-    # Find min/max values and convert them back to floats.
-    min_y = float(np.min(new_coords[0, :])) / image_height
-    min_x = float(np.min(new_coords[1, :])) / image_width
-    max_y = float(np.max(new_coords[0, :])) / image_height
-    max_x = float(np.max(new_coords[1, :])) / image_width
-
-    # Clip the bboxes to be sure the fall between [0, 1].
-    min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
-    min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
-    return np.stack([min_y, min_x, max_y, max_x])
-
-
-def shear_with_bboxes(image, bboxes, level, replace, shear_horizontal):
-    """Applies Shear Transformation to the image and shifts the bboxes.
-
-    Args:
-        image: 3D uint8 Tensor.
-        bboxes: 2D Tensor that is a list of the bboxes in the image. Each bbox
-            has 4 elements (min_y, min_x, max_y, max_x) of type float with values
-            between [0, 1].
-        level: Float. How much to shear the image. This value will be between
-            -0.3 to 0.3.
-        replace: A one or three value 1D tensor to fill empty pixels.
-        shear_horizontal: Boolean. If true then shear in X dimension else shear in
-            the Y dimension.
-
-    Returns:
-        A tuple containing a 3D uint8 Tensor that will be the result of shearing
-        image by level. The second element of the tuple is bboxes, where now
-        the coordinates will be shifted to reflect the sheared image.
-    """
-    if shear_horizontal:
-        image = shear_x(image, level, replace)
-    else:
-        image = shear_y(image, level, replace)
-
-    # Convert bbox coordinates to pixel values.
-    image_height, image_width = image.shape[:2]
-    # pylint:disable=g-long-lambda
-    wrapped_shear_bbox = lambda bbox: _shear_bbox(bbox, image_height, image_width, level, shear_horizontal)
-    # pylint:enable=g-long-lambda
-    new_bboxes = deepcopy(bboxes)
-    num_bboxes = len(bboxes)
-    for idx in range(num_bboxes):
-        new_bboxes[idx] = wrapped_shear_bbox(bboxes[idx])
-    return image.astype(np.uint8), new_bboxes
-
-
-def autocontrast(image):
-    """Implements Autocontrast function from PIL.
-
-    Args:
-        image: A 3D uint8 tensor.
-
-    Returns:
-        The image after it has had autocontrast applied to it and will be of type
-        uint8.
-    """
-
-    def scale_channel(image):
-        """Scale the 2D image using the autocontrast rule."""
-        # A possibly cheaper version can be done using cumsum/unique_with_counts
-        # over the histogram values, rather than iterating over the entire image.
-        # to compute mins and maxes.
-        lo = float(np.min(image))
-        hi = float(np.max(image))
-
-        # Scale the image, making the lowest value 0 and the highest value 255.
-        def scale_values(im):
-            scale = 255.0 / (hi - lo)
-            offset = -lo * scale
-            im = im.astype(np.float32) * scale + offset
-            img = np.clip(im, a_min=0, a_max=255.0)
-            return im.astype(np.uint8)
-
-        result = scale_values(image) if hi > lo else image
-        return result
-
-    # Assumes RGB for now.    Scales each channel independently
-    # and then stacks the result.
-    s1 = scale_channel(image[:, :, 0])
-    s2 = scale_channel(image[:, :, 1])
-    s3 = scale_channel(image[:, :, 2])
-    image = np.stack([s1, s2, s3], 2)
-    return image
-
-
-def sharpness(image, factor):
-    """Implements Sharpness function from PIL."""
-    orig_image = image
-    image = image.astype(np.float32)
-    # Make image 4D for conv operation.
-    # SMOOTH PIL Kernel.
-    kernel = np.array(
-        [[1, 1, 1], [1, 5, 1], [1, 1, 1]], dtype=np.float32) / 13.
-    result = cv2.filter2D(image, -1, kernel).astype(np.uint8)
-
-    # Blend the final result.
-    return blend(result, orig_image, factor)
-
-
-def equalize(image):
-    """Implements Equalize function from PIL using."""
-
-    def scale_channel(im, c):
-        """Scale the data in the channel to implement equalize."""
-        im = im[:, :, c].astype(np.int32)
-        # Compute the histogram of the image channel.
-        histo, _ = np.histogram(im, range=[0, 255], bins=256)
-
-        # For the purposes of computing the step, filter out the nonzeros.
-        nonzero = np.where(np.not_equal(histo, 0))
-        nonzero_histo = np.reshape(np.take(histo, nonzero), [-1])
-        step = (np.sum(nonzero_histo) - nonzero_histo[-1]) // 255
-
-        def build_lut(histo, step):
-            # Compute the cumulative sum, shifting by step // 2
-            # and then normalization by step.
-            lut = (np.cumsum(histo) + (step // 2)) // step
-            # Shift lut, prepending with 0.
-            lut = np.concatenate([[0], lut[:-1]], 0)
-            # Clip the counts to be in range.    This is done
-            # in the C code for image.point.
-            return np.clip(lut, a_min=0, a_max=255).astype(np.uint8)
-
-        # If step is zero, return the original image.    Otherwise, build
-        # lut from the full histogram and step and then index from it.
-        if step == 0:
-            result = im
-        else:
-            result = np.take(build_lut(histo, step), im)
-
-        return result.astype(np.uint8)
-
-    # Assumes RGB for now.    Scales each channel independently
-    # and then stacks the result.
-    s1 = scale_channel(image, 0)
-    s2 = scale_channel(image, 1)
-    s3 = scale_channel(image, 2)
-    image = np.stack([s1, s2, s3], 2)
-    return image
-
-
-def wrap(image):
-    """Returns 'image' with an extra channel set to all 1s."""
-    shape = image.shape
-    extended_channel = 255 * np.ones([shape[0], shape[1], 1], image.dtype)
-    extended = np.concatenate([image, extended_channel], 2).astype(image.dtype)
-    return extended
-
-
-def unwrap(image, replace):
-    """Unwraps an image produced by wrap.
-
-    Where there is a 0 in the last channel for every spatial position,
-    the rest of the three channels in that spatial dimension are grayed
-    (set to 128).    Operations like translate and shear on a wrapped
-    Tensor will leave 0s in empty locations.    Some transformations look
-    at the intensity of values to do preprocessing, and we want these
-    empty pixels to assume the 'average' value, rather than pure black.
-
-
-    Args:
-        image: A 3D Image Tensor with 4 channels.
-        replace: A one or three value 1D tensor to fill empty pixels.
-
-    Returns:
-        image: A 3D image Tensor with 3 channels.
-    """
-    image_shape = image.shape
-    # Flatten the spatial dimensions.
-    flattened_image = np.reshape(image, [-1, image_shape[2]])
-
-    # Find all pixels where the last channel is zero.
-    alpha_channel = flattened_image[:, 3]
-
-    replace = np.concatenate([replace, np.ones([1], image.dtype)], 0)
-
-    # Where they are zero, fill them in with 'replace'.
-    alpha_channel = np.reshape(alpha_channel, (-1, 1))
-    alpha_channel = np.tile(alpha_channel, reps=(1, flattened_image.shape[1]))
-
-    flattened_image = np.where(
-        np.equal(alpha_channel, 0),
-        np.ones_like(
-            flattened_image, dtype=image.dtype) * replace,
-        flattened_image)
-
-    image = np.reshape(flattened_image, image_shape)
-    image = image[:, :, :3]
-    return image.astype(np.uint8)
-
-
-def _cutout_inside_bbox(image, bbox, pad_fraction):
-    """Generates cutout mask and the mean pixel value of the bbox.
-
-    First a location is randomly chosen within the image as the center where the
-    cutout mask will be applied. Note this can be towards the boundaries of the
-    image, so the full cutout mask may not be applied.
-
-    Args:
-        image: 3D uint8 Tensor.
-        bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
-            of type float that represents the normalized coordinates between 0 and 1.
-        pad_fraction: Float that specifies how large the cutout mask should be in
-            in reference to the size of the original bbox. If pad_fraction is 0.25,
-            then the cutout mask will be of shape
-            (0.25 * bbox height, 0.25 * bbox width).
-
-    Returns:
-        A tuple. Fist element is a tensor of the same shape as image where each
-        element is either a 1 or 0 that is used to determine where the image
-        will have cutout applied. The second element is the mean of the pixels
-        in the image where the bbox is located.
-        mask value: [0,1]
-    """
-    image_height, image_width = image.shape[0], image.shape[1]
-    # Transform from shape [1, 4] to [4].
-    bbox = np.squeeze(bbox)
-
-    min_y = int(float(image_height) * bbox[0])
-    min_x = int(float(image_width) * bbox[1])
-    max_y = int(float(image_height) * bbox[2])
-    max_x = int(float(image_width) * bbox[3])
-
-    # Calculate the mean pixel values in the bounding box, which will be used
-    # to fill the cutout region.
-    mean = np.mean(image[min_y:max_y + 1, min_x:max_x + 1], axis=(0, 1))
-    # Cutout mask will be size pad_size_heigh * 2 by pad_size_width * 2 if the
-    # region lies entirely within the bbox.
-    box_height = max_y - min_y + 1
-    box_width = max_x - min_x + 1
-    pad_size_height = int(pad_fraction * (box_height / 2))
-    pad_size_width = int(pad_fraction * (box_width / 2))
-
-    # Sample the center location in the image where the zero mask will be applied.
-    cutout_center_height = np.random.randint(min_y, max_y + 1, dtype=np.int32)
-    cutout_center_width = np.random.randint(min_x, max_x + 1, dtype=np.int32)
-
-    lower_pad = np.maximum(0, cutout_center_height - pad_size_height)
-    upper_pad = np.maximum(
-        0, image_height - cutout_center_height - pad_size_height)
-    left_pad = np.maximum(0, cutout_center_width - pad_size_width)
-    right_pad = np.maximum(0,
-                           image_width - cutout_center_width - pad_size_width)
-
-    cutout_shape = [
-        image_height - (lower_pad + upper_pad),
-        image_width - (left_pad + right_pad)
-    ]
-    padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]
-
-    mask = np.pad(np.zeros(
-        cutout_shape, dtype=image.dtype),
-                  padding_dims,
-                  'constant',
-                  constant_values=1)
-
-    mask = np.expand_dims(mask, 2)
-    mask = np.tile(mask, [1, 1, 3])
-    return mask, mean
-
-
-def bbox_cutout(image, bboxes, pad_fraction, replace_with_mean):
-    """Applies cutout to the image according to bbox information.
-
-    This is a cutout variant that using bbox information to make more informed
-    decisions on where to place the cutout mask.
-
-    Args:
-        image: 3D uint8 Tensor.
-        bboxes: 2D Tensor that is a list of the bboxes in the image. Each bbox
-            has 4 elements (min_y, min_x, max_y, max_x) of type float with values
-            between [0, 1].
-        pad_fraction: Float that specifies how large the cutout mask should be in
-            in reference to the size of the original bbox. If pad_fraction is 0.25,
-            then the cutout mask will be of shape
-            (0.25 * bbox height, 0.25 * bbox width).
-        replace_with_mean: Boolean that specified what value should be filled in
-            where the cutout mask is applied. Since the incoming image will be of
-            uint8 and will not have had any mean normalization applied, by default
-            we set the value to be 128. If replace_with_mean is True then we find
-            the mean pixel values across the channel dimension and use those to fill
-            in where the cutout mask is applied.
-
-    Returns:
-        A tuple. First element is a tensor of the same shape as image that has
-        cutout applied to it. Second element is the bboxes that were passed in
-        that will be unchanged.
-    """
-
-    def apply_bbox_cutout(image, bboxes, pad_fraction):
-        """Applies cutout to a single bounding box within image."""
-        # Choose a single bounding box to apply cutout to.
-        random_index = np.random.randint(0, bboxes.shape[0], dtype=np.int32)
-        # Select the corresponding bbox and apply cutout.
-        chosen_bbox = np.take(bboxes, random_index, axis=0)
-        mask, mean = _cutout_inside_bbox(image, chosen_bbox, pad_fraction)
-
-        # When applying cutout we either set the pixel value to 128 or to the mean
-        # value inside the bbox.
-        replace = mean if replace_with_mean else [128] * 3
-
-        # Apply the cutout mask to the image. Where the mask is 0 we fill it with
-        # `replace`.
-        image = np.where(
-            np.equal(mask, 0),
-            np.ones_like(
-                image, dtype=image.dtype) * replace,
-            image).astype(image.dtype)
-        return image
-
-    # Check to see if there are boxes, if so then apply boxcutout.
-    if len(bboxes) != 0:
-        image = apply_bbox_cutout(image, bboxes, pad_fraction)
-
-    return image, bboxes
-
-
-NAME_TO_FUNC = {
-        'AutoContrast': autocontrast,
-        'Equalize': equalize,
-        'Posterize': posterize,
-        'Solarize': solarize,
-        'SolarizeAdd': solarize_add,
-        'Color': color,
-        'Contrast': contrast,
-        'Brightness': brightness,
-        'Sharpness': sharpness,
-        'Cutout': cutout,
-        'BBox_Cutout': bbox_cutout,
-        'Rotate_BBox': rotate_with_bboxes,
-        # pylint:disable=g-long-lambda
-        'TranslateX_BBox': lambda image, bboxes, pixels, replace: translate_bbox(
-                image, bboxes, pixels, replace, shift_horizontal=True),
-        'TranslateY_BBox': lambda image, bboxes, pixels, replace: translate_bbox(
-                image, bboxes, pixels, replace, shift_horizontal=False),
-        'ShearX_BBox': lambda image, bboxes, level, replace: shear_with_bboxes(
-                image, bboxes, level, replace, shear_horizontal=True),
-        'ShearY_BBox': lambda image, bboxes, level, replace: shear_with_bboxes(
-                image, bboxes, level, replace, shear_horizontal=False),
-        # pylint:enable=g-long-lambda
-        'Rotate_Only_BBoxes': rotate_only_bboxes,
-        'ShearX_Only_BBoxes': shear_x_only_bboxes,
-        'ShearY_Only_BBoxes': shear_y_only_bboxes,
-        'TranslateX_Only_BBoxes': translate_x_only_bboxes,
-        'TranslateY_Only_BBoxes': translate_y_only_bboxes,
-        'Flip_Only_BBoxes': flip_only_bboxes,
-        'Solarize_Only_BBoxes': solarize_only_bboxes,
-        'Equalize_Only_BBoxes': equalize_only_bboxes,
-        'Cutout_Only_BBoxes': cutout_only_bboxes,
-}
-
-
-def _randomly_negate_tensor(tensor):
-    """With 50% prob turn the tensor negative."""
-    should_flip = np.floor(np.random.rand() + 0.5) >= 1
-    final_tensor = tensor if should_flip else -tensor
-    return final_tensor
-
-
-def _rotate_level_to_arg(level):
-    level = (level / _MAX_LEVEL) * 30.
-    level = _randomly_negate_tensor(level)
-    return (level, )
-
-
-def _shrink_level_to_arg(level):
-    """Converts level to ratio by which we shrink the image content."""
-    if level == 0:
-        return (1.0, )  # if level is zero, do not shrink the image
-    # Maximum shrinking ratio is 2.9.
-    level = 2. / (_MAX_LEVEL / level) + 0.9
-    return (level, )
-
-
-def _enhance_level_to_arg(level):
-    return ((level / _MAX_LEVEL) * 1.8 + 0.1, )
-
-
-def _shear_level_to_arg(level):
-    level = (level / _MAX_LEVEL) * 0.3
-    # Flip level to negative with 50% chance.
-    level = _randomly_negate_tensor(level)
-    return (level, )
-
-
-def _translate_level_to_arg(level, translate_const):
-    level = (level / _MAX_LEVEL) * float(translate_const)
-    # Flip level to negative with 50% chance.
-    level = _randomly_negate_tensor(level)
-    return (level, )
-
-
-def _bbox_cutout_level_to_arg(level, hparams):
-    cutout_pad_fraction = (
-        level / _MAX_LEVEL) * 0.75  # hparams.cutout_max_pad_fraction
-    return (cutout_pad_fraction, False
-            )  # hparams.cutout_bbox_replace_with_mean
-
-
-def level_to_arg(hparams):
-    return {
-        'AutoContrast': lambda level: (),
-        'Equalize': lambda level: (),
-        'Posterize': lambda level: (int((level / _MAX_LEVEL) * 4), ),
-        'Solarize': lambda level: (int((level / _MAX_LEVEL) * 256), ),
-        'SolarizeAdd': lambda level: (int((level / _MAX_LEVEL) * 110), ),
-        'Color': _enhance_level_to_arg,
-        'Contrast': _enhance_level_to_arg,
-        'Brightness': _enhance_level_to_arg,
-        'Sharpness': _enhance_level_to_arg,
-        'Cutout':
-        lambda level: (int((level / _MAX_LEVEL) * 100), ),  # hparams.cutout_const=100
-        # pylint:disable=g-long-lambda
-        'BBox_Cutout': lambda level: _bbox_cutout_level_to_arg(level, hparams),
-        'TranslateX_BBox':
-        lambda level: _translate_level_to_arg(level, 250),  # hparams.translate_const=250
-        'TranslateY_BBox':
-        lambda level: _translate_level_to_arg(level, 250),  # hparams.translate_cons
-        # pylint:enable=g-long-lambda
-        'ShearX_BBox': _shear_level_to_arg,
-        'ShearY_BBox': _shear_level_to_arg,
-        'Rotate_BBox': _rotate_level_to_arg,
-        'Rotate_Only_BBoxes': _rotate_level_to_arg,
-        'ShearX_Only_BBoxes': _shear_level_to_arg,
-        'ShearY_Only_BBoxes': _shear_level_to_arg,
-        # pylint:disable=g-long-lambda
-        'TranslateX_Only_BBoxes':
-        lambda level: _translate_level_to_arg(level, 120),  # hparams.translate_bbox_const
-        'TranslateY_Only_BBoxes':
-        lambda level: _translate_level_to_arg(level, 120),  # hparams.translate_bbox_const
-        # pylint:enable=g-long-lambda
-        'Flip_Only_BBoxes': lambda level: (),
-        'Solarize_Only_BBoxes':
-        lambda level: (int((level / _MAX_LEVEL) * 256), ),
-        'Equalize_Only_BBoxes': lambda level: (),
-        # pylint:disable=g-long-lambda
-        'Cutout_Only_BBoxes':
-        lambda level: (int((level / _MAX_LEVEL) * 50), ),  # hparams.cutout_bbox_const
-        # pylint:enable=g-long-lambda
-    }
-
-
-def bbox_wrapper(func):
-    """Adds a bboxes function argument to func and returns unchanged bboxes."""
-
-    def wrapper(images, bboxes, *args, **kwargs):
-        return (func(images, *args, **kwargs), bboxes)
-
-    return wrapper
-
-
-def _parse_policy_info(name, prob, level, replace_value, augmentation_hparams):
-    """Return the function that corresponds to `name` and update `level` param."""
-    func = NAME_TO_FUNC[name]
-    args = level_to_arg(augmentation_hparams)[name](level)
-
-    # Check to see if prob is passed into function. This is used for operations
-    # where we alter bboxes independently.
-    # pytype:disable=wrong-arg-types
-    if 'prob' in inspect.getfullargspec(func)[0]:
-        args = tuple([prob] + list(args))
-    # pytype:enable=wrong-arg-types
-
-    # Add in replace arg if it is required for the function that is being called.
-    if 'replace' in inspect.getfullargspec(func)[0]:
-        # Make sure replace is the final argument
-        assert 'replace' == inspect.getfullargspec(func)[0][-1]
-        args = tuple(list(args) + [replace_value])
-
-    # Add bboxes as the second positional argument for the function if it does
-    # not already exist.
-    if 'bboxes' not in inspect.getfullargspec(func)[0]:
-        func = bbox_wrapper(func)
-    return (func, prob, args)
-
-
-def _apply_func_with_prob(func, image, args, prob, bboxes):
-    """Apply `func` to image w/ `args` as input with probability `prob`."""
-    assert isinstance(args, tuple)
-    assert 'bboxes' == inspect.getfullargspec(func)[0][1]
-
-    # If prob is a function argument, then this randomness is being handled
-    # inside the function, so make sure it is always called.
-    if 'prob' in inspect.getfullargspec(func)[0]:
-        prob = 1.0
-
-    # Apply the function with probability `prob`.
-    should_apply_op = np.floor(np.random.rand() + 0.5) >= 1
-    if should_apply_op:
-        augmented_image, augmented_bboxes = func(image, bboxes, *args)
-    else:
-        augmented_image, augmented_bboxes = (image, bboxes)
-    return augmented_image, augmented_bboxes
-
-
-def select_and_apply_random_policy(policies, image, bboxes):
-    """Select a random policy from `policies` and apply it to `image`."""
-    policy_to_select = np.random.randint(0, len(policies), dtype=np.int32)
-    # policy_to_select = 6 # for test
-    for (i, policy) in enumerate(policies):
-        if i == policy_to_select:
-            image, bboxes = policy(image, bboxes)
-    return (image, bboxes)
-
-
-def build_and_apply_nas_policy(policies, image, bboxes, augmentation_hparams):
-    """Build a policy from the given policies passed in and apply to image.
-
-    Args:
-        policies: list of lists of tuples in the form `(func, prob, level)`, `func`
-            is a string name of the augmentation function, `prob` is the probability
-            of applying the `func` operation, `level` is the input argument for
-            `func`.
-        image: numpy array that the resulting policy will be applied to.
-        bboxes:
-        augmentation_hparams: Hparams associated with the NAS learned policy.
-
-    Returns:
-        A version of image that now has data augmentation applied to it based on
-        the `policies` pass into the function. Additionally, returns bboxes if
-        a value for them is passed in that is not None
-    """
-    replace_value = [128, 128, 128]
-
-    # func is the string name of the augmentation function, prob is the
-    # probability of applying the operation and level is the parameter associated
-
-    # tf_policies are functions that take in an image and return an augmented
-    # image.
-    tf_policies = []
-    for policy in policies:
-        tf_policy = []
-        # Link string name to the correct python function and make sure the correct
-        # argument is passed into that function.
-        for policy_info in policy:
-            policy_info = list(
-                policy_info) + [replace_value, augmentation_hparams]
-
-            tf_policy.append(_parse_policy_info(*policy_info))
-        # Now build the tf policy that will apply the augmentation procedue
-        # on image.
-        def make_final_policy(tf_policy_):
-            def final_policy(image_, bboxes_):
-                for func, prob, args in tf_policy_:
-                    image_, bboxes_ = _apply_func_with_prob(func, image_, args,
-                                                            prob, bboxes_)
-                return image_, bboxes_
-
-            return final_policy
-
-        tf_policies.append(make_final_policy(tf_policy))
-
-    augmented_images, augmented_bboxes = select_and_apply_random_policy(
-        tf_policies, image, bboxes)
-    # If no bounding boxes were specified, then just return the images.
-    return (augmented_images, augmented_bboxes)
-
-
-# TODO(barretzoph): Add in ArXiv link once paper is out.
-def distort_image_with_autoaugment(image, bboxes, augmentation_name):
-    """Applies the AutoAugment policy to `image` and `bboxes`.
-
-    Args:
-        image: `Tensor` of shape [height, width, 3] representing an image.
-        bboxes: `Tensor` of shape [N, 4] representing ground truth boxes that are
-            normalized between [0, 1].
-        augmentation_name: The name of the AutoAugment policy to use. The available
-            options are `v0`, `v1`, `v2`, `v3` and `test`. `v0` is the policy used for
-            all of the results in the paper and was found to achieve the best results
-            on the COCO dataset. `v1`, `v2` and `v3` are additional good policies
-            found on the COCO dataset that have slight variation in what operations
-            were used during the search procedure along with how many operations are
-            applied in parallel to a single image (2 vs 3).
-
-    Returns:
-        A tuple containing the augmented versions of `image` and `bboxes`.
-    """
-    available_policies = {
-        'v0': policy_v0,
-        'v1': policy_v1,
-        'v2': policy_v2,
-        'v3': policy_v3,
-        'test': policy_vtest
-    }
-    if augmentation_name not in available_policies:
-        raise ValueError('Invalid augmentation_name: {}'.format(
-            augmentation_name))
-
-    policy = available_policies[augmentation_name]()
-    augmentation_hparams = {}
-    return build_and_apply_nas_policy(policy, image, bboxes,
-                                      augmentation_hparams)