Update app.py

app.py

@@ -1,3 +1,8 @@
+import os
+import io
+import re
+import json
+import requests
 import numpy as np
 import gradio as gr
 from PIL import Image
@@ -6,9 +11,13 @@ import torchxrayvision as xrv
 from torchvision import transforms
 from skimage.transform import resize as sk_resize
 import matplotlib.pyplot as plt
-import io
-import re
-import json
+from dotenv import load_dotenv
+
+# -----------------------------
+# Environment setup
+# -----------------------------
+load_dotenv()
+MOONSHOT_API_KEY = os.getenv("MOONSHOT_API_KEY")
 
 # -----------------------------
 # Imaging Agent (Chest X-ray)
@@ -23,33 +32,32 @@ def imaging_agent(image_path: str):
     if not image_path:
         return "No image provided.", None, None
     try:
-        # Load grayscale X-ray
         img = Image.open(image_path).convert("L")
         arr = np.array(img).astype(np.float32)
         if arr.max() > 1:
             arr /= 255.0
         arr = xrv.datasets.normalize(arr, 4096)
 
-        # Manual center crop & resize
+        # Center crop and resize
         h, w = arr.shape
         min_dim = min(h, w)
         startx = w // 2 - (min_dim // 2)
         starty = h // 2 - (min_dim // 2)
-        arr = arr[starty:starty+min_dim, startx:startx+min_dim]
+        arr = arr[starty:starty + min_dim, startx:startx + min_dim]
         arr = sk_resize(arr, (224, 224), preserve_range=True)
 
         tensor_img = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).to(DEVICE)
 
-        # Inference
+        # Predictions
         with torch.no_grad():
             preds = MODEL(tensor_img)[0]
             probs = torch.sigmoid(preds).cpu().numpy().tolist()
 
-        # Focus on lung pathologies
+        # Focus on key lung pathologies
         focus_labels = ["Lung Opacity", "Mass", "Nodule"]
         focus = [(l, probs[PATHOLOGIES.index(l)]) for l in focus_labels if l in PATHOLOGIES]
 
-        # Generate heatmap overlay (using feature maps)
+        # Generate heatmap
         fmap = MODEL.features(tensor_img).detach().cpu().numpy()[0]
         heatmap = np.mean(fmap, axis=0)
         heatmap = sk_resize(heatmap, arr.shape, preserve_range=True)
@@ -64,10 +72,9 @@ def imaging_agent(image_path: str):
         buf.seek(0)
         heatmap_img = Image.open(buf)
 
-        # Format outputs
-        lines = [f"{name}: {p*100:.1f}%" for name, p in sorted(focus, key=lambda x: x[1], reverse=True)]
+        # Output probabilities
+        lines = [f"{name}: {p * 100:.1f}%" for name, p in sorted(focus, key=lambda x: x[1], reverse=True)]
         prob_text = "\n".join(lines)
-
         return (
             "🖼️ Imaging Agent (Chest X-ray for cancer risk)\n" + "\n".join(lines),
             prob_text,
@@ -112,14 +119,47 @@ def lab_agent(text: str):
 
 
 # -----------------------------
-# Coordinator
+# Moonshot.ai API Integration
+# -----------------------------
+def moonshot_summary(prompt: str):
+    try:
+        url = "https://api.moonshot.cn/v1/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {MOONSHOT_API_KEY}",
+            "Content-Type": "application/json"
+        }
+        payload = {
+            "model": "moonshot-v1",
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": 0.6,
+            "max_tokens": 600
+        }
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        data = response.json()
+        return data["choices"][0]["message"]["content"].strip()
+    except Exception as e:
+        return f"Moonshot API error: {e}"
+
+
+# -----------------------------
+# Coordinator (LLM-enhanced)
 # -----------------------------
 def coordinator(imaging_txt, lab_txt):
-    summary = "📋 Coordinator Summary (Early Cancer Screening)\n"
-    if imaging_txt:
-        summary += "\n" + imaging_txt
-    if lab_txt:
-        summary += "\n" + lab_txt
+    prompt = f"""
+    You are an AI medical coordinator. Based on the imaging and lab findings below,
+    generate a clear, patient-friendly summary assessing potential cancer risk,
+    highlighting abnormal findings, and suggesting appropriate next steps.
+    Avoid technical jargon. Keep it concise and empathetic.
+
+    Imaging findings:
+    {imaging_txt}
+
+    Lab results:
+    {lab_txt}
+    """
+    ai_summary = moonshot_summary(prompt)
+    summary = "📋 AI Coordinator Summary (LLM-generated)\n\n" + ai_summary
     summary += "\n\n⚠️ Disclaimer: Research demo only. Not for clinical use."
     return summary
 
@@ -153,32 +193,24 @@ def run_all(image, labs):
 # -----------------------------
 with gr.Blocks(theme="soft") as demo:
     gr.Markdown("# 🏥 AI Diagnostics Agent: Early Cancer Discovery (Demo)")
-    gr.Markdown("Upload a chest X-ray, or pick a demo sample. Paste tumor marker labs.\n\n⚠️ Research demo only. Not for clinical use.")
+    gr.Markdown("Upload a chest X-ray or choose a sample. Paste or load lab / MRI / CT reports. "
+                "\n\n⚠️ Research demo only. Not for clinical use.")
 
     with gr.Row():
         with gr.Column():
-            sample_dropdown = gr.Dropdown(
-                choices=list(SAMPLES.keys()),
-                value="Normal X-ray",
-                label="Select Sample X-ray"
-            )
+            sample_dropdown = gr.Dropdown(choices=list(SAMPLES.keys()), value="Normal X-ray", label="Select Sample X-ray")
             img_in = gr.Image(type="filepath", label="Chest X-ray (PNG/JPG)")
             imaging_out = gr.Textbox(label="Imaging Agent Output")
             imaging_raw = gr.Textbox(label="Probabilities (%)", lines=6)
             imaging_heatmap = gr.Image(label="Heatmap Overlay")
         with gr.Column():
-            text_dropdown = gr.Dropdown(
-                choices=list(SAMPLE_TEXTS.keys()),
-                value="Lab Results",
-                label="Select Sample Report"
-            )
+            text_dropdown = gr.Dropdown(choices=list(SAMPLE_TEXTS.keys()), value="Lab Results", label="Select Sample Report")
             lab_in = gr.Textbox(lines=6, label="Lab / Report Input")
             lab_out = gr.Textbox(label="Lab Agent Output")
 
     run_btn = gr.Button("Run Agents")
-    coord_out = gr.Textbox(label="Coordinator Summary", lines=10)
+    coord_out = gr.Textbox(label="AI Coordinator Summary", lines=10)
 
-    # Link dropdowns
     def load_sample(choice):
         return SAMPLES.get(choice, None)
 
@@ -191,13 +223,7 @@ with gr.Blocks(theme="soft") as demo:
 
     sample_dropdown.change(load_sample, inputs=sample_dropdown, outputs=img_in)
     text_dropdown.change(load_text, inputs=text_dropdown, outputs=lab_in)
-
-    # Main button
-    run_btn.click(
-        run_all,
-        inputs=[img_in, lab_in],
-        outputs=[imaging_out, imaging_raw, imaging_heatmap, lab_out, coord_out],
-    )
+    run_btn.click(run_all, inputs=[img_in, lab_in], outputs=[imaging_out, imaging_raw, imaging_heatmap, lab_out, coord_out])
 
 demo.launch()
 
@@ -210,3 +236,4 @@ demo.launch()
 
 
 
+