import gradio as gr
import rembg
from PIL import Image
from io import BytesIO

def remove_background(image):
    """
    Removes the background from an uploaded image using rembg.
    """
    # Ensure image is in RGBA mode
    image = image.convert("RGBA")
    
    # Remove background
    output_image = rembg.remove(image)

    return output_image  # Directly return PIL image

# Gradio Interface
demo = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Image(type="pil", label="Processed Image with Transparent Background"),
    title="AI Image Background Remover",
    description="Upload an image and remove its background using AI-powered rembg.",
)

demo.launch()