Update app.py

app.py

@@ -1,220 +1,144 @@
+"""
+RadiologyScan AI – X-ray & Report analyser
+Author : <you>
+▶ requirements.txt needs:
+    torch torchvision torchxrayvision==1.2.0
+    pillow gradio pymupdf torchcam==0.4.0
+    transformers>=4.40.0 accelerate
+"""
+
+import os, re, logging, tempfile
 import gradio as gr
 from PIL import Image
 import torch
-from torchvision import models, transforms
-import fitz  # PyMuPDF
-import logging
-import os
-import re  # For better pattern matching
+import torch.nn.functional as F
+from torchvision import transforms
+import torchxrayvision as xrv                    # CheXNet-style models
+import fitz                                       # PyMuPDF
+from torchcam.methods import SmoothGradCAMpp      # visual explainability
+from transformers import pipeline
 
-# Setup logging
 logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Condition list (add your specific diseases here)
-conditions = [
-    "Normal", "Pneumonia", "Cancer", "TB", "Other",
-    "Coronary Artery Disease", "Aortic Aneurysm", "Stroke", "Peripheral Artery Disease",
-    "Brain Tumor", "Alzheimer's Disease", "Multiple Sclerosis", "Epilepsy",
-    "COPD", "Lung Cancer", "Pulmonary Embolism",
-    "Fractures", "Arthritis", "Osteoporosis",
-    "Appendicitis", "Gallstones", "Kidney Stones", "Infections", "Abdominal Aortic Aneurysm", "Diverticulitis"
-]
-
-# Condition details for diagnosis (can be expanded with real data)
-condition_details = {
-    "Normal": {"description": "No abnormal signs detected.", "recommendation": "Routine check-ups recommended."},
-    "Pneumonia": {"description": "Lung inflammation detected, possibly infectious.", "recommendation": "Seek medical attention for treatment."},
-    "Cancer": {"description": "Suspicious masses suggest cancer; further imaging needed.", "recommendation": "Consult an oncologist."},
-    "TB": {"description": "Cavitary lesions or nodular opacities indicate tuberculosis.", "recommendation": "Immediate medical evaluation and anti-TB therapy required."},
-    "Other": {"description": "Unclear abnormality; further investigation needed.", "recommendation": "Consult a radiologist."},
-    "Coronary Artery Disease": {"description": "Narrowing of coronary arteries detected.", "recommendation": "Cardiology consultation required."},
-    "Aortic Aneurysm": {"description": "Abnormal enlargement of the aorta.", "recommendation": "Vascular surgery evaluation."},
-    "Stroke": {"description": "Signs of brain ischemia or hemorrhage.", "recommendation": "Urgent neurological evaluation."},
-    "Peripheral Artery Disease": {"description": "Reduced blood flow in peripheral arteries.", "recommendation": "Vascular specialist consultation."},
-    "Brain Tumor": {"description": "Abnormal mass in the brain detected.", "recommendation": "Consult a neurosurgeon."},
-    "Alzheimer's Disease": {"description": "Signs of neurodegenerative changes.", "recommendation": "Neurology consultation."},
-    "Multiple Sclerosis": {"description": "Demyelinating lesions in the CNS.", "recommendation": "Neurology consultation."},
-    "Epilepsy": {"description": "Signs of seizure activity.", "recommendation": "Neurology consultation."},
-    "COPD": {"description": "Lung damage from COPD observed.", "recommendation": "Pulmonary consultation."},
-    "Lung Cancer": {"description": "Malignant lung masses detected.", "recommendation": "Oncology consultation."},
-    "Pulmonary Embolism": {"description": "Blockage in pulmonary arteries.", "recommendation": "Urgent medical attention."},
-    "Fractures": {"description": "Bone break detected.", "recommendation": "Orthopedic evaluation."},
-    "Arthritis": {"description": "Joint inflammation detected.", "recommendation": "Rheumatology consultation."},
-    "Osteoporosis": {"description": "Reduced bone density observed.", "recommendation": "Bone health specialist consultation."},
-    "Appendicitis": {"description": "Inflammation of the appendix.", "recommendation": "Surgical evaluation."},
-    "Gallstones": {"description": "Stones in the gallbladder detected.", "recommendation": "Gastroenterology consultation."},
-    "Kidney Stones": {"description": "Stones in the kidneys detected.", "recommendation": "Urology consultation."},
-    "Infections": {"description": "Signs of infection observed.", "recommendation": "Infectious disease consultation."},
-    "Abdominal Aortic Aneurysm": {"description": "Enlargement of the abdominal aorta.", "recommendation": "Vascular surgery evaluation."},
-    "Diverticulitis": {"description": "Inflammation of diverticula in the colon.", "recommendation": "Gastroenterology consultation."}
+log = logging.getLogger(__name__)
+
+# ------------------------------------------------------------------
+# 1. Load model – 18-label denseNet trained on multiple X-ray sets
+# ------------------------------------------------------------------
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL  = xrv.models.get_model("densenet121-res224-all").to(DEVICE).eval()
+LABELS = MODEL.pathologies                       # 18 canonical labels
+
+TRANSFORM = transforms.Compose([
+    transforms.Resize(224),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+])
+
+# ------------------------ helper ----------------------------------
+def preprocess(pil_img: Image.Image) -> torch.Tensor:
+    if pil_img.mode != "RGB":
+        pil_img = pil_img.convert("RGB")
+    return TRANSFORM(pil_img).unsqueeze(0).to(DEVICE)
+
+# ------------------------------------------------------------------
+# 2. X-ray prediction with Grad-CAM + textual advice
+# ------------------------------------------------------------------
+cam_extractor = SmoothGradCAMpp(MODEL)
+
+def analyse_xray(img: Image.Image):
+    if img is None: return "Please upload an image."
+
+    x = preprocess(img)
+    with torch.no_grad():
+        logits = MODEL(x)
+        probs  = torch.sigmoid(logits)[0] * 100   # multi-label %
+    topk = torch.topk(probs, 3)                   # show best 3
+
+    # Grad-CAM heat-map for the highest score
+    target = topk.indices[0].item()
+    activation_map = cam_extractor(target, logits)[0]   # H×W
+    heatmap = cam_extractor.overlay(torch.squeeze(x).cpu(), activation_map)
+
+    # Build HTML summary
+    rows = "".join(
+        f"<tr><td>{LABELS[i]}</td><td>{probs[i]:.1f}%</td></tr>"
+        for i in topk.indices
+    )
+    advice = medical_advice(LABELS[target])
+    html  = f"""
+    <h3>AI findings</h3>
+    <table>{rows}</table>
+    <p><b>Advice:</b> {advice}</p>
+    """
+
+    return html, Image.fromarray(heatmap)
+
+# simple rule-based advice (extend or swap for knowledge-graph)
+ADVICE = {
+    "Pneumonia": "Consult a pulmonologist; antibiotics or antivirals as indicated.",
+    "Cardiomegaly": "Recommend echocardiography; refer to cardiology.",
+    "Fracture": "Orthopaedic consultation; consider CT if uncertain.",
 }
-
-# Load a pre-trained DenseNet121 model
-model = models.densenet121(pretrained=True)
-model.classifier = torch.nn.Linear(model.classifier.in_features, len(conditions))  # Adjust the classifier for our condition count
-model.eval()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-# Image preprocessing function
-def preprocess_image(image):
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),  # Resize to match the model input size
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # Standard ImageNet normalization
-    ])
-    return transform(image).unsqueeze(0).to(device)
-
-# X-ray prediction function with summary output
-def predict_xray(image):
-    try:
-        if image is None:
-            return "Please upload an image."
-
-        img_tensor = preprocess_image(image)
-        with torch.no_grad():
-            output = model(img_tensor)
-
-        # Get probabilities for each condition
-        probs = torch.nn.functional.softmax(output, dim=1)[0]
-        results = {conditions[i]: probs[i].item() * 100 for i in range(len(conditions))}
-
-        # Identify the condition with the highest probability
-        top_condition = max(results, key=results.get)
-        confidence = results[top_condition]
-
-        # Construct a summary based on prediction
-        if confidence < 50:
-            return f"""
-            <div style="font-family:Arial">
-                <h3>Summary</h3>
-                <p><b>Disease Identified:</b> Uncertain</p>
-                <p><b>Cause/Status:</b> The model is not confident enough to provide a clear diagnosis.</p>
-                <p><b>Treatment/Recommendation:</b> Please consult a radiologist or upload a better-quality image for better accuracy.</p>
-            </div>
-            """
-
-        # Fetch details for the identified condition
-        info = condition_details.get(top_condition, condition_details["Other"])
-        return f"""
-        <div style="font-family:Arial">
-            <h3>Summary</h3>
-            <p><b>Disease Identified:</b> {top_condition}</p>
-            <p><b>Cause/Status:</b> {info['description']}</p>
-            <p><b>Treatment/Recommendation:</b> {info['recommendation']}</p>
-        </div>
-        """
-
-    except Exception as e:
-        logger.error(f"Error in prediction: {e}")
-        return f"Error: {str(e)}"
-
-# Enhanced PDF analysis function to extract patient details and disease
-def analyze_report(file):
-    if not file or not file.name.endswith(".pdf"):
-        return "Please upload a valid PDF file."
-    
-    try:
-        doc = fitz.open(file.name)
-        text = "".join(page.get_text() for page in doc)  # Extract all text from the PDF
-        doc.close()
-
-        # Initialize variables for patient details and condition
-        patient_info = {
-            "Name": "Not found",
-            "Age": "Not found",
-            "Gender": "Not found",
-            "Medical Record Number": "Not found",
-            "Date of Exam": "Not found"
-        }
-        condition = "Unclear"
-        confidence = "Not specified"
-
-        # Extract patient details using regex patterns
-        name_match = re.search(r'Name:\s*([^\n]+)', text, re.IGNORECASE)
-        if name_match:
-            patient_info["Name"] = name_match.group(1).strip()
-
-        age_match = re.search(r'Age:\s*(\d+)', text, re.IGNORECASE)
-        if age_match:
-            patient_info["Age"] = age_match.group(1).strip()
-
-        gender_match = re.search(r'Gender:\s*(Male|Female|Other)', text, re.IGNORECASE)
-        if gender_match:
-            patient_info["Gender"] = gender_match.group(1).strip()
-
-        mrn_match = re.search(r'Medical Record Number:\s*([^\n]+)', text, re.IGNORECASE)
-        if mrn_match:
-            patient_info["Medical Record Number"] = mrn_match.group(1).strip()
-
-        date_match = re.search(r'Date of Exam:\s*([^\n]+)', text, re.IGNORECASE)
-        if date_match:
-            patient_info["Date of Exam"] = date_match.group(1).strip()
-
-        # Improved condition matching with regex
-        if re.search(r'\btuberculosis\b|\bTB\b', text, re.IGNORECASE):
-            condition = "TB"
-            # Look for percentage of lung involvement
-            percentage_match = re.search(r'lung involvement:\s*approximately\s*(\d+)%', text, re.IGNORECASE)
-            if percentage_match:
-                confidence = f"{percentage_match.group(1)}% lung involvement"
-        elif re.search(r'\bstroke\b', text, re.IGNORECASE):
-            condition = "Stroke"
-        elif re.search(r'\bcancer\b', text, re.IGNORECASE):
-            condition = "Cancer"
-        elif re.search(r'\bfracture\b', text, re.IGNORECASE):
-            condition = "Fractures"
-        elif re.search(r'\bpneumonia\b', text, re.IGNORECASE):
-            condition = "Pneumonia"
-        # Add more conditions as needed from the conditions list
-
-        # Fetch condition details
-        info = condition_details.get(condition, condition_details["Other"])
-
-        # Construct summary output
-        summary = f"""
-        <div style="font-family:Arial">
-            <h3>Summary</h3>
-            <p><b>Disease Identified:</b> {condition}</p>
-            <p><b>Cause/Status:</b> {info['description']} {f'({confidence})' if confidence != 'Not specified' else ''}</p>
-            <p><b>Treatment/Recommendation:</b> {info['recommendation']}</p>
-            <h4>Patient Details</h4>
-            <p><b>Name:</b> {patient_info['Name']}</p>
-            <p><b>Age:</b> {patient_info['Age']}</p>
-            <p><b>Gender:</b> {patient_info['Gender']}</p>
-            <p><b>Medical Record Number:</b> {patient_info['Medical Record Number']}</p>
-            <p><b>Date of Exam:</b> {patient_info['Date of Exam']}</p>
-        </div>
-        """
-        return summary
-
-    except Exception as e:
-        logger.error(f"Failed to process PDF: {e}")
-        return f"Failed to process PDF: {str(e)}"
-
-# Gradio interface
-def create_interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("<h1 style='text-align:center;'>🩻 RadiologyScan AI</h1><p style='text-align:center;'>AI-powered X-ray and Report Analysis</p>")
-
-        with gr.Tabs():
-            with gr.TabItem("X-ray Analysis"):
-                img_input = gr.Image(label="Upload Chest X-ray", type="pil")
-                summary_output = gr.HTML(label="Summary Result")
-                with gr.Row():
-                    gr.Button("Analyze X-ray", elem_id="analyze_button", scale=0.3).click(predict_xray, inputs=img_input, outputs=summary_output)
-                    gr.Button("Clear", elem_id="clear_button", scale=0.3).click(lambda: [None, ""], inputs=None, outputs=[img_input, summary_output])
-
-            with gr.TabItem("Report Analysis"):
-                pdf_input = gr.File(label="Upload PDF Report", file_types=[". PDFs"])
-                summary_output_report = gr.HTML(label="Summary Result")
-                with gr.Row():
-                    gr.Button("Analyze Report", elem_id="analyze_button", scale=0.3).click(analyze_report, inputs=pdf_input, outputs=summary_output_report)
-                    gr.Button("Clear", elem_id="clear_button", scale=0.3).click(lambda: [None, ""], inputs=None, outputs=[pdf_input, summary_output_report])
-
-        return demo
+def medical_advice(label): return ADVICE.get(label, "Discuss with a radiologist for next steps.")
+
+# ------------------------------------------------------------------
+# 3. PDF report summariser (LLM pipeline fallback)
+# ------------------------------------------------------------------
+# Regex first → else call an LLM summariser (small DistilBART)
+summariser = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+
+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:
+        # fallback LLM summary
+        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):
+    patterns = {
+        "Pneumonia"    : r"\bpneumonia\b",
+        "Cardiomegaly" : r"\bcardiomegaly\b",
+        "Fracture"     : r"\bfracture\b",
+    }
+    for k,v in patterns.items():
+        if re.search(v, t, flags=re.I): return k
+    return None
+
+# ------------------------------------------------------------------
+# 4. Gradio UI
+# ------------------------------------------------------------------
+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"])
+            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])
 
 if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch(server_port=7860, ssr_mode=False)
+    demo.launch(show_error=True, server_port=int(os.getenv("PORT",7860)))