Update app.py

app.py

@@ -1,8 +1,15 @@
+import os
+import io
 import gradio as gr
 from PIL import Image
 import torch
 import torchvision.transforms as T
 import torchvision.models as models
+from openai import OpenAI
+
+# ======================
+# Config / Model loading
+# ======================
 
 MODEL_PATH = "resnet50-corrosion-classifier-v1.pth"
 
@@ -39,7 +46,27 @@ transform = T.Compose([
                 std=[0.229, 0.224, 0.225]),
 ])
 
-def predict(image: Image.Image):
+# ======================
+# OpenAI client
+# ======================
+
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    # Gradio UI mostrerà l'errore se mancano le secrets
+    print("WARNING: OPENAI_API_KEY not set. Set it in HF Space Secrets.")
+
+client = OpenAI(api_key=OPENAI_API_KEY)
+
+# Puoi forzare la 2a lingua con env, altrimenti auto → in base al testo utente, fallback IT
+APP_FORCE_LANG = os.environ.get("APP_FORCE_LANG", "").strip()
+
+OPENAI_MODEL = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")
+
+# ======================
+# Inference functions
+# ======================
+
+def predict_image(image: Image.Image):
     img = image.convert("RGB")
     x = transform(img).unsqueeze(0).to(device)
     with torch.no_grad():
@@ -48,16 +75,93 @@ def predict(image: Image.Image):
     idx = int(probs.argmax())
     label = IDX2LABEL.get(idx, f"class_{idx}")
     confidence = float(probs[idx])
-    return {"label": label, "confidence": confidence}
+    return label, confidence
+
+def build_prompt(label: str, confidence: float, user_note: str, second_lang_hint: str):
+    # Istruzioni:
+    # 1) Prima parte sempre in EN
+    # 2) Seconda parte nella lingua dell'input utente (se presente),
+    #    altrimenti in IT, o forzata via APP_FORCE_LANG
+    second_lang_clause = ""
+    if APP_FORCE_LANG:
+        second_lang_clause = f"Then provide the same content in {APP_FORCE_LANG}."
+    else:
+        # Se l'utente ha scritto qualcosa, usa la sua lingua
+        if user_note and user_note.strip():
+            second_lang_clause = "Then repeat the same content in the same language used in the user note."
+        else:
+            second_lang_clause = "Then repeat the same content in Italian."
+
+    prompt = f"""
+You are a marine coatings technical assistant (PPG style). You have an image classification result:
+- Corrosion type: {label}
+- Model confidence: {round(confidence*100,2)}%
+
+Provide a technical advisory in two parts:
+
+1) English section (concise, structured as bullet points):
+- Diagnosis
+- Likely causes
+- Generic paint system suggestion (no unrealistic promises)
+- Warnings about substrate condition / prep
+
+2) {second_lang_clause}
+
+Context from user (optional): {user_note or "(none)"}
+
+Keep it pragmatic, accurate, and brief.
+"""
+    return prompt
+
+def ask_openai(label: str, confidence: float, user_note: str):
+    prompt = build_prompt(label, confidence, user_note, APP_FORCE_LANG)
+    try:
+        resp = client.chat.completions.create(
+            model=OPENAI_MODEL,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0.3
+        )
+        return resp.choices[0].message.content.strip()
+    except Exception as e:
+        return f"OpenAI error: {e}"
+
+def pipeline(image, note):
+    if image is None:
+        return "No image received."
+
+    label, conf = predict_image(image)
+    header = f"**Model result:** `{label}` — confidence **{round(conf*100,2)}%**\n\n"
+    advisory = ask_openai(label, conf, note or "")
+    return header + advisory
+
+# ======================
+# UI
+# ======================
+
+with gr.Blocks(title="Corrosion Assistant", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # Corrosion Assistant
+        Upload or take a photo. The model classifies corrosion type and an AI assistant gives a short technical advisory.
+        """
+    )
+    with gr.Row():
+        with gr.Column(scale=2):
+            img = gr.Image(type="pil", sources=["upload","webcam"], label="Upload or webcam")
+            note = gr.Textbox(label="Notes / Context (optional)", placeholder="Write in your language. EN comes first, then your language.")
+            btn = gr.Button("Analyze image", variant="primary")
+        with gr.Column(scale=3):
+            out = gr.Markdown(label="Analysis")
+    gr.Markdown(
+        """
+        <small>Tip: keep photos under ~2MB. Secrets are stored server-side.
+        UI is lightly themed; HF outer shell isn't customizable.</small>
+        """
+    )
 
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Upload o scatta"),
-    outputs=gr.JSON(label="Risultato"),
-    title="Corrosion Classifier API",
-    description="Restituisce label e confidence per l'immagine"
-)
+    btn.click(pipeline, inputs=[img, note], outputs=[out])
 
+# Enable programmatic API access to `pipeline`
 demo.api_mode = "enabled"
 
 if __name__ == "__main__":