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

def augment_images(image_file, num_duplicates):
    """
    Augments uploaded images based on specified transformations and number of duplicates,
    then compresses them into a ZIP file for download.
    """
    # Augmentation configuration
    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')
    
    # Initialize ZIP file in memory
    zip_buffer = io.BytesIO()
    
    with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED) as zipf:
        img = Image.open(image_file.name).convert('RGB')
        img = img.resize((256, 256)) # Resize image
        x = img_to_array(img)  # Convert image to np.array
        x = np.expand_dims(x, axis=0)  # Add batch dimension
        
        # Generate and save augmented images
        for i in range(num_duplicates):
            it = datagen.flow(x, batch_size=1)
            batch = next(it)
            image = Image.fromarray(batch[0].astype('uint8'), 'RGB')
            img_byte_arr = io.BytesIO()
            image.save(img_byte_arr, format='JPEG')
            img_byte_arr.seek(0)
            zipf.writestr(f"augmented_image_{i}.jpeg", img_byte_arr.getvalue())
    
    # Prepare ZIP file for download
    zip_buffer.seek(0)
    
    return zip_buffer, 'augmented_images.zip'

# Create the Gradio interface
iface = gr.Interface(
    fn=augment_images,
    inputs=[gr.inputs.Image(type='file', label="Upload Image"), gr.inputs.Number(default=5, label="Number of Augmented Images")],
    outputs=[gr.outputs.File(label="Download ZIP of Augmented Images"), "text"],
    title="Image Augmentation Tool",
    description="Upload an image and select the number of augmented images to generate. Download the result as a ZIP file."
)

if __name__ == "__main__":
    iface.launch()