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

def remove_bg(image):
    try:
        img = Image.open(image) if isinstance(image, str) else Image.fromarray(image)
        img_byte_arr = io.BytesIO()
        img.save(img_byte_arr, format='PNG')
        output_bytes = remove(img_byte_arr.getvalue())
        return Image.open(io.BytesIO(output_bytes))
    except Exception as e:
        raise gr.Error(f"Error processing image: {str(e)}")

with gr.Blocks() as demo:
    gr.Markdown("# Background Remover")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="filepath", label="Input Image")
            submit_btn = gr.Button("Remove Background")
        with gr.Column():
            output_image = gr.Image(label="Output Image", type="pil")
    
    submit_btn.click(remove_bg, inputs=input_image, outputs=output_image)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)