import streamlit as st
from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation, pipeline
from PIL import Image, ImageFilter, ImageOps
import numpy as np

# Function to apply Gaussian Blur
def apply_gaussian_blur(image, predicted_mask, sigma=15):
    mask = Image.fromarray(
        (predicted_mask.cpu().numpy() * 255).astype(np.uint8)
    ).resize(image.size).convert('L')

    foreground = Image.composite(
        image, Image.new('RGB', image.size), mask
    )
    background = Image.composite(
        image.filter(ImageFilter.GaussianBlur(sigma)), image,
        ImageOps.invert(mask)
    )

    final_image = Image.composite(foreground, background, mask)
    return final_image

# Function to load and process image for segmentation
def get_segmentation_mask(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
    predicted_mask = logits.argmax(1).squeeze(0)
    return predicted_mask

def get_depth_mask(image):
    pipe = pipeline(task="depth-estimation", model="Intel/dpt-beit-base-384")

    result = pipe(image)
    depth_map = result["depth"]

    return np.array(depth_map)

def add_depth_based_blur(depth_array, image):
    depth_normalized = (depth_array - depth_array.min()) / (depth_array.max() - depth_array.min()) * 15

    image_array = np.array(image)

    blurred_images = [
        np.array(image.filter(ImageFilter.GaussianBlur(radius)))
        for radius in range(16)
    ]

    depth_blurred_array = np.zeros_like(image_array)

    for i in range(depth_normalized.shape[0]):
        for j in range(depth_normalized.shape[1]):
            blur_lvl = 15 - int(depth_normalized[i, j])
            depth_blurred_array[i, j] = blurred_images[blur_lvl][i, j]

    return Image.fromarray(depth_blurred_array.astype(np.uint8))

# 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)
    image = image.resize((512, 512))
    st.image(image, caption="Uploaded Image", use_container_width=True)

    predicted_mask = get_segmentation_mask(image)
    
    # Apply Gaussian Blur
    # sigma = st.slider("Gaussian Blur Intensity", 5, 50, 15)
    blurred_image = apply_gaussian_blur(image, predicted_mask)
    st.image(blurred_image, caption="Gaussian Blurred Image", use_container_width=True)
    
    # Perform lens blur
    # st.write("Calculating depth and applying lens blur...")
    with st.spinner("Applying lens blur... This might take a few moments."):
        depth_array = get_depth_mask(image)
        lens_blurred_img = add_depth_based_blur(depth_array, image)
    st.image(lens_blurred_img, caption="Lens Blur Effect Image", use_container_width=True)