Update app.py

app.py

@@ -1,23 +1,21 @@
-import os
 import json
 import numpy as np
 import gradio as gr
 from PIL import Image
 
-# Imaging (Chest X-ray) — TorchXRayVision
 import torch
 import torchxrayvision as xrv
 from torchvision import transforms
 
 # -----------------------------
-# Imaging Agent (Chest X-ray)
+# Imaging Agent (Chest X-ray, proxy for lung cancer risk)
 # -----------------------------
-_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-_MODEL = xrv.models.DenseNet(weights="densenet121-res224-all").to(_DEVICE)
-_MODEL.eval()
-_PATHOLOGIES = _MODEL.pathologies
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL = xrv.models.DenseNet(weights="densenet121-res224-all").to(DEVICE)
+MODEL.eval()
+PATHOLOGIES = MODEL.pathologies
 
-_resize = transforms.Compose([
+resize = transforms.Compose([
     xrv.datasets.XRayCenterCrop(),
     xrv.datasets.XRayResizer(224)
 ])
@@ -31,33 +29,30 @@ def imaging_agent(image_path: str):
         if arr.max() > 1:
             arr /= 255.0
         arr = xrv.datasets.normalize(arr, 4096)
-        arr = _resize(arr)
-        arr = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).to(_DEVICE)
+        arr = resize(arr)
+        arr = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).to(DEVICE)
 
         with torch.no_grad():
-            preds = _MODEL(arr)[0]
+            preds = MODEL(arr)[0]
             probs = torch.sigmoid(preds).cpu().numpy().tolist()
 
-        top = sorted(zip(_PATHOLOGIES, probs), key=lambda x: x[1], reverse=True)[:5]
-        lines = [f"{name}: {p*100:.1f}%" for name, p in top]
-        table = {name: round(p, 4) for name, p in zip(_PATHOLOGIES, probs)}
-        return "🖼️ Imaging Agent (Chest X-ray)\n" + "\n".join(lines), json.dumps(table, indent=2)
+        # Focus on lung pathologies relevant to cancer
+        focus_labels = ["Lung Opacity", "Mass", "Nodule"]
+        focus = [(l, probs[PATHOLOGIES.index(l)]) for l in focus_labels if l in PATHOLOGIES]
+
+        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)
     except Exception as e:
         return f"Imaging agent error: {e}", None
 
 # -----------------------------
-# Stubbed Signal Agent (ECG)
-# -----------------------------
-def signal_agent(file, sample_rate):
-    return "💓 Signal Agent (stub): ECG analysis not enabled in demo."
-
-# -----------------------------
-# Lab Agent (simple parsing)
+# Lab Agent (tumor markers, stub thresholds)
 # -----------------------------
-_THRESHOLDS = {
-    "glucose": {"unit": "mg/dL", "high": 126, "low": 70},
-    "hemoglobin": {"unit": "g/dL", "low": 12.0},
-    "spo2": {"unit": "%", "low": 92},
+CANCER_MARKERS = {
+    "psa": {"unit": "ng/mL", "high": 4},
+    "ca125": {"unit": "U/mL", "high": 35},
+    "afp": {"unit": "ng/mL", "high": 10},
 }
 
 import re
@@ -67,64 +62,57 @@ def lab_agent(text: str):
     results = []
     flags = []
     for line in text.splitlines():
-        m = re.findall(r'([a-z]+)\s*[:=]?\s*([\d\.]+)', line.lower())
+        m = re.findall(r'([a-z0-9]+)\s*[:=]?\s*([\d\.]+)', line.lower())
         for label, val in m:
-            if label in _THRESHOLDS:
+            if label in CANCER_MARKERS:
                 v = float(val)
-                thr = _THRESHOLDS[label]
+                thr = CANCER_MARKERS[label]
                 status = "ok"
-                if "high" in thr and v > thr["high"]:
-                    status = "high"; flags.append(f"{label} high")
-                if "low" in thr and v < thr["low"]:
-                    status = "low"; flags.append(f"{label} low")
-                results.append(f"{label.capitalize()}: {v} {thr['unit']} → {status}")
+                if v > thr["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 labs."
-    return "🧪 Lab Agent\n" + "\n".join(results) + ("\nFlags: " + ", ".join(flags) if flags else "\nFlags: none")
+        return "Could not parse tumor markers."
+    return "🧪 Lab Agent (Tumor Markers)\n" + "\n".join(results) + ("\nFlags: " + ", ".join(flags) if flags else "\nFlags: none")
 
 # -----------------------------
 # Coordinator
 # -----------------------------
-def coordinator(imaging_txt, signal_txt, lab_txt):
-    summary = "📋 Coordinator Summary\n"
+def coordinator(imaging_txt, lab_txt):
+    summary = "📋 Coordinator Summary (Early Cancer Screening)\n"
     if imaging_txt: summary += "\n" + imaging_txt
-    if signal_txt: summary += "\n" + signal_txt
     if lab_txt: summary += "\n" + lab_txt
-    summary += "\n\n⚠️ Disclaimer: Demo only. Not for clinical use."
+    summary += "\n\n⚠️ Disclaimer: Research demo only. Not for clinical use."
     return summary
 
 # -----------------------------
 # Gradio UI
 # -----------------------------
 with gr.Blocks(theme="soft") as demo:
-    gr.Markdown("# 🏥 AI Diagnostics Agents (Demo)")
-    gr.Markdown("Upload a chest X-ray or paste labs. ECG agent is stubbed for now.\n\n⚠️ Not for clinical use.")
+    gr.Markdown("# 🏥 AI Diagnostics Agent: Early Cancer Discovery (Demo)")
+    gr.Markdown("Upload a chest X-ray or paste tumor marker labs.\n\n⚠️ Research demo only. Not for clinical use.")
 
     with gr.Row():
         with gr.Column():
             img_in = gr.Image(type="filepath", label="Chest X-ray (PNG/JPG)")
-            imaging_out = gr.Textbox(label="Imaging Agent")
+            imaging_out = gr.Textbox(label="Imaging Agent Output")
             imaging_raw = gr.Code(label="Probabilities JSON", language="json")
         with gr.Column():
-            ecg_in = gr.File(label="ECG / Biosignal (stubbed)")
-            sr = gr.Number(label="Sampling Rate (Hz)", value=250)
-            signal_out = gr.Textbox(label="Signal Agent Output")
-        with gr.Column():
-            lab_in = gr.Textbox(lines=6, label="Lab Results (e.g., 'glucose: 180')")
+            lab_in = gr.Textbox(lines=6, label="Lab Results (e.g., 'PSA: 8 ng/mL')")
             lab_out = gr.Textbox(label="Lab Agent Output")
 
     run_btn = gr.Button("Run Agents")
     coord_out = gr.Textbox(label="Coordinator Summary", lines=10)
 
-    def run_all(image, ecg, rate, labs):
+    def run_all(image, labs):
         txt, raw = imaging_agent(image) if image else ("No image.", None)
-        sig = signal_agent(ecg, rate)
         lab = lab_agent(labs)
-        coord = coordinator(txt, sig, lab)
-        return txt, raw, sig, lab, coord
+        coord = coordinator(txt, lab)
+        return txt, raw, lab, coord
 
-    run_btn.click(run_all, inputs=[img_in, ecg_in, sr, lab_in], outputs=[imaging_out, imaging_raw, signal_out, lab_out, coord_out])
+    run_btn.click(run_all, inputs=[img_in, lab_in], outputs=[imaging_out, imaging_raw, lab_out, coord_out])
 
 demo.launch()
 
 
+