import streamlit as st
from PIL import Image
import numpy as np

def mock_encoder(image):
    """Simulates encoding an image into a latent representation."""
    return np.random.normal(0, 1, (1, 2)), np.random.normal(0, 1, (1, 2)), np.random.normal(0, 1, (1, 2))

def mock_decoder(latent_representation):
    """Simulates decoding a latent representation back into an image."""
    return np.random.rand(28, 28, 1) * 255  # Random image for demonstration

def latent_space_augmentation(image, encoder, decoder, noise_scale=0.1):
    z_mean, z_log_var, _ = encoder(image)
    epsilon = np.random.normal(size=z_mean.shape)
    z_augmented = z_mean + np.exp(0.5 * z_log_var) * epsilon * noise_scale
    augmented_image = decoder(z_augmented)
    return np.squeeze(augmented_image)

st.title("VAE-based 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).convert('L')
    image = image.resize((28, 28))  # Resizing for simplicity
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    if st.button("Augment Image"):
        # Convert PIL image to numpy array
        image_array = np.array(image) / 255.0  # Normalize the image
        image_array = image_array.reshape((28, 28, 1))  # Reshape for the mock encoder/decoder
        
        # Perform augmentation
        augmented_image = latent_space_augmentation(image_array, mock_encoder, mock_decoder)
        
        # Display the augmented image
        st.image(augmented_image, caption="Augmented Image", clamp=True, use_column_width=True)