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

def augment_images(images, num_duplicates):
    # Initialize the image data generator for augmentation
    datagen = ImageDataGenerator(
        rotation_range=40,
        width_shift_range=0.2,
        height_shift_range=0.2,
        zoom_range=0.2,
        fill_mode='nearest')
    
    # Create a zip buffer in memory
    zip_buffer = io.BytesIO()
    
    with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED, False) as zipf:
        for i, img_file in enumerate(images):
            # Open image and prepare for augmentation
            img = Image.open(img_file).convert('RGB')
            img = img.resize((256, 256))  # Resize for consistency
            x = img_to_array(img)  # Convert to array
            x = np.expand_dims(x, axis=0)  # Add batch dimension
            
            # Generate augmented images and add them to the zip
            for j in range(num_duplicates):
                aug_iter = datagen.flow(x, batch_size=1)
                aug_image = next(aug_iter)[0].astype('uint8')
                aug_image_pil = Image.fromarray(aug_image)
                
                # Save augmented image to zip
                img_byte_arr = io.BytesIO()
                aug_image_pil.save(img_byte_arr, format='PNG')
                img_byte_arr = img_byte_arr.getvalue()
                zipf.writestr(f"augmented_image_{i}_{j}.png", img_byte_arr)
    
    # Prepare the zip file for downloading
    zip_buffer.seek(0)
    
    return zip_buffer

# Define Gradio interface
demo = gr.Interface(
    fn=augment_images,
    inputs=[
        gr.Files(label="Upload Images", multiple=True),
        gr.Slider(minimum=1, maximum=10, default=5, label="Number of Augmented Images per Original")
    ],
    outputs=gr.File(label="Download Augmented Images Zip"),
    title="Image Augmentation App",
    description="Upload images to generate augmented versions. Adjust the number of augmented images per original image using the slider."
)

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