app.py

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')
    
    # Convert to RGB and resize the image
    img = Image.open(image_file).convert('RGB')
    img = img.resize((256, 256))
    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 = os.path.join(tempfile.gettempdir(), 'augmented_images.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(tool="editor", label="Upload Image"),
        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()