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

# First, ensure all required dependencies are installed
try:
    import torchvision
    import skimage
except ImportError:
    print("Installing required packages...")
    import subprocess
    subprocess.check_call(["pip", "install", "torchvision", "scikit-image"])
    import torchvision
    import skimage

def remove_background(input_image):
    try:
        # Initialize the pipeline with correct parameters and dependencies
        segmentor = pipeline(
            "image-segmentation", 
            model="briaai/RMBG-1.4",
            trust_remote_code=True,
            device="cpu",
            framework="pt"
        )
        
        # 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 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()