import streamlit as st from PIL import Image, ImageEnhance, ImageOps import numpy as np import io import zipfile 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("Ready-To-Use Synthetic Image Dataset Generation with Few-shots") 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("Generate Synthetic Dataset") and all_augmented_images: 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" )