import gradio as gr
from rembg import remove
from PIL import Image

def remove_background(image):
    """
    Takes a PIL image as input and returns the image with background removed.
    """
    return remove(image)

# Create the Gradio interface
app = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Image(type="pil", label="Image with Background Removed"),
    title="Background Removal App",
    description="Upload any image to remove its background using the U²-Net model.",
    examples=[
        ["example_images/person.jpg"],
        ["example_images/object.png"]
    ],
    cache_examples=False,
    theme="default"
)

# Launch the app
if __name__ == "__main__":
    app.launch()