Update app.py

app.py

@@ -1,41 +1,76 @@
-import numpy as np
 import gradio as gr
 from tensorflow.keras.preprocessing.image import img_to_array, ImageDataGenerator
 from PIL import Image
+import os
+import zipfile
 
-def augment_images(input_img):
-    # Define data augmentation parameters
+# 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,
-        shear_range=0.2,
         zoom_range=0.2,
-        horizontal_flip=True,
-        fill_mode='nearest'
-    )
+        fill_mode='nearest')
+
+    # Process each uploaded image
+    for image_file in images:
+        augment_image(image_file, datagen, num_duplicates)
     
-    # Convert input image to numpy array
-    img = Image.open(input_img).convert('RGB')
-    img = img.resize((256, 256))  # Resize image
-    x = img_to_array(img)
-    x = x.reshape((1,) + x.shape)
+    # Create a zip file with all augmented images
+    zip_file = create_zip_from_temp()
     
-    # Generate augmented images
-    augmented_images = []
-    for _ in datagen.flow(x, batch_size=1, save_to_dir=None, save_prefix='', save_format='jpeg'):
-        augmented_images.append(_.squeeze())
-        if len(augmented_images) >= 5:  # Generate 5 augmented samples
-            break
+    # 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 augmented_images
-
-iface = gr.Interface(
-    fn=augment_images,
-    inputs=gr.inputs.Image(label="Upload Image"),
-    outputs=gr.outputs.Image(type="numpy"),
-    title="Image Data Augmentation App",
-    description="Upload an image to generate augmented versions."
+    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."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()