import gradio as gr
from transparent_background import Remover

# Load model
remover = Remover()  # default setting
remover = Remover(mode='fast', jit=True, device='cpu', ckpt='latest.pth')  # custom setting
remover = Remover(mode='base-nightly')  # nightly release checkpoint


def run(inp_image, sod_type):
    return remover.process(inp_image.convert("RGB"), type=sod_type)


with gr.Blocks() as demo:
    with gr.Row():
        inp_image = gr.Image(type="pil")
        res_image = gr.Image(type="pil")
    sod_type = gr.Radio(["rgba", "map", "green", "white", "blur", "overlay"], value="rgba")
    submit_btn = gr.Button(value="Submit")
    submit_btn.click(run, inputs=[inp_image, sod_type], outputs=[res_image])

if __name__ == "__main__":
    demo.launch(debug=False, server_name="0.0.0.0", server_port=8189)