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

def remove_background(input_image):
    # Convert to PIL Image if not already
    if not isinstance(input_image, Image.Image):
        input_image = Image.fromarray(input_image)
    
    output_image = remove(input_image)
    return output_image

# Create the Gradio interface using the new syntax
interface = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Remove Background",
    description="This App removes the background from an image",
    examples=[
        "examples/input/1.jpeg",
        "examples/input/2.jpeg",
        "examples/input/3.jpeg",
    ],
    cache_examples=True,
)

if __name__ == "__main__":
    interface.launch()