background / app.py
0notexist0's picture
Update app.py
4b267d2 verified
raw
history blame contribute delete
883 Bytes
import gradio as gr
from transformers import pipeline
from PIL import Image
import numpy as np
# 1) Carica il modello di segmentazione
seg = pipeline("image-segmentation", model="BritishWerewolf/U-2-Netp", framework="pt")
def remove_bg(image: Image.Image):
# 2) Ottieni la maschera
outputs = seg(image)
mask = outputs[0]["mask"].astype(np.uint8) * 255
# 3) Applica la maschera all’alpha channel
rgba = image.convert("RGBA")
rgba_arr = np.array(rgba)
rgba_arr[..., 3] = mask # canale alpha = mask
return Image.fromarray(rgba_arr)
# 4) Definisci l’interfaccia “chat”
with gr.Blocks() as demo:
chatbot = gr.Chatbot(label="BG-Removal Bot")
img_in = gr.Image(type="pil", label="Invia la tua foto")
img_out = gr.Image(type="pil", label="Risultato")
img_in.change(fn=remove_bg, inputs=img_in, outputs=img_out)
demo.launch()