File size: 1,264 Bytes
05d99b1 4877588 05d99b1 4877588 f6c4736 4877588 f6c4736 05d99b1 04afab2 05d99b1 4877588 51a33fe 4877588 5b68504 57135e2 005d5cc 790e21a 4877588 490d964 f6c4736 790e21a 4877588 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import logging
import gradio as gr
from rembg import new_session
from cutter import remove, make_label
from utils import *
remove_bg_models = {
"U2NET": "u2net",
"U2NET Human Seg": "u2net_human_seg",
"U2NET Cloth Seg": "u2net_cloth_seg"
}
default_model = "U2NET"
def predict(image):
session = new_session(remove_bg_models[default_model])
try:
return remove(session, image, False, None, None)
except ValueError as err:
logging.error(err)
return make_label(str(err)), None
with gr.Blocks(title="Remove background") as app:
gr.HTML("<center><h1>Background Remover</h1></center>")
with gr.Row(equal_height=False):
with gr.Column():
input_img = gr.Image(type="pil", label="Input image")
with gr.Column():
output_img = gr.Image(type="pil", label="Result image")
with gr.Row(equal_height=True):
run_btn = gr.Button(value="Remove background", variant="primary")
clear_btn = gr.Button(value="Clear", variant="secondary")
run_btn.click(predict, inputs=[input_img], outputs=[output_img])
clear_btn.click(lambda: (None, None), inputs=None, outputs=[input_img, output_img])
app.launch(share=False, debug=True, enable_queue=True, show_error=True)
|