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

# -------------------------------
# Function to remove background
# -------------------------------
def remove_bg(image: Image.Image):
    if image is None:
        return None

    # Convert PIL image to bytes
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    img_bytes = img_byte_arr.getvalue()

    # Remove background using rembg
    output_bytes = remove(img_bytes)

    # Convert bytes back to PIL image
    output_image = Image.open(io.BytesIO(output_bytes))

    return output_image

# -------------------------------
# Gradio UI
# -------------------------------
demo = gr.Interface(
    fn=remove_bg,
    inputs=gr.Image(type="pil", label="Upload any image"),
    outputs=gr.Image(type="pil", label="Background Removed"),
    title="Universal Background Remover",
    description="Upload any image and the background will be automatically removed using rembg.",
    allow_flagging="never",
)

# -------------------------------
# Launch
# -------------------------------
if __name__ == "__main__":
    demo.launch()