Update app.py

app.py

@@ -1,49 +1,96 @@
-import gradio as gr
+import streamlit as st
 import numpy as np
-from PIL import Image, ImageDraw
+import torch
+from PIL import Image
+import cv2
+import matplotlib.pyplot as plt
 from huggingface_hub import hf_hub_download
-# carica SAM e GroundingDINO direttamente dall’HF Hub
-SAM_CHECKPOINT = hf_hub_download("facebook/sam-vit-base", "sam_vit_b.pth")
-GDINO_CONFIG    = hf_hub_download("IDEA-Research/GroundingDINO", "GroundingDINO_SwinT_OGC.py")
-GDINO_CHECKPT   = hf_hub_download("IDEA-Research/GroundingDINO", "groundingdino_swint_ogc.pth")
-
-# — pseudocodice di import —
-from segment_anything import sam_model
-from groundingdino.util.inference import load_model, predict
-
-sam    = sam_model.load_from_checkpoint(SAM_CHECKPOINT)
-gdino  = load_model(GDINO_CONFIG, GDINO_CHECKPT)
-
-def recognize(img, prompt, conf):
-    masks = sam.segment(img)       # 1. segmentazione zero‑shot
-    out   = Image.fromarray(img).convert("RGBA")
-    draw  = ImageDraw.Draw(out, "RGBA")
-    results = []
-    for m in masks:
-        label, score = predict(gdino, img, m, prompt)  # 2. classificazione zero‑shot
-        if score<conf: continue
-        yy, xx = np.where(m)
-        bbox = (xx.min(), yy.min(), xx.max(), yy.max())
-        area = int(m.sum())
-        draw.rectangle(bbox, outline=(255,0,0,180), width=3)
-        draw.text((bbox[0], bbox[1]-10), f"{label} {score:.2f}", fill=(255,0,0,180))
-        results.append({"label":label, "score":score, "area":area, "bbox":bbox})
-    return np.array(out), results
-
-app = gr.Interface(
-    fn=recognize,
-    inputs=[
-      gr.Image(type="numpy", label="Upload Image"),
-      gr.Textbox(label="Prompt (comma‑separated)"),
-      gr.Slider(0,1,0.25, label="Confidence Threshold"),
-    ],
-    outputs=[
-      gr.Image(label="Overlay"),
-      gr.JSON(label="Detections")
-    ],
-    title="Zero‑Shot Component Recognition",
-    description="Segmenta e classifica componenti meccanici da foto, senza training specifico."
-)
-
-if __name__=="__main__":
-    app.launch()
+from segment_anything import SamPredictor, sam_model_registry
+from groundingdino.util.inference import load_model, predict, annotate
+
+# Titolo dell'app
+st.title("🔍 Riconoscimento Zero-Shot con GroundingDINO + SAM")
+
+# Configurazione dei modelli da Hugging Face Hub
+@st.cache_resource
+def load_sam():
+    checkpoint = hf_hub_download(
+        repo_id="SegmentAnything/sam_vit_b",
+        filename="sam_vit_b_01ec64.pth"
+    )
+    model = sam_model_registry["vit_b"](checkpoint=checkpoint)
+    return SamPredictor(model.to("cuda" if torch.cuda.is_available() else "cpu"))
+
+@st.cache_resource
+def load_grounding_dino():
+    config_path = hf_hub_download(
+        repo_id="IDEA-Research/grounding-dino-tiny",
+        filename="GroundingDINO_SwinT_OGC.py"
+    )
+    checkpoint_path = hf_hub_download(
+        repo_id="IDEA-Research/grounding-dino-tiny",
+        filename="groundingdino_tiny.pth"
+    )
+    model = load_model(config_path, checkpoint_path)
+    return model
+
+sam = load_sam()
+grounding_dino = load_grounding_dino()
+
+# Caricamento immagine da parte dell'utente
+uploaded_image = st.file_uploader("📷 Carica un'immagine", type=['jpg', 'jpeg', 'png'])
+
+prompt = st.text_input("📝 Inserisci le classi da riconoscere (separate da virgola)", 
+                       value="lamiera, foro circolare, vite, bullone, scanalatura")
+
+if uploaded_image is not None:
+    image = Image.open(uploaded_image).convert("RGB")
+    img_array = np.array(image)
+
+    st.image(image, caption="Immagine caricata", use_column_width=True)
+
+    if st.button("▶️ Avvia riconoscimento"):
+        # GroundingDINO prediction
+        boxes, logits, phrases = predict(
+            model=grounding_dino,
+            image=img_array,
+            caption=prompt,
+            box_threshold=0.3,
+            text_threshold=0.25,
+            device="cuda" if torch.cuda.is_available() else "cpu"
+        )
+
+        annotated_frame = annotate(image_source=img_array, boxes=boxes, logits=logits, phrases=phrases)
+
+        st.subheader("Risultato GroundingDINO")
+        st.image(annotated_frame, caption="Annotazione GroundingDINO")
+
+        # SAM segmentation
+        sam.set_image(img_array)
+        H, W, _ = img_array.shape
+        boxes_scaled = boxes * torch.tensor([W, H, W, H], device=boxes.device)
+        boxes_scaled = boxes_scaled.cpu().numpy()
+
+        masks, scores, _ = sam.predict_torch(
+            point_coords=None,
+            point_labels=None,
+            boxes=torch.tensor(boxes_scaled, device=sam.device),
+            multimask_output=False,
+        )
+
+        # Visualizza maschere segmentate
+        st.subheader("Risultato Segment Anything (SAM)")
+        plt.figure(figsize=(10, 10))
+        plt.imshow(img_array)
+        for mask in masks:
+            mask_np = mask[0].cpu().numpy()
+            plt.contour(mask_np, colors='red', linewidths=1.5)
+        plt.axis('off')
+
+        st.pyplot(plt.gcf())
+        plt.close()
+
+        # Tabella risultati
+        st.subheader("🔖 Tabella Risultati")
+        result_data = [{"Classe": phrase, "Confidenza": round(logit.item(), 2)} for phrase, logit in zip(phrases, logits)]
+        st.table(result_data)