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

# Image Augmentation Function
def augment_images(image_files, num_duplicates):
    datagen = ImageDataGenerator(
        rotation_range=40,
        width_shift_range=0.2,
        height_shift_range=0.2,
        zoom_range=0.2,
        shear_range=0.2,
        horizontal_flip=True,
        fill_mode='nearest')

    # Create a temporary directory to store augmented images
    with tempfile.TemporaryDirectory() as temp_dir:
        for image_file in image_files:
            img = Image.open(image_file).convert('RGB')  # Convert to RGB for consistency
            img = img.resize((256, 256))  # Resize image
            x = img_to_array(img)  # Convert the image to a numpy array
            x = x.reshape((1,) + x.shape)  # Reshape for the data generator

            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, dirs, files in os.walk(temp_dir):
                for file in files:
                    zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.join(temp_dir, '..')))
        return zip_name

# Gradio Interface
def gradio_interface(image_files, num_duplicates):
    zip_file_path = augment_images(image_files, num_duplicates)
    return zip_file_path

iface = gr.Interface(fn=gradio_interface,
                     inputs=[gr.inputs.Image(type="file", label="Upload Images", accept="image/*", multiple=True),
                             gr.inputs.Number(default=5, label="Number of Duplicates", min_value=1, max_value=20)],
                     outputs=gr.outputs.File(label="Download Augmented Images"),
                     title="Image Augmentation App",
                     description="Upload images to generate augmented versions.")

iface.launch()