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

def augment_images(image_file, num_duplicates):
    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')
    
    img = Image.open(image_file).convert('RGB')  # Convert to RGB
    img = img.resize((256, 256))  # Resize image
    x = img_to_array(img)  # Convert image to numpy array
    x = x.reshape((1,) + x.shape)  # Reshape for data generator
    
    with tempfile.TemporaryDirectory() as temp_dir:
        i = 0
        for _ in datagen.flow(x, batch_size=1, save_to_dir=temp_dir, save_prefix='aug', save_format='jpeg'):
            i += 1
            if i >= num_duplicates:
                break

        # Zip the augmented images
        zip_name = tempfile.mktemp(suffix='.zip')
        with zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, _, files in os.walk(temp_dir):
                for file in files:
                    zipf.write(os.path.join(root, file), arcname=file)
    
    return zip_name

iface = gr.Interface(
    fn=augment_images,
    inputs=[gr.Image(shape=None, label="Upload Image", source="upload"), gr.Slider(minimum=1, maximum=20, default=5, label="Number of Augmented Samples")],
    outputs=gr.File(label="Download Augmented Images"),
    title="Image Augmentation App",
    description="Upload an image to generate augmented versions. Select the number of augmented duplicates you want for the image."
)

iface.launch()