import gradio as gr
import imgaug.augmenters as iaa
import cv2
import numpy as np
import matplotlib.pyplot as plt

def augment_image(image, flip, rotate, brightness, noise_scale, elastic_alpha, elastic_sigma):
    image = np.array(image)
    
    # apply augs based on user inputs
    if flip:
        flip_aug = iaa.Fliplr(1.0)  # flips horizontally 100% of the time
        image = flip_aug.augment_image(image)

    rotate_aug = iaa.Affine(rotate=rotate)  # rotate by specified num of degrees (if you pass in a touple it will select a random value between options)
    image = rotate_aug.augment_image(image)

    brightness_aug = iaa.Multiply(brightness)  # adjust brightness
    image = brightness_aug.augment_image(image)

    noise_aug = iaa.AdditiveGaussianNoise(scale=(noise_scale))  # eaussian noise
    image = noise_aug.augment_image(image)

    elastic_aug = iaa.ElasticTransformation(alpha=elastic_alpha, sigma=elastic_sigma)  # elastic transformation
    image = elastic_aug.augment_image(image)
    
    return image

def gradio_interface(image, flip, rotate, brightness, noise_scale, elastic_alpha, elastic_sigma):
    augmented_image = augment_image(image, flip, rotate, brightness, noise_scale, elastic_alpha, elastic_sigma)
    return augmented_image

inputs = [
    gr.Image(type="pil"),  # Image input
    gr.Checkbox(label="Flip Image Horizontally"),  # Flip input
    gr.Slider(minimum=-180, maximum=180, step=1, value=0, label="Rotate Image (degrees)"),  # Rotation input
    gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Adjust Brightness"),  # Brightness input
    gr.Slider(minimum=0, maximum=100, step=1, value=10, label="Gaussian Noise Scale"),  # Noise input
    gr.Slider(minimum=0, maximum=200, step=10, value=100, label="Elastic Transformation Alpha"),  # Elastic Alpha input
    gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=3.0, label="Elastic Transformation Sigma")  # Elastic Sigma input
]

iface = gr.Interface(
    fn=gradio_interface, 
    inputs=inputs, 
    outputs=gr.Image(type="numpy"),
    title="Image Augmentation Demo",
    description="Try out different data augmentation techniques on your image.",
)

iface.launch()