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 io

# 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,
        fill_mode='nearest')

    augmented_images = []
    for image_file in image_files:
        try:
            img = Image.open(image_file).convert('RGB')
            img = img.resize((256, 256))
            x = img_to_array(img)
            x = x.reshape((1,) + x.shape)
            i = 0
            for _ in datagen.flow(x, batch_size=1):
                i += 1
                augmented_images.append(x[0])
                if i >= num_duplicates:
                    break
        except Exception as e:
            print(f"Error processing image: {e}")

    return augmented_images

# Gradio UI
demo = gr.Interface(
    fn=augment_images,
    inputs=gr.File(label="Upload Images", multiple=True, file_types=["jpg", "jpeg", "png"]),
    outputs=gr.Image(label="Augmented Images"),
    examples=[["images/cat.jpg"], ["images/dog.jpg"]],
    description="Image Augmentation App",
    allow_flagging=False)

demo.launch()