Create app.py

app.py

@@ -0,0 +1,126 @@
+import io, os, json, time
+import streamlit as st
+import numpy as np
+from PIL import Image
+import torch
+from huggingface_hub import hf_hub_download
+
+# ────────── impostazioni Streamlit
+st.set_page_config(page_title="♻️ UPA – Upcycling Parts Analyzer",
+                   page_icon="♻️", layout="wide")
+torch.set_float32_matmul_precision("medium")        # meno VRAM fp32
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+st.sidebar.info(f"Device attivo: **{DEVICE}**")
+
+# ────────── controlli UI
+sam_backbone = st.sidebar.selectbox(
+    "Backbone SAM", ("vit_h", "vit_l", "vit_b"), index=0,
+    help="Con la GPU T4 puoi tenere vit_h; vit_b è più leggero per CPU."
+)
+points_per_side = st.sidebar.slider("Punti per lato", 0, 128, 32)
+iou_th   = st.sidebar.slider("Soglia IoU",          0.0, 1.0, 0.8)
+stab_th  = st.sidebar.slider("Soglia Stabilità",    0.0, 1.0, 0.9)
+min_area = st.sidebar.number_input("Area minima (px)", 0, 10_000, 100)
+
+# ────────── helper model loader
+@st.cache_resource(show_spinner="📥 Carico SAM …")
+def load_sam(backbone: str):
+    from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+    repo, fname = {
+        "vit_h": ("sam-vit-h", "sam_vit_h_4b8939.pth"),
+        "vit_l": ("sam-vit-l", "sam_vit_l_0b3195.pth"),
+        "vit_b": ("sam-vit-b", "sam_vit_b_01ec64.pth"),
+    }[backbone]
+    ckpt = hf_hub_download(repo_id=f"facebook/{repo}", filename=fname)
+    sam  = sam_model_registry[backbone](checkpoint=ckpt).to(DEVICE).eval()
+    if DEVICE.type == "cuda":
+        sam = torch.compile(sam)
+    gen = SamAutomaticMaskGenerator(
+        sam,
+        points_per_side=points_per_side,
+        pred_iou_thresh=iou_th,
+        stability_score_thresh=stab_th,
+        crop_n_layers=1,
+        crop_n_points_downscale_factor=2
+    )
+    return gen
+mask_generator = load_sam(sam_backbone)
+
+@st.cache_resource(show_spinner="🔤 Carico CLIP …")
+def load_clip():
+    from transformers import CLIPModel, CLIPProcessor
+    name  = "laion/CLIP-ViT-L-14-laion2B-s32B-b82K"
+    proc  = CLIPProcessor.from_pretrained(name)
+    model = CLIPModel.from_pretrained(name).to(DEVICE).eval()
+    return model, proc
+clip_model, clip_proc = load_clip()
+
+# ────────── funzioni utility
+MAX_SIDE = 1024
+LABELS = ["lamiera", "foro circolare", "scanalatura rettangolare"]
+
+def preprocess(img: Image.Image) -> Image.Image:
+    if max(img.size) > MAX_SIDE:
+        img = img.copy()
+        img.thumbnail((MAX_SIDE, MAX_SIDE))
+    return img
+
+def classify_piece(np_img: np.ndarray, mask_u8: np.ndarray):
+    from cv2 import bitwise_and
+    crop = Image.fromarray(bitwise_and(np_img, np_img, mask=mask_u8))
+    inputs = clip_proc(text=LABELS, images=crop, return_tensors="pt", padding=True)
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+    with torch.no_grad():
+        logits = clip_model(**inputs).logits_per_image
+        probs  = logits.softmax(dim=1)
+        conf, idx = torch.max(probs, 1)
+        return LABELS[idx.item()], float(conf.item())
+
+def draw_overlay(np_img: np.ndarray, results):
+    import cv2
+    overlay = np_img.copy()
+    rng = np.random.default_rng(seed=42)
+    for lbl, _, mask in results:
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        color = tuple(int(c) for c in rng.integers(50, 255, 3))
+        cv2.drawContours(overlay, contours, -1, color, 2)
+        # scrive la label sul centroide
+        for cnt in contours:
+            M = cv2.moments(cnt)
+            if M["m00"]:
+                cx, cy = int(M["m10"]/M["m00"]), int(M["m01"]/M["m00"])
+                cv2.putText(overlay, lbl, (cx, cy), cv2.FONT_HERSHEY_SIMPLEX,
+                            0.6, color, 2, cv2.LINE_AA)
+    return overlay
+
+# ────────── interfaccia principale
+st.title("🔍 Riconoscimento Parti Meccaniche")
+files = st.file_uploader("Carica immagini JPG / PNG", type=["jpg","jpeg","png"],
+                         accept_multiple_files=True)
+
+if files:
+    if st.button("Segmenta + Classifica"):
+        for f in files:
+            img_pil = preprocess(Image.open(f).convert("RGB"))
+            np_img  = np.array(img_pil)
+
+            with st.spinner(f"Segmentazione {f.name}…"):
+                t0 = time.time()
+                masks = [m for m in mask_generator.generate(np_img) if m["area"] >= min_area]
+                t1 = time.time()
+
+            results = []
+            for m in masks:
+                mask_u8 = (m["segmentation"] * 255).astype(np.uint8)
+                lbl, conf = classify_piece(np_img, mask_u8)
+                results.append((lbl, conf, mask_u8))
+
+            col1, col2 = st.columns([2,1])
+            with col1:
+                overlay = draw_overlay(np_img, results)
+                st.image(overlay, caption=f"{f.name} – {len(results)} pezzi ∙ {t1-t0:.1f}s",
+                         use_container_width=True)
+            with col2:
+                st.subheader("Dettagli")
+                for lbl, cf, _ in results:
+                    st.write(f"{lbl}: {cf:.2%}")