Update app.py

app.py

@@ -1,15 +1,58 @@
-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 io
+import zipfile
+
+def augment_images(images, num_duplicates):
+    # Initialize the image 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')
+    
+    # Create a zip buffer in memory
+    zip_buffer = io.BytesIO()
+    
+    with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED, False) as zipf:
+        for i, img_file in enumerate(images):
+            # Open image and prepare for augmentation
+            img = Image.open(img_file).convert('RGB')
+            img = img.resize((256, 256))  # Resize for consistency
+            x = img_to_array(img)  # Convert to array
+            x = np.expand_dims(x, axis=0)  # Add batch dimension
+            
+            # Generate augmented images and add them to the zip
+            for j in range(num_duplicates):
+                aug_iter = datagen.flow(x, batch_size=1)
+                aug_image = next(aug_iter)[0].astype('uint8')
+                aug_image_pil = Image.fromarray(aug_image)
+                
+                # Save augmented image to zip
+                img_byte_arr = io.BytesIO()
+                aug_image_pil.save(img_byte_arr, format='PNG')
+                img_byte_arr = img_byte_arr.getvalue()
+                zipf.writestr(f"augmented_image_{i}_{j}.png", img_byte_arr)
+    
+    # Prepare the zip file for downloading
+    zip_buffer.seek(0)
+    
+    return zip_buffer
 
-def sepia(input_img):
-    sepia_filter = np.array([
-        [0.393, 0.769, 0.189], 
-        [0.349, 0.686, 0.168], 
-        [0.272, 0.534, 0.131]
-    ])
-    sepia_img = input_img.dot(sepia_filter.T)
-    sepia_img /= sepia_img.max()
-    return sepia_img
+# Define Gradio interface
+demo = gr.Interface(
+    fn=augment_images,
+    inputs=[
+        gr.Files(label="Upload Images", multiple=True),
+        gr.Slider(minimum=1, maximum=10, default=5, label="Number of Augmented Images per Original")
+    ],
+    outputs=gr.File(label="Download Augmented Images Zip"),
+    title="Image Augmentation App",
+    description="Upload images to generate augmented versions. Adjust the number of augmented images per original image using the slider."
+)
 
-demo = gr.Interface(sepia, gr.Image(), "image")
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()