import gradio as gr
import numpy as np
from PIL import Image
from rembg import remove, new_session

# Initialize session with proper settings to prevent cropping
session = new_session("u2net")
bg_removal_kwargs = {
    "alpha_matting": False,  # Disable advanced features that cause cropping
    "session": session,
    "only_mask": False,
    "post_process_mask": True  # Clean edges without cropping
}

def remove_background(input_image):
    try:
        # Convert any input to PIL Image
        if isinstance(input_image, np.ndarray):
            img = Image.fromarray(input_image)
        elif isinstance(input_image, dict):  # Handle paste/drop
            img = Image.open(input_image["name"])
        else:
            img = input_image
        
        # Preserve original size (disable auto-resizing)
        result = remove(img, **bg_removal_kwargs)
        return result
    
    except Exception as e:
        print(f"Error: {str(e)}")
        return input_image  # Return original if fails

# Gradio interface with proper image handling
with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ Background Remover (No Cropping)")
    
    with gr.Row():
        input_img = gr.Image(
            label="Original", 
            type="pil",  # Ensures we get PIL Images
            height=400
        )
        output_img = gr.Image(
            label="Result", 
            type="pil",
            height=400
        )
    
    gr.Button("Remove Background").click(
        remove_background,
        inputs=input_img,
        outputs=output_img
    )

demo.launch()