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

# Function to remove the background
def remove_background(image):
    # Convert the uploaded image to bytes
    img_bytes = io.BytesIO()
    image.save(img_bytes, format="PNG")
    img_bytes = img_bytes.getvalue()
    
    # Use rembg to remove the background
    output_bytes = remove(img_bytes)
    
    # Convert the output bytes back to an image
    output_image = Image.open(io.BytesIO(output_bytes)).convert("RGBA")
    return output_image

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Background Remover AI")
    gr.Markdown("Upload an image, and the AI will remove the background for you.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Upload Image")
        with gr.Column():
            output_image = gr.Image(type="pil", label="Output Image")
    
    remove_button = gr.Button("Remove Background")
    
    # Link the button to the function
    remove_button.click(remove_background, inputs=input_image, outputs=output_image)

# Launch the app
demo.launch()