File size: 1,626 Bytes
05d99b1
4877588
05d99b1
 
4877588
c79bc95
4877588
 
 
 
 
 
 
f6c4736
4877588
f6c4736
 
e2aec45
 
 
05d99b1
7388e1c
c79bc95
 
7388e1c
05d99b1
 
7388e1c
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
40
41
42
43
44
45
46
import logging
import gradio as gr
from rembg import new_session
from cutter import remove, make_label
from utils import *
from PIL import Image

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 = (0, 0, 0)  # Color de fondo predeterminado (negro)
    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)  # Convertir el array de numpy a una imagen PIL
        return result
    except ValueError as err:
        logging.error(err)
        return 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)