Update app.py

app.py

@@ -1,50 +1,63 @@
 import streamlit as st
 from PIL import Image, ImageEnhance, ImageOps
 import numpy as np
+import os
+import zipfile
+import io
 
 def apply_basic_augmentations(image):
     """Applies basic augmentations such as rotation and color jitter."""
-    # Rotate the image
     image = image.rotate(np.random.uniform(-30, 30))
-    
-    # Apply color jitter
     enhancer = ImageEnhance.Color(image)
     image = enhancer.enhance(np.random.uniform(0.75, 1.25))
-    
-    # Mirror image randomly
     if np.random.rand() > 0.5:
         image = ImageOps.mirror(image)
-    
     return image
 
 def simulate_latent_space_noising(image, noise_scale=25):
-    """Simulates latent space manipulation by adding noise to the image."""
-    # Convert image to numpy array
+    """Simulates latent space manipulation by adding noise."""
     image_array = np.array(image)
-    
-    # Generate noise
     noise = np.random.normal(0, noise_scale, image_array.shape)
-    
-    # Add noise to image
     noised_image_array = np.clip(image_array + noise, 0, 255).astype(np.uint8)
-    
-    # Convert back to PIL image
-    noised_image = Image.fromarray(noised_image_array)
-    
-    return noised_image
-
-st.title("Hybrid Image Augmentation Demo")
-
-uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-if uploaded_image is not None:
-    image = Image.open(uploaded_image)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    
-    if st.button("Augment Image"):
-        # Apply basic transformations
-        transformed_image = apply_basic_augmentations(image)
-        
-        # Simulate VAE latent space manipulation by adding noise
-        augmented_image = simulate_latent_space_noising(transformed_image)
-        
-        st.image(augmented_image, caption="Augmented Image", use_column_width=True)
+    return Image.fromarray(noised_image_array)
+
+def augment_image(image, augmentations_count):
+    """Generates augmented versions of a single image."""
+    augmented_images = []
+    for _ in range(augmentations_count):
+        augmented_image = apply_basic_augmentations(image)
+        augmented_image = simulate_latent_space_noising(augmented_image)
+        augmented_images.append(augmented_image)
+    return augmented_images
+
+def create_downloadable_zip(augmented_images):
+    """Creates a ZIP file in memory for downloading."""
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED, False) as zip_file:
+        for idx, image in enumerate(augmented_images):
+            img_byte_arr = io.BytesIO()
+            image.save(img_byte_arr, format="JPEG")
+            zip_file.writestr(f"augmented_image_{idx+1}.jpg", img_byte_arr.getvalue())
+    zip_buffer.seek(0)
+    return zip_buffer
+
+st.title("Batch Image Augmentation for Dataset Creation")
+
+uploaded_files = st.file_uploader("Choose images (1-10)", accept_multiple_files=True, type=["jpg", "jpeg", "png"])
+augmentations_count = st.number_input("Number of augmented samples per image", min_value=1, max_value=10, value=3)
+
+if uploaded_files:
+    all_augmented_images = []
+    for uploaded_file in uploaded_files:
+        image = Image.open(uploaded_file).convert("RGB")
+        augmented_images = augment_image(image, augmentations_count)
+        all_augmented_images.extend(augmented_images)
+
+    if st.button("Download Augmented Dataset"):
+        zip_buffer = create_downloadable_zip(all_augmented_images)
+        st.download_button(
+            label="Download ZIP",
+            data=zip_buffer,
+            file_name="augmented_images.zip",
+            mime="application/zip"
+        )