import numpy as np
import gradio as gr
from tensorflow.keras.preprocessing.image import img_to_array, ImageDataGenerator
from PIL import Image

def augment_images(input_img):
    # Define data augmentation parameters
    datagen = ImageDataGenerator(
        rotation_range=40,
        width_shift_range=0.2,
        height_shift_range=0.2,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True,
        fill_mode='nearest'
    )
    
    # Convert input image to numpy array
    img = Image.open(input_img).convert('RGB')
    img = img.resize((256, 256))  # Resize image
    x = img_to_array(img)
    x = x.reshape((1,) + x.shape)
    
    # Generate augmented images
    augmented_images = []
    for _ in datagen.flow(x, batch_size=1, save_to_dir=None, save_prefix='', save_format='jpeg'):
        augmented_images.append(_.squeeze())
        if len(augmented_images) >= 5:  # Generate 5 augmented samples
            break
    
    return augmented_images

iface = gr.Interface(
    fn=augment_images,
    inputs=gr.inputs.Image(label="Upload Image"),
    outputs=gr.outputs.Image(type="numpy"),
    title="Image Data Augmentation App",
    description="Upload an image to generate augmented versions."
)

iface.launch()