import tensorflow as tf

import gin
import numpy as np
import random
import tensorflow as tf

FLAGS = flags.FLAGS

# augmentation functions

# augment

def random_crop_and_resize(images, ratio=0.8):
  b, h, w, c = images.get_shape().as_list()
  ch, cw = map(lambda x: int(x * ratio), (h, w))
  crop = tf.random_crop(images, size=[b, ch, cw, 3])
  crop = tf.image.resize(crop, [h, w])
  return crop

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)
