Spaces:
Runtime error
Runtime error
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" | |
) | |