Update app.py

app.py

@@ -1,12 +1,10 @@
+import numpy as np
 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):
+def augment_images(input_img):
+    # Define data augmentation parameters
     datagen = ImageDataGenerator(
         rotation_range=40,
         width_shift_range=0.2,
@@ -14,35 +12,30 @@ def augment_images(image_file, num_duplicates):
         shear_range=0.2,
         zoom_range=0.2,
         horizontal_flip=True,
-        fill_mode='nearest')
+        fill_mode='nearest'
+    )
     
-    img = Image.open(image_file).convert('RGB')  # Convert to RGB
+    # 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)  # Convert image to numpy array
-    x = x.reshape((1,) + x.shape)  # Reshape for data generator
+    x = img_to_array(img)
+    x = x.reshape((1,) + x.shape)
     
-    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 = tempfile.mktemp(suffix='.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)
+    # 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
     
-    return zip_name
+    return augmented_images
 
 iface = gr.Interface(
     fn=augment_images,
-    inputs=[gr.inputs.Image(label="Upload Image"), gr.inputs.Slider(minimum=1, maximum=20, default=5, label="Number of Augmented Samples")],
-    outputs=gr.outputs.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."
+    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."
 )
 
 iface.launch()