import gradio as gr
from transformers import pipeline
import torch
import numpy as np
from PIL import Image
import io

def remove_background(input_image):
    try:
        # Initialize the pipeline
        segmentor = pipeline("image-segmentation", 
                           model="briaai/RMBG-1.4",
                           device=-1)  # CPU inference
        
        # Process the image
        result = segmentor(input_image)
        
        return result['output_image']
    except Exception as e:
        raise gr.Error(f"Error processing image: {str(e)}")

# Create the Gradio interface with custom styling
css = """
.gradio-container {
    font-family: 'Segoe UI', sans-serif;
    background: linear-gradient(135deg, #1a1a1a 0%, #2d2d2d 100%);
}
.gr-image {
    border-radius: 12px;
    border: 2px solid rgba(255, 215, 0, 0.3);
}
.gr-button {
    background: linear-gradient(45deg, #FFD700, #FFA500);
    border: none;
    color: black;
}
.gr-button:hover {
    transform: translateY(-2px);
    box-shadow: 0 4px 12px rgba(255, 215, 0, 0.3);
}
"""

with gr.Blocks(css=css) as demo:
    gr.HTML(
        """
        <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
            <h1 style="font-size: 2.5rem; margin-bottom: 1rem; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
                AI Background Remover
            </h1>
            <p style="color: #cccccc; font-size: 1.1rem; margin-bottom: 2rem;">
                Powered by RMBG V1.4 model from BRIA AI
            </p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload Image",
                type="pil",
                sources=["upload", "clipboard"]
            )
            
        with gr.Column():
            output_image = gr.Image(
                label="Result",
                type="pil"
            )
    
    with gr.Row():
        clear_btn = gr.Button("Clear")
        process_btn = gr.Button("Remove Background", variant="primary")

    # Event handlers
    process_btn.click(
        fn=remove_background,
        inputs=[input_image],
        outputs=[output_image]
    )
    
    clear_btn.click(
        lambda: (None, None),
        outputs=[input_image, output_image]
    )

demo.launch()