import tensorflow as tf


def random_apply(fn, image, prob=1.):
  b, *_ = image.get_shape().as_list()
  chance = tf.less(tf.random_uniform([b], 0, 1.0), prob)
  return tf.where(chance, fn(image), tf.identity(image))

def color_distortion(image, s=1.0):
  lower, upper, x = (1 - 0.8 * s), (1 + 0.8 * s), image
  x = tf.image.random_brightness(x, max_delta=0.8*s)
  x = tf.image.random_contrast(x, lower=lower, upper=upper)
  x = tf.image.random_saturation(x, lower=lower, upper=upper)
  x = tf.image.random_hue(x, max_delta=0.2*s)
  x = tf.clip_by_value(x, 0, 1)
  return x

def color_drop(image):
