app.py

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

# Define where augmented images will be temporarily stored
TEMP_DIR = "temp_augmented_images"

# Ensure the temp directory exists
if not os.path.exists(TEMP_DIR):
    os.makedirs(TEMP_DIR)

# Image Augmentation Function
def augment_image(image_file, datagen, num_duplicates):
    try:
        img = Image.open(image_file).convert('RGB')  # Convert to RGB
        img = img.resize((256, 256))  # Resize for consistency
        x = img_to_array(img)  # Image to array
        x = x.reshape((1,) + x.shape)  # Reshape for data generator
        
        # Augment image
        i = 0
        for batch 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
    except Exception as e:
        print(f"Error in augmenting image: {e}")

def create_zip_from_temp(directory=TEMP_DIR):
    zip_path = f"{directory}/augmented_images.zip"
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, _, files in os.walk(directory):
            for file in files:
                if file.endswith(".jpeg"):  # Ensure only augmented images are added
                    zipf.write(os.path.join(root, file), arcname=file)
    return zip_path

def process_images(images, num_duplicates):
    # 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')

    # Process each uploaded image
    for image_file in images:
        augment_image(image_file, datagen, num_duplicates)
    
    # Create a zip file with all augmented images
    zip_file = create_zip_from_temp()
    
    # Clean up augmented images to avoid clutter
    for file in os.listdir(TEMP_DIR):
        if file.endswith(".jpeg"):  # Clean up only augmented images, not the zip
            os.remove(os.path.join(TEMP_DIR, file))
    
    return zip_file

# Gradio Interface
demo = gr.Interface(
    fn=process_images, 
    inputs=[
        gr.Files(type="file", label="Upload Images", accept=["image/jpeg", "image/png"], multiple=True),
        gr.Slider(minimum=1, maximum=20, default=5, label="Number of Duplicates per Image")
    ],
    outputs=gr.File(label="Download Augmented Images"),
    title="Image Augmentation App",
    description="Upload images to augment them with random transformations. Download the augmented images as a zip file."
)

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