Update app.py

app.py

@@ -1,12 +1,23 @@
 import os
+import io
 import time
+import traceback
+from typing import Optional, Tuple
+
 import gradio as gr
 from PIL import Image
+
 import torch
 import torchvision.transforms as T
 import torchvision.models as models
-from openai import OpenAI
-import spaces  # per ZeroGPU
+
+try:
+    from openai import OpenAI
+except Exception:
+    OpenAI = None  # gestiamo assenza pacchetto elegantemente
+
+import spaces  # ZeroGPU decorator
+
 
 # ======================
 # Config / Model / Classes
@@ -23,7 +34,7 @@ IDX2LABEL = {
     5: "pitting_corrosion",
     6: "stress_corrosion",
     7: "under_insulation_corrosion",
-    8: "uniform_corrosion"
+    8: "uniform_corrosion",
 }
 
 ZONES = [
@@ -36,9 +47,13 @@ ZONES = [
     "Cargo holds / Dry bulk",
     "Engine room / Hot surfaces",
     "Pipes / Under insulation (UIC/CUI)",
-    "Other / Not sure"
+    "Other / Not sure",
 ]
 
+# ======================
+# Model load (CPU default)
+# ======================
+
 def load_model_cpu():
     m = models.resnet50(weights=None)
     num_ftrs = m.fc.in_features
@@ -48,6 +63,7 @@ def load_model_cpu():
     m.eval()
     return m
 
+print("[BOOT] Loading model on CPU…")
 model_cpu = load_model_cpu()
 
 transform = T.Compose([
@@ -59,24 +75,32 @@ transform = T.Compose([
 ])
 
 # ======================
-# OpenAI Assistant
+# OpenAI Assistant (optional)
 # ======================
 
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-if not OPENAI_API_KEY:
-    print("WARNING: OPENAI_API_KEY not set. Add it in HF Space Secrets.")
-
-client = OpenAI(api_key=OPENAI_API_KEY)
-
 ASSISTANT_ID = os.environ.get("PPG_ASSISTANT_ID", "asst_20DNMEENkfBsYupFjPCwfijZ")
 VECTOR_STORE_ID = os.environ.get("PPG_VECTOR_STORE_ID", "")
 APP_FORCE_LANG = os.environ.get("APP_FORCE_LANG", "").strip()
 
+client = None
+assistant_enabled = False
+if OPENAI_API_KEY and OpenAI is not None:
+    try:
+        client = OpenAI(api_key=OPENAI_API_KEY)
+        assistant_enabled = True
+        print("[BOOT] OpenAI client initialized.")
+    except Exception as e:
+        print("[BOOT][WARN] OpenAI init failed:", e)
+
+def _assistant_safe() -> bool:
+    return bool(assistant_enabled and client is not None and ASSISTANT_ID)
+
 # ======================
 # Inference utils (CPU/GPU)
 # ======================
 
-def predict_on_cpu(img_pil: Image.Image):
+def predict_on_cpu(img_pil: Image.Image) -> Tuple[str, float]:
     x = transform(img_pil.convert("RGB")).unsqueeze(0)
     with torch.no_grad():
         logits = model_cpu(x)
@@ -85,7 +109,7 @@ def predict_on_cpu(img_pil: Image.Image):
     return IDX2LABEL.get(idx, f"class_{idx}"), float(probs[idx])
 
 @spaces.GPU(duration=60)
-def predict_on_gpu(img_pil: Image.Image):
+def predict_on_gpu(img_pil: Image.Image) -> Tuple[str, float]:
     device = "cuda"
     m = models.resnet50(weights=None)
     num_ftrs = m.fc.in_features
@@ -101,134 +125,221 @@ def predict_on_gpu(img_pil: Image.Image):
     idx = int(probs.argmax())
     return IDX2LABEL.get(idx, f"class_{idx}"), float(probs[idx])
 
-def predict_image(image: Image.Image):
+def predict_image(image: Image.Image) -> Tuple[str, float]:
     try:
         if torch.cuda.is_available():
             return predict_on_gpu(image)
-    except Exception:
-        pass
+    except Exception as e:
+        print("[GPU][WARN] Falling back to CPU:", e)
     return predict_on_cpu(image)
 
 # ======================
-# Assistant calls
+# Assistant calls (with optional image)
 # ======================
 
-def call_assistant(label, confidence, zone, note, user_question, thread_id=None):
-    if not thread_id:
-        if VECTOR_STORE_ID:
-            thread = client.beta.threads.create(
-                tool_resources={"file_search": {"vector_store_ids": [VECTOR_STORE_ID]}}
-            )
-        else:
-            thread = client.beta.threads.create()
-        thread_id = thread.id
-
-    core_context = f"""
-Classification: {label} ({round(confidence*100,2)}%).
-Zone: {zone or "Not specified"}.
-User note: {note or "(none)"}.
-"""
-    user_payload = core_context + "\nUser question:\n" + (user_question or "Provide initial advisory.")
+def call_assistant(
+    label: str,
+    confidence: float,
+    zone: str,
+    note: str,
+    user_question: str,
+    image: Optional[Image.Image],
+    thread_id: Optional[str] = None,
+    max_wait_s: int = 45,
+) -> Tuple[str, str]:
+    """
+    Ritorna (reply_text, thread_id). Non lancia eccezioni.
+    """
+    if not _assistant_safe():
+        return ("[Assistant disabled] No OPENAI_API_KEY or client not available. "
+                "Model classification shown above.", thread_id or "")
 
-    client.beta.threads.messages.create(
-        thread_id=thread_id,
-        role="user",
-        content=user_payload,
-    )
-
-    second_lang_clause = f"Then provide the same content in {APP_FORCE_LANG}." if APP_FORCE_LANG else \
-                         "Then repeat in the user's language if detectable from note; else in Italian."
+    try:
+        # crea thread se serve
+        if not thread_id:
+            if VECTOR_STORE_ID:
+                thread = client.beta.threads.create(
+                    tool_resources={"file_search": {"vector_store_ids": [VECTOR_STORE_ID]}}
+                )
+            else:
+                thread = client.beta.threads.create()
+            thread_id = thread.id
+
+        # prepara testo
+        core_context = (
+            f"Classification: {label} ({round(confidence*100,2)}%).\n"
+            f"Zone: {zone or 'Not specified'}.\n"
+            f"User note: {note or '(none)'}.\n"
+        )
+        user_payload = core_context + "\nUser question:\n" + (user_question or "Provide initial advisory.")
+
+        # costruisci contenuto multi-part con immagine
+        content = [{"type": "input_text", "text": user_payload}]
+
+        if image is not None:
+            buf = io.BytesIO()
+            image.convert("RGB").save(buf, format="PNG", optimize=True)
+            buf.seek(0)
+            uploaded = client.files.create(file=buf, purpose="assistants")
+            content.append({"type": "input_image", "image_file": {"file_id": uploaded.id}})
+
+        client.beta.threads.messages.create(
+            thread_id=thread_id,
+            role="user",
+            content=content
+        )
+
+        second_lang_clause = (
+            f"Then provide the same content in {APP_FORCE_LANG}."
+            if APP_FORCE_LANG else
+            "Then repeat in the user's language if detectable from note; else in Italian."
+        )
+
+        extra_instructions = (
+            "Act as a PPG marine coatings technical specialist for ships (marine environments only). "
+            "Answer ONLY using information found in the attached docs via File Search. "
+            "If docs lack details, reply 'Not in docs'. "
+            "ALWAYS ask for the zone if missing before prescribing. "
+            "Structure: Diagnosis; Surface Preparation; System; Notes; Disclaimer. "
+            "Provide first in English. " + second_lang_clause
+        )
+
+        run = client.beta.threads.runs.create(
+            thread_id=thread_id,
+            assistant_id=ASSISTANT_ID,
+            instructions=extra_instructions,
+        )
+
+        # polling con timeout hard
+        t0 = time.time()
+        while True:
+            r = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
+            if r.status in ["completed", "failed", "cancelled", "expired"]:
+                break
+            if time.time() - t0 > max_wait_s:
+                print("[Assistant][WARN] Timeout waiting run.")
+                break
+            time.sleep(0.7)
+
+        msgs = client.beta.threads.messages.list(thread_id=thread_id)
+        reply = None
+        for m in msgs.data:
+            if m.role == "assistant":
+                for part in m.content:
+                    if getattr(part, "type", "") == "text":
+                        reply = part.text.value
+                        break
+                if reply:
+                    break
 
-    extra_instructions = (
-        "Act as a PPG marine coatings technical specialist for ships (marine environments only). "
-        "Answer ONLY using information found in the attached docs via File Search. "
-        "If docs lack details, reply 'Not in docs'. "
-        "ALWAYS ask for the zone if missing before prescribing. "
-        "Structure: Diagnosis; Surface Preparation; System; Notes; Disclaimer. "
-        "Provide first in English. " + second_lang_clause
-    )
+        if not reply:
+            reply = "[Assistant] No reply received."
 
-    run = client.beta.threads.runs.create(
-        thread_id=thread_id,
-        assistant_id=ASSISTANT_ID,
-        instructions=extra_instructions,
-    )
+        return reply, thread_id
 
-    while True:
-        r = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
-        if r.status in ["completed", "failed", "cancelled", "expired"]:
-            break
-        time.sleep(0.6)
-
-    msgs = client.beta.threads.messages.list(thread_id=thread_id)
-    reply = None
-    for m in msgs.data:
-        if m.role == "assistant":
-            for part in m.content:
-                if getattr(part, "type", "") == "text":
-                    reply = part.text.value
-                    break
-            if reply:
-                break
-    return reply or "No reply from Assistant.", thread_id
+    except Exception as e:
+        print("[Assistant][ERROR]", e)
+        traceback.print_exc()
+        return ("[Assistant error] " + str(e) + "\nProceed using model result only.", thread_id or "")
 
 # ======================
-# Pipelines (generator) con barra di avanzamento
+# Pipelines (generator) SAFE
 # ======================
 
 def run_analysis(image, note, zone, chat_history, thread_state):
-    # La progress bar mostra avanzamenti deterministici lato server.
+    """
+    Generator sicuro: intercetta ogni eccezione e restituisce sempre qualcosa,
+    così l’overlay non resta appeso e l’utente non rimane a fissare il vuoto cosmico.
+    """
     with gr.Progress() as prog:
-        prog(0.02, desc="Ricezione immagine")
-        if image is None:
-            yield "No image received.", chat_history, thread_state
-            return
-
-        prog(0.06, desc="Validazione input")
-        if not zone or zone == "Other / Not sure":
-            yield "**Please select the area/zone first.**", chat_history, thread_state
-            return
-
-        # messaggio intermedio per feedback immediato
-        yield "**Analyzing image...** Please wait.", chat_history, thread_state
-
-        prog(0.15, desc="Preprocessing")
-        # un micro ritardo solo per dare feedback visivo
-        time.sleep(0.05)
-
-        prog(0.45, desc="Classificazione (CNN)")
-        label, conf = predict_image(image)
+        try:
+            prog(0.03, desc="Checking input")
+            if image is None:
+                yield "No image received.", chat_history, thread_state
+                return
+
+            if not zone or zone == "Other / Not sure":
+                yield "**Please select the area/zone first.**", chat_history, thread_state
+                return
+
+            # feedback immediato
+            yield "**Analyzing image...** Please wait.", chat_history, thread_state
+
+            prog(0.18, desc="Preprocessing")
+            time.sleep(0.05)
+
+            prog(0.50, desc="Classifying (ResNet50)")
+            label, conf = predict_image(image)
+
+            prog(0.72, desc="Consulting PPG Assistant")
+            reply, thread_id = call_assistant(
+                label=label,
+                confidence=conf,
+                zone=zone,
+                note=note or "",
+                user_question="Provide initial advisory.",
+                image=image,
+                thread_id=(thread_state or {}).get("thread_id")
+            )
 
-        prog(0.70, desc="Consulto PPG Assistant")
-        reply, thread_id = call_assistant(label, conf, zone, note, "Provide initial advisory.")
+            header = f"**Model result:** `{label}` — confidence **{round(conf*100,2)}%**\n\n"
+            out_text = header + (reply or "")
+            new_history = (chat_history[:] if chat_history else [])
+            if reply:
+                new_history.append(("", reply))
 
-        prog(0.98, desc="Composizione risposta")
+            prog(1.0, desc="Done")
 
-        header = f"**Model result:** `{label}` — confidence **{round(conf*100,2)}%**\n\n"
-        out_text = header + (reply or "")
-        new_history = chat_history[:] if chat_history else []
-        new_history.append(("", reply))
-        prog(1.0, desc="Fatto")
+            yield out_text, new_history, {
+                "thread_id": thread_id,
+                "label": label,
+                "confidence": conf,
+                "zone": zone or "",
+            }
 
-        yield out_text, new_history, {"thread_id": thread_id, "label": label, "confidence": conf, "zone": zone or ""}
+        except Exception as e:
+            # non lasciamo l’overlay attivo in eterno
+            print("[Pipeline][ERROR]", e)
+            traceback.print_exc()
+            err = f"**Error during analysis**:\n```\n{e}\n```\nCheck logs/keys and try again."
+            yield err, chat_history, thread_state or {}
 
 def continue_chat(user_msg, chat_history, thread_state, note, zone):
-    if not user_msg.strip():
+    if not user_msg or not user_msg.strip():
         return chat_history, ""
     with gr.Progress() as prog:
-        prog(0.1, desc="Invio messaggio")
-        label = (thread_state or {}).get("label") or "unknown"
-        conf = (thread_state or {}).get("confidence") or 0.0
-        current_zone = zone or (thread_state or {}).get("zone") or "Not specified"
-        thread_id = (thread_state or {}).get("thread_id")
-
-        prog(0.6, desc="Consulto PPG Assistant")
-        reply, thread_id = call_assistant(label, conf, current_zone, note, user_msg, thread_id)
-
-        chat_history.append((user_msg, reply))
-        thread_state["thread_id"] = thread_id
-        prog(1.0, desc="Fatto")
-        return chat_history, ""
+        try:
+            prog(0.2, desc="Sending")
+            label = (thread_state or {}).get("label") or "unknown"
+            conf = (thread_state or {}).get("confidence") or 0.0
+            current_zone = zone or (thread_state or {}).get("zone") or "Not specified"
+            thread_id = (thread_state or {}).get("thread_id")
+
+            prog(0.7, desc="Consulting PPG Assistant")
+            reply, thread_id = call_assistant(
+                label=label,
+                confidence=conf,
+                zone=current_zone,
+                note=note or "",
+                user_question=user_msg,
+                image=None,  # la thread ha già il file dell'ultima analisi
+                thread_id=thread_id
+            )
+
+            chat_history = chat_history or []
+            chat_history.append((user_msg, reply))
+            if isinstance(thread_state, dict):
+                thread_state["thread_id"] = thread_id
+
+            prog(1.0, desc="Done")
+            return chat_history, ""
+
+        except Exception as e:
+            print("[Chat][ERROR]", e)
+            traceback.print_exc()
+            chat_history = chat_history or []
+            chat_history.append((user_msg, f"[Error] {e}"))
+            return chat_history, ""
 
 # ======================
 # UI
@@ -238,13 +349,10 @@ WELCOME = """
 # Corrosion Assistant — Beta
 
 **Welcome!**  
-This model is trained for educational purpose only. At this time is locally trained with c.a. 9000 Images, some of them AI generated to give better instruction to the Model 
-and our proprietary AI Analyzer. More data/Pictures are needed to increase accuracy. At this time Crevice corrosion and Galvanic corrosion are still not well
-recognized so use this Model at your own risk. 
-**Disclaimer**: research & experimental only. Made with love by JQ.
+This model is trained for educational purpose only. Some classes still weak (crevice, galvanic).
+**Disclaimer**: research & experimental only. Validate with official PPG specs.
 """
 
-# Overlay HTML/CSS per stato di caricamento/analisi
 LOADER_HTML = """
 <div id="overlay-mask" style="
   position: fixed; inset: 0; background: rgba(0,0,0,0.55);
@@ -253,22 +361,19 @@ LOADER_HTML = """
 ">
   <div style="background:#111; color:#fff; padding:24px 28px; border-radius:16px;
               font-family: ui-sans-serif, system-ui, -apple-system; text-align:center;
-              box-shadow: 0 10px 30px rgba(0,0,0,0.5);">
+              box-shadow: 0 10px 30px rgba(0,0,0,0.5); max-width: 360px;">
     <div class="spinner" style="
       width:48px;height:48px;border:4px solid #444;border-top-color:#fff;border-radius:50%;
-      margin:0 auto 14px; animation: spin 1s linear infinite;"></div>
-    <div style="font-size:16px; font-weight:600;">Elaborazione in corso…</div>
-    <div style="opacity:0.85; font-size:12px; margin-top:6px;">Non toccare nulla, grazie.</div>
+      margin:0 auto 14px; animation: spin 0.9s linear infinite;"></div>
+    <div style="font-size:16px; font-weight:700;">Elaborazione in corso…</div>
+    <div style="opacity:0.9; font-size:12px; margin-top:6px;">Potrebbe richiedere alcuni secondi.</div>
   </div>
 </div>
-<style>
-@keyframes spin { to { transform: rotate(360deg); } }
-</style>
+<style>@keyframes spin { to { transform: rotate(360deg); } }</style>
 """
 
-# piccoli helper per mostrare/nascondere overlay e bloccare/sbloccare bottone
 def _show_overlay_and_busy():
-    return gr.update(visible=True), gr.update(interactive=False, value="🔄 Analisi in corso…")
+    return gr.update(visible=True), gr.update(interactive=False, value="🔄 Analyzing…")
 
 def _hide_overlay_and_idle():
     return gr.update(visible=False), gr.update(interactive=True, value="Analyze image")
@@ -276,15 +381,11 @@ def _hide_overlay_and_idle():
 with gr.Blocks(title="Corrosion Assistant", theme=gr.themes.Soft()) as demo:
     gr.Markdown(WELCOME)
 
-    # overlay nascosto di default
     overlay = gr.HTML(LOADER_HTML, visible=False)
-    # NB: l’overlay sta in cima all’app grazie a position:fixed
 
     with gr.Row():
         with gr.Column(scale=2):
-            # Lato “upload”: non posso mostrare percentuale reale dell’upload,
-            # ma il bottone mostrerà spinner e l’overlay partirà subito al click.
-            img = gr.Image(type="pil", sources=["upload","webcam"], label="Upload or webcam")
+            img = gr.Image(type="pil", sources=["upload", "webcam"], label="Upload or webcam")
             note = gr.Textbox(label="Notes / Context (optional)")
             zone = gr.Dropdown(choices=ZONES, label="Zone (indicative)", value="Other / Not sure")
             analyze_btn = gr.Button("Analyze image", variant="primary")
@@ -301,14 +402,14 @@ with gr.Blocks(title="Corrosion Assistant", theme=gr.themes.Soft()) as demo:
             clear_btn = gr.Button("Clear chat")
         with gr.Column(scale=2):
             gr.Markdown(
-                "> **Disclaimer:** Research & experimental use only. Validate with official PPG specs. "
-                "No professional advice or responsibility assumed."
+                "> **Privacy note:** If enabled, the image is sent to OpenAI to allow visual analysis. "
+                "Disable API key to skip assistant."
             )
 
     chat_state = gr.State([])
     thread_state = gr.State({"thread_id": None, "label": None, "confidence": 0.0, "zone": ""})
 
-    # catena: mostra overlay + disabilita bottone -> run -> nascondi overlay + riabilita -> aggiorna chat
+    # Catena robusta: overlay ON -> run -> overlay OFF -> sync chat
     analyze_btn.click(
         fn=_show_overlay_and_busy,
         inputs=[],
@@ -318,7 +419,7 @@ with gr.Blocks(title="Corrosion Assistant", theme=gr.themes.Soft()) as demo:
         fn=run_analysis,
         inputs=[img, note, zone, chat_state, thread_state],
         outputs=[out_md, chat_state, thread_state],
-        show_progress=True  # barra deterministica server-side
+        show_progress=True
     ).then(
         fn=_hide_overlay_and_idle,
         inputs=[],
@@ -348,5 +449,6 @@ with gr.Blocks(title="Corrosion Assistant", theme=gr.themes.Soft()) as demo:
     demo.api_mode = "enabled"
 
 if __name__ == "__main__":
+    # in Space, Gradio gestisce host/porta; local dev ok
     demo.launch()