cutout tensorflow 2

cutout.py

@tf.function
def cutout_channel(image, probability=0.5, iterations=3):
    image = tf.cast(image, tf.float32)
    r, g, b = tf.split(image, 3, axis=2)

    # for cc in tf.range(iterations):
    r = random_function(r, cutout_base, probability, None, **{"sl": 0.02, "sh": 0.08, "r1": 0.3})
    g = random_function(g, cutout_base, probability, None, **{"sl": 0.02, "sh": 0.08, "r1": 0.3})
    b = random_function(b, cutout_base, probability, None, **{"sl": 0.02, "sh": 0.08, "r1": 0.3})

    image = tf.concat([r, g, b], axis=2)
    return image


@tf.function
def cutout(image, probability=0.5, iterations=3):
    # for cc in tf.range(iterations):
    image = random_function(image, cutout_base, probability, None, **{"sl": 0.0002, "sh": 0.0005, "r1": 1.0})
    return image


@tf.function
def cutout_base(image, probability=1.0, sl=0.02, sh=0.08, r1=1.0):
    """
    Implementation was taken and then refactored to TF2 from:
        https://github.com/uranusx86/Random-Erasing-tensorflow/blob/master/random_erasing.py
    """
    shape = tf.shape(image)

    height = shape[0]
    width = shape[1]
    channel = shape[2]

    image = tf.cast(image, tf.float32)
    area = tf.cast(width * height, tf.float32)

    erase_area_low_bound = tf.cast(tf.round(tf.sqrt(sl * area * r1)), tf.int32)
    erase_area_up_bound = tf.cast(tf.round(tf.sqrt((sh * area) / r1)), tf.int32)

    h_upper_bound = tf.minimum(erase_area_up_bound, height)
    w_upper_bound = tf.minimum(erase_area_up_bound, width)

    h = tf.random.uniform([], erase_area_low_bound, h_upper_bound, tf.int32)
    w = tf.random.uniform([], erase_area_low_bound, w_upper_bound, tf.int32)

    x1 = tf.random.uniform([], 0, height + 1 - h, tf.int32)
    y1 = tf.random.uniform([], 0, width + 1 - w, tf.int32)

    # todo: this
    erase_area = tf.cast(tf.random.uniform([h, w, channel], 0, 255, tf.int32), tf.uint8)
    erase_area = tf.cast(tf.ones([h, w, channel], tf.int32), tf.float32) * 255.0

    erasing_img = tf.image.pad_to_bounding_box(erase_area, x1, y1, height, width)
    mask = tf.equal(erasing_img, 255.0)
    erasing_img = tf.where(mask, erasing_img, image)

    return tf.cast(
        tf.cond(tf.random.uniform([], 0.0, 1.0) > probability, lambda: image, lambda: erasing_img), tf.float32
    )