File size: 1,693 Bytes
05d99b1 7d38977 4877588 05d99b1 4877588 f6c4736 4877588 f6c4736 e2aec45 e110031 05d99b1 e110031 c79bc95 7d38977 7388e1c 05d99b1 e110031 05d99b1 9c77858 51a33fe 4877588 5b68504 57135e2 005d5cc 790e21a e110031 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 40 41 42 43 44 45 46 |
import logging
import numpy as np
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])
smoot = False
matting = (0, 0, 0) # Valores predeterminados para matting
bg_color = False # Color de fondo predeterminado (no cambiar color)
try:
result, _ = remove(session, image, smoot, matting, bg_color)
if isinstance(result, np.ndarray): # Verificar si la salida es un array de numpy
result = Image.fromarray(result.astype('uint8')) # Convertir el array de numpy a una imagen PIL
return result
except ValueError as err:
logging.error(err)
return make_label(str(err)), None
with gr.Blocks(css="custom.css", 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)
|