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

def remove_background(image):
    output = remove(image)
    white_bg = Image.new("RGBA", output.size, (255, 255, 255, 255))
    white_bg.paste(output, mask=output.split()[3])
    return white_bg.convert("RGB")

demo = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(type="pil", label="Garment Image"),
    outputs=gr.Image(type="pil", label="Clean Garment"),
    api_name="remove_background"
)

demo.launch()