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()