import streamlit as st
from PIL import Image, ImageEnhance, ImageOps
import numpy as np

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
    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)