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

def remove_background(input_image):
    # Convert Gradio input image data to bytes
    image_array = input_image.astype(np.uint8)  # Ensure the image is in the correct format
    image_pil = Image.fromarray(image_array)

    # Save the PIL Image to an in-memory buffer
    image_buffer = io.BytesIO()
    image_pil.save(image_buffer, format='PNG')

    # Use rembg to remove the background
    output_data = rembg.remove(image_buffer.getvalue())

    # Convert the output data to a PIL Image
    output_image = Image.open(io.BytesIO(output_data))

    # Return the output image directly
    return np.array(output_image)

iface = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(),
    outputs=gr.Image()
)

iface.launch()