Update app.py

app.py

@@ -1,9 +1,12 @@
 """
 RadiologyScan AI – X-ray & Report analyser
-Author : <you>
 """
 
-import os, re, logging, tempfile
+import os
+# Fix for PyTorch 2.6 weights_only issue
+os.environ['TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD'] = '1'
+
+import re, logging, tempfile
 import gradio as gr
 from PIL import Image
 import torch
@@ -59,7 +62,7 @@ def analyse_xray(img: Image.Image):
     advice = medical_advice(LABELS[target])
     html  = f"""
     <h3>AI findings</h3>
-    <table>{rows}</table>
+    <table border="1"><tr><th>Condition</th><th>Probability</th></tr>{rows}</table>
     <p><b>Advice:</b> {advice}</p>
     """
 
@@ -69,42 +72,67 @@ def analyse_xray(img: Image.Image):
 ADVICE = {
     "Pneumonia": "Consult a pulmonologist; antibiotics or antivirals as indicated.",
     "Cardiomegaly": "Recommend echocardiography; refer to cardiology.",
+    "Atelectasis": "Further imaging may be needed; consult pulmonologist.",
+    "Consolidation": "Likely infection or inflammation; seek medical attention.",
+    "Pleural_Thickening": "Monitor for progression; pulmonology consultation.",
+    "Edema": "Evaluate for heart failure; cardiology consultation.",
+    "Effusion": "Thoracentesis may be needed; pulmonology consultation.",
     "Fracture": "Orthopaedic consultation; consider CT if uncertain.",
 }
-def medical_advice(label): return ADVICE.get(label, "Discuss with a radiologist for next steps.")
+
+def medical_advice(label): 
+    return ADVICE.get(label, "Discuss with a radiologist for next steps.")
 
 # PDF report summariser
-summariser = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+try:
+    summariser = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+except:
+    summariser = None
+    log.warning("Could not load summarization model")
 
 def analyse_report(file):
     if file is None: return "Please upload a PDF."
 
-    with tempfile.NamedTemporaryFile(delete=False) as tmp:
-        tmp.write(file.read())
-        path = tmp.name
-    doc   = fitz.open(path)
-    text  = "\n".join(page.get_text() for page in doc)
-    doc.close()
-
-    disease = regex_find_disease(text)
-    if not disease:
-        short = summariser(text[:4000], max_length=120, min_length=30, do_sample=False)[0]["summary_text"]
-        return f"<h3>Report summary</h3><p>{short}</p>"
-
-    advice = medical_advice(disease)
-    return f"""
-    <h3>Disease detected:</h3><p>{disease}</p>
-    <p><b>Recommendation:</b> {advice}</p>
-    """
-
-def regex_find_disease(t:str):
+    try:
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp:
+            tmp.write(file.read())
+            tmp_path = tmp.name
+        
+        doc = fitz.open(tmp_path)
+        text = "\n".join(page.get_text() for page in doc)
+        doc.close()
+        os.unlink(tmp_path)
+
+        disease = regex_find_disease(text)
+        if disease:
+            advice = medical_advice(disease)
+            return f"""
+            <h3>Disease detected:</h3><p>{disease}</p>
+            <p><b>Recommendation:</b> {advice}</p>
+            """
+        elif summariser:
+            # fallback LLM summary
+            short_text = text[:4000] if len(text) > 4000 else text
+            summary = summariser(short_text, max_length=120, min_length=30, do_sample=False)
+            return f"<h3>Report summary</h3><p>{summary[0]['summary_text']}</p>"
+        else:
+            return "<h3>Report processed</h3><p>No specific conditions detected. Please consult with a medical professional for interpretation.</p>"
+    except Exception as e:
+        return f"Error processing PDF: {str(e)}"
+
+def regex_find_disease(text: str):
     patterns = {
-        "Pneumonia"    : r"\bpneumonia\b",
-        "Cardiomegaly" : r"\bcardiomegaly\b",
-        "Fracture"     : r"\bfracture\b",
+        "Pneumonia": r"\b(pneumonia|lung infection)\b",
+        "Cardiomegaly": r"\b(cardiomegaly|enlarged heart)\b",
+        "Atelectasis": r"\b(atelectasis|lung collapse)\b",
+        "Consolidation": r"\b(consolidation|lung consolidation)\b",
+        "Fracture": r"\b(fracture|broken bone|break)\b",
+        "Edema": r"\b(edema|fluid buildup)\b",
+        "Effusion": r"\b(effusion|fluid collection)\b",
     }
-    for k,v in patterns.items():
-        if re.search(v, t, flags=re.I): return k
+    for condition, pattern in patterns.items():
+        if re.search(pattern, text, flags=re.I): 
+            return condition
     return None
 
 # Gradio UI
@@ -112,18 +140,28 @@ with gr.Blocks(title="🩻 RadiologyScan AI") as demo:
     gr.Markdown("## 🩻 RadiologyScan AI – Chest X-ray & Report Analyser")
 
     with gr.Tabs():
-        with gr.Tab("X-ray image"):
-            in_img  = gr.Image(label="Upload chest X-ray")
-            out_html= gr.HTML()
-            out_cam = gr.Image(label="Explainability map")
-            gr.Button("Analyse").click(analyse_xray, in_img, outputs=[out_html,out_cam])
-            gr.Button("Clear").click(lambda: (None,"",None), None, [in_img,out_html,out_cam])
-
-        with gr.Tab("Radiology report"):
-            in_pdf  = gr.File(label="Upload PDF report", file_types=[".pdf"])
+        with gr.Tab("X-ray Analysis"):
+            in_img = gr.Image(label="Upload chest X-ray", type="pil")
+            out_html = gr.HTML()
+            out_cam = gr.Image(label="Attention Map")
+            
+            with gr.Row():
+                analyze_btn = gr.Button("Analyze X-ray", variant="primary")
+                clear_btn = gr.Button("Clear")
+            
+            analyze_btn.click(analyse_xray, inputs=in_img, outputs=[out_html, out_cam])
+            clear_btn.click(lambda: (None, "", None), inputs=None, outputs=[in_img, out_html, out_cam])
+
+        with gr.Tab("Report Analysis"):
+            in_pdf = gr.File(label="Upload PDF report", file_types=[".pdf"])
             out_rep = gr.HTML()
-            gr.Button("Analyse").click(analyse_report, in_pdf, out_rep)
-            gr.Button("Clear").click(lambda: (None,""), None, [in_pdf,out_rep])
+            
+            with gr.Row():
+                analyze_rep_btn = gr.Button("Analyze Report", variant="primary")
+                clear_rep_btn = gr.Button("Clear")
+            
+            analyze_rep_btn.click(analyse_report, inputs=in_pdf, outputs=out_rep)
+            clear_rep_btn.click(lambda: (None, ""), inputs=None, outputs=[in_pdf, out_rep])
 
 if __name__ == "__main__":
-    demo.launch(show_error=True, server_port=int(os.getenv("PORT",7860)))
+    demo.launch(show_error=True, server_port=int(os.getenv("PORT", 7860)))