Update app.py

app.py

@@ -7,6 +7,7 @@ import torch
 import torchxrayvision as xrv
 from torchvision import transforms
 import re
+import os
 
 # -----------------------------
 # Imaging Agent (Chest X-ray, proxy for lung cancer risk)
@@ -30,7 +31,12 @@ def imaging_agent(image_path: str):
         if arr.max() > 1:
             arr /= 255.0
         arr = xrv.datasets.normalize(arr, 4096)
-        arr = resize(arr)
+
+        # TorchXRayVision expects dict
+        sample = {"img": arr}
+        sample = resize(sample)
+        arr = sample["img"]
+
         arr = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).to(DEVICE)
 
         with torch.no_grad():
@@ -86,24 +92,38 @@ def coordinator(imaging_txt, lab_txt):
     return summary
 
 # -----------------------------
-# Gradio UI
+# Demo samples
 # -----------------------------
-SAMPLE_XRAY = "samples/sample_xray.jpg"
+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"
 
+# -----------------------------
+# Runner
+# -----------------------------
 def run_all(image, labs):
     txt, raw = imaging_agent(image) if image else ("No image.", None)
     lab = lab_agent(labs)
     coord = coordinator(txt, lab)
     return txt, raw, lab, coord
 
+# -----------------------------
+# Gradio UI
+# -----------------------------
 with gr.Blocks(theme="soft") as demo:
     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.")
+    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.")
 
     with gr.Row():
         with gr.Column():
-            img_in = gr.Image(type="filepath", label="Chest X-ray (PNG/JPG)", value=SAMPLE_XRAY)
+            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.Code(label="Probabilities JSON", language="json")
         with gr.Column():
@@ -113,6 +133,12 @@ with gr.Blocks(theme="soft") as demo:
     run_btn = gr.Button("Run Agents")
     coord_out = gr.Textbox(label="Coordinator Summary", lines=10)
 
+    # Link dropdown to image input
+    def load_sample(choice):
+        return SAMPLES.get(choice, None)
+    sample_dropdown.change(load_sample, inputs=sample_dropdown, outputs=img_in)
+
+    # Main button
     run_btn.click(run_all, inputs=[img_in, lab_in], outputs=[imaging_out, imaging_raw, lab_out, coord_out])
 
 demo.launch()
@@ -120,3 +146,4 @@ demo.launch()
 
 
 
+