Update app.py

app.py

@@ -1,18 +1,28 @@
-from transformers import AutoProcessor, AutoModelForImageSegmentation
+import gradio as gr
+from transformers import pipeline
 from PIL import Image
-import torch
+import numpy as np
 
-# Carica modello e processor
-processor = AutoProcessor.from_pretrained("BritishWerewolf/U-2-Netp")
-model     = AutoModelForImageSegmentation.from_pretrained("BritishWerewolf/U-2-Netp")
+# 1) Carica il modello di segmentazione
+seg = pipeline("image-segmentation", model="BritishWerewolf/U-2-Netp", framework="pt")
 
-# Prepara l’immagine
-img = Image.open("input.jpg").convert("RGB")
-inputs = processor(images=img, return_tensors="pt")
+def remove_bg(image: Image.Image):
+    # 2) Ottieni la maschera
+    outputs = seg(image)
+    mask = outputs[0]["mask"].astype(np.uint8) * 255
 
-# Inferenzia maschera
-with torch.no_grad():
-    outputs = model(**inputs)
-    mask = outputs.logits.argmax(dim=1)[0].cpu().numpy()
+    # 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)
 
-# Applica maschera all’immagine originale...
+# 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()