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:
        # Convert input to PIL Image if it's not already
        if not isinstance(input_image, Image.Image):
            input_image = Image.fromarray(input_image)
            
        # Initialize the pipeline
        segmentor = pipeline(
            task="image-segmentation",
            model="briaai/RMBG-1.4",
            trust_remote_code=True
        )
        
        # Process the image and get mask
        result = segmentor(
            input_image, 
            return_mask=True
        )
        
        # Create output image with transparent background
        output_image = Image.new('RGBA', input_image.size, (0, 0, 0, 0))
        
        # Convert input to RGBA if it's not already
        if input_image.mode != 'RGBA':
            input_image = input_image.convert('RGBA')
            
        # Apply mask to create transparent background
        mask = result['mask'] if isinstance(result, dict) else result
        output_image.paste(input_image, mask=mask)
        
        return output_image
        
    except Exception as e:
        raise gr.Error(f"Error processing image: {str(e)}")

# Create Gradio interface
with gr.Blocks() 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;">
                AI Background Remover
            </h1>
            <p style="color: #666; font-size: 1.1rem;">
                Remove backgrounds instantly using RMBG V1.4 model
            </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", variant="secondary")
        process_btn = gr.Button("Remove Background", variant="primary")

    # Status message
    status_msg = gr.Textbox(
        label="Status",
        placeholder="Ready to process your image...",
        interactive=False
    )

    # Event handlers
    def process_and_update(image):
        if image is None:
            return None, "Please upload an image first"
        try:
            result = remove_background(image)
            return result, "✨ Background removed successfully!"
        except Exception as e:
            return None, f"❌ Error: {str(e)}"

    process_btn.click(
        fn=process_and_update,
        inputs=[input_image],
        outputs=[output_image, status_msg],
    )
    
    clear_btn.click(
        fn=lambda: (None, None, "Ready to process your image..."),
        outputs=[input_image, output_image, status_msg],
    )

# Launch the app
demo.launch()