import numpy as np
import torch
from PIL import Image
import cv2
from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
import gradio as gr

# Initialize the SegFormer model for segmentation
segformer_processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
segformer_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")

# Function to segment the person in the image
def segment_person(image_input):
    # Convert input image (numpy array in RGB) to PIL Image
    image = Image.fromarray(image_input).convert("RGB")
    original_width, original_height = image.size
    
    # Resize image to 512x512 for the model
    model_input = image.resize((512, 512), Image.Resampling.LANCZOS)
    
    # Prepare the image for SegFormer
    inputs = segformer_processor(images=model_input, return_tensors="pt")
    
    # Perform inference
    with torch.no_grad():
        outputs = segformer_model(**inputs)
        logits = outputs.logits
    
    # Upsample logits to 512x512
    upsampled_logits = torch.nn.functional.interpolate(
        logits, size=(512, 512), mode="bilinear", align_corners=False
    )
    
    # Get the predicted segmentation mask (person class = 12 in ADE20K dataset)
    person_class_id = 12
    predicted_mask = upsampled_logits.argmax(dim=1)[0]  # Shape: (512, 512)
    binary_mask = (predicted_mask == person_class_id).cpu().numpy()  # Boolean mask
    
    # Post-process the mask
    mask_uint8 = (binary_mask * 255).astype(np.uint8)
    kernel = np.ones((5, 5), np.uint8)
    mask_cleaned = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel, iterations=2)
    mask_cleaned = cv2.morphologyEx(mask_cleaned, cv2.MORPH_OPEN, kernel, iterations=2)
    mask_smoothed = cv2.GaussianBlur(mask_cleaned, (7, 7), 0)
    _, mask_final = cv2.threshold(mask_smoothed, 127, 255, cv2.THRESH_BINARY)
    
    # Resize mask back to original dimensions
    mask_pil = Image.fromarray(mask_final)
    mask_resized = mask_pil.resize((original_width, original_height), Image.Resampling.LANCZOS)
    mask_array = np.array(mask_resized) > 0  # Boolean mask
    
    return mask_array

# Function to apply background blur
def blur_background(image_input, blur_strength):
    # Ensure image is in numpy array format (RGB)
    image_array = np.array(image_input)
    
    # Segment the person
    mask = segment_person(image_array)
    
    # Apply Gaussian blur to the entire image
    sigma = blur_strength
    blurred_image = cv2.GaussianBlur(image_array, (0, 0), sigmaX=sigma, sigmaY=sigma)
    
    # Composite the original foreground with the blurred background
    mask_3d = mask[:, :, np.newaxis]  # Add channel dimension for broadcasting
    result = np.where(mask_3d, image_array, blurred_image).astype(np.uint8)
    
    return result

# Gradio interface function
def gradio_interface(image, blur_strength):
    if image is None:
        raise ValueError("Please upload an image.")
    
    # Process the image
    output_image = blur_background(image, blur_strength)
    
    return output_image

# Create the Gradio app
app = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Image(type="numpy", label="Upload Image"),
        gr.Slider(minimum=1, maximum=25, value=10, step=1, label="Blur Strength (Sigma)")
    ],
    outputs=gr.Image(type="numpy", label="Output Image"),
    title="Person Segmentation and Background Blur",
    description="Upload an image to segment the person and blur the background. Adjust the blur strength using the slider."
)

# Launch the app
app.launch()