import gradio as gr
from transformers import pipeline
from PIL import Image, ImageFilter
import numpy as np

# Load the segmentation pipeline
pipe = pipeline("image-segmentation", model="mattmdjaga/segformer_b2_clothes")

# Simplified refine_mask function
def refine_mask(mask):
    """Simplify and smooth the segmentation mask."""
    mask_array = np.array(mask)
    mask_array = (mask_array > 128).astype(np.uint8) * 255  # Threshold to binary mask
    refined_mask = Image.fromarray(mask_array).filter(ImageFilter.GaussianBlur(0.5))  # Smooth edges
    return refined_mask

# Function to blur the background
def blur_background(image, blur_radius):
    # Perform segmentation
    result = pipe(image)
    
    # Extract the background mask
    background_mask = None
    for entry in result:
        if entry["label"] == "Background":
            background_mask = refine_mask(entry["mask"])  # Refine the background mask
            break
    
    if background_mask is None:
        return image  # If no background is detected, return the original image
    
    # Convert the image and mask to NumPy arrays
    image_np = np.array(image)
    background_mask_np = np.array(background_mask)
    
    # Create a blurred version of the entire image
    blurred_image = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
    blurred_np = np.array(blurred_image)
    
    # Combine the original image and the blurred background
    final_image = np.where(background_mask_np[..., None] == 255, blurred_np, image_np).astype(np.uint8)
    
    # Convert back to PIL image
    return Image.fromarray(final_image)

# Example inputs for Gradio
examples = [["1.jpg", 10] ,["2.jpg", 10]  ,["3.jpg", 10] ] # Example: Image with a blur intensity of 10

# Gradio interface
interface = gr.Interface(
    fn=blur_background,
    inputs=[
        gr.Image(type="pil"),  # Input image
        gr.Slider(1, 50, step=1, label="Blur Intensity")  # Slider for blur radius
    ],
    outputs=gr.Image(type="pil"),  # Output image with blurred background
    examples=examples,  # Provide examples as a nested list
    title="Blur Background with Refined Mask",
    description="Upload an image and adjust the slider to control the background blur level. The background edges are smoothed for better blending."
)

# Launch the app
interface.launch()