import streamlit as st
from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation
from PIL import Image
import numpy as np
import cv2
import torch

# Function to apply Gaussian Blur
def apply_gaussian_blur(image, sigma=15):
    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    blurred = cv2.GaussianBlur(image_cv, (0, 0), sigma)
    return Image.fromarray(cv2.cvtColor(blurred, cv2.COLOR_BGR2RGB))

# Function to load and process image for segmentation
def segment_image(image):
    feature_extractor = MobileViTFeatureExtractor.from_pretrained("apple/mobilevit-small")
    model = MobileViTForSemanticSegmentation.from_pretrained("apple/mobilevit-small")
    
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    
    # Get segmentation mask
    logits = outputs.logits
    upsampled_logits = torch.nn.functional.interpolate(
        logits, size=image.size[::-1], mode="bilinear", align_corners=False
    )
    segmentation = upsampled_logits.argmax(dim=1).squeeze().detach().cpu().numpy()
    return segmentation

# Streamlit interface
st.title("Image Segmentation and Blur Effects")
st.write("Upload an image to apply segmentation, Gaussian blur, and depth-based blur.")

uploaded_file = st.file_uploader("Upload an Image (PNG, JPG, JPEG)", type=["png", "jpg", "jpeg"])

if uploaded_file:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    # Apply Gaussian Blur
    sigma = st.slider("Gaussian Blur Intensity", 5, 50, 15)
    blurred_image = apply_gaussian_blur(image, sigma)
    st.image(blurred_image, caption="Gaussian Blurred Image", use_column_width=True)
    
    # Perform segmentation
    if st.button("Perform Segmentation"):
        st.write("Segmenting the image...")
        segmentation = segment_image(image)
        st.image(segmentation, caption="Segmentation Mask", use_column_width=True)