Update app.py

app.py

@@ -2,69 +2,84 @@ import json
 import numpy as np
 import gradio as gr
 from PIL import Image
-
 import torch
 import torchxrayvision as xrv
-import re
-import os
+from torchvision import transforms
 from skimage.transform import resize as sk_resize
+import matplotlib.pyplot as plt
+import io
+import re
 
 # -----------------------------
-# Imaging Agent (Chest X-ray, proxy for lung cancer risk)
+# Imaging Agent (Chest X-ray)
 # -----------------------------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 MODEL = xrv.models.DenseNet(weights="densenet121-res224-all").to(DEVICE)
 MODEL.eval()
 PATHOLOGIES = MODEL.pathologies
 
+
 def imaging_agent(image_path: str):
     if not image_path:
-        return "No image provided.", None
+        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)
-
-        # Normalize to [0,1]
         if arr.max() > 1:
             arr /= 255.0
-
-        # TorchXRayVision normalization
         arr = xrv.datasets.normalize(arr, 4096)
 
-        # --- Manual preprocessing instead of dict transforms ---
+        # Manual center crop & 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]  # center crop
-
-        arr = sk_resize(arr, (224, 224), preserve_range=True)   # resize to 224x224
+        arr = arr[starty:starty+min_dim, startx:startx+min_dim]
+        arr = sk_resize(arr, (224, 224), preserve_range=True)
 
-        # Convert to tensor
-        arr = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).to(DEVICE)
+        tensor_img = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).to(DEVICE)
 
         # Inference
         with torch.no_grad():
-            preds = MODEL(arr)[0]
+            preds = MODEL(tensor_img)[0]
             probs = torch.sigmoid(preds).cpu().numpy().tolist()
 
-        # Focus on lung pathologies relevant to cancer
+        # Focus on 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)
+        fmap = MODEL.features(tensor_img).detach().cpu().numpy()[0]
+        heatmap = np.mean(fmap, axis=0)
+        heatmap = sk_resize(heatmap, arr.shape, preserve_range=True)
+
+        plt.figure(figsize=(4, 4))
+        plt.imshow(arr, cmap="gray")
+        plt.imshow(heatmap, cmap="jet", alpha=0.4)
+        plt.axis("off")
+        buf = io.BytesIO()
+        plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+        plt.close()
+        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)]
         table = {name: round(p, 4) for name, p in focus}
 
-        return "🖼️ Imaging Agent (Chest X-ray for cancer risk)\n" + "\n".join(lines), json.dumps(table, indent=2)
+        return (
+            "🖼️ Imaging Agent (Chest X-ray for cancer risk)\n" + "\n".join(lines),
+            json.dumps(table, indent=2),
+            heatmap_img
+        )
 
     except Exception as e:
-        return f"Imaging agent error: {e}", None
+        return f"Imaging agent error: {e}", None, None
 
 
 # -----------------------------
-# Lab Agent (tumor markers, thresholds stub)
+# Lab Agent (tumor markers)
 # -----------------------------
 CANCER_MARKERS = {
     "psa": {"unit": "ng/mL", "high": 4},
@@ -72,6 +87,7 @@ CANCER_MARKERS = {
     "afp": {"unit": "ng/mL", "high": 10},
 }
 
+
 def lab_agent(text: str):
     if not text.strip():
         return "No lab text provided."
@@ -85,11 +101,14 @@ def lab_agent(text: str):
                 thr = CANCER_MARKERS[label]
                 status = "ok"
                 if v > thr["high"]:
-                    status = "elevated"; flags.append(f"{label.upper()} high")
+                    status = "elevated"
+                    flags.append(f"{label.upper()} high")
                 results.append(f"{label.upper()}: {v} {thr['unit']} → {status}")
     if not results:
         return "Could not parse tumor markers."
-    return "🧪 Lab Agent (Tumor Markers)\n" + "\n".join(results) + ("\nFlags: " + ", ".join(flags) if flags else "\nFlags: none")
+    return "🧪 Lab Agent (Tumor Markers)\n" + "\n".join(results) + (
+        "\nFlags: " + ", ".join(flags) if flags else "\nFlags: none"
+    )
 
 
 # -----------------------------
@@ -97,8 +116,10 @@ def lab_agent(text: str):
 # -----------------------------
 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
+    if imaging_txt:
+        summary += "\n" + imaging_txt
+    if lab_txt:
+        summary += "\n" + lab_txt
     summary += "\n\n⚠️ Disclaimer: Research demo only. Not for clinical use."
     return summary
 
@@ -110,17 +131,21 @@ SAMPLES = {
     "Normal X-ray": "samples/sample_xray1.png",
     "Suspicious X-ray": "samples/sample_xray2.png",
 }
-SAMPLE_LABS = "PSA: 8 ng/mL\nCA125: 20 U/mL\nAFP: 15 ng/mL"
+SAMPLE_TEXTS = {
+    "Lab Results": "samples/sample_labs.txt",
+    "MRI Report": "samples/sample_mri.txt",
+    "CT Report": "samples/sample_ct.txt",
+}
 
 
 # -----------------------------
 # Runner
 # -----------------------------
 def run_all(image, labs):
-    txt, raw = imaging_agent(image) if image else ("No image.", None)
+    txt, raw, heatmap = imaging_agent(image) if image else ("No image.", None, None)
     lab = lab_agent(labs)
     coord = coordinator(txt, lab)
-    return txt, raw, lab, coord
+    return txt, raw, heatmap, lab, coord
 
 
 # -----------------------------
@@ -140,20 +165,39 @@ with gr.Blocks(theme="soft") as demo:
             img_in = gr.Image(type="filepath", label="Chest X-ray (PNG/JPG)")
             imaging_out = gr.Textbox(label="Imaging Agent Output")
             imaging_raw = gr.Code(label="Probabilities JSON", language="json")
+            imaging_heatmap = gr.Image(label="Heatmap Overlay")
         with gr.Column():
-            lab_in = gr.Textbox(lines=6, label="Lab Results", value=SAMPLE_LABS)
+            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)
 
-    # Link dropdown to image input
+    # Link dropdowns
     def load_sample(choice):
         return SAMPLES.get(choice, None)
+
+    def load_text(choice):
+        path = SAMPLE_TEXTS.get(choice, None)
+        if path and path.endswith(".txt"):
+            with open(path, "r") as f:
+                return f.read()
+        return ""
+
     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, 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()
 
@@ -164,3 +208,4 @@ demo.launch()
 
 
 
+