Rahatara's picture
Update app.py
1fd6cdb verified
raw
history blame
2.56 kB
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."""
image = image.rotate(np.random.uniform(-30, 30))
enhancer = ImageEnhance.Color(image)
image = enhancer.enhance(np.random.uniform(0.75, 1.25))
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."""
image_array = np.array(image)
noise = np.random.normal(0, noise_scale, image_array.shape)
noised_image_array = np.clip(image_array + noise, 0, 255).astype(np.uint8)
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"
)