Update app.py

app.py

@@ -15,27 +15,82 @@ 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
 
-# --- Image Preprocessing ---
-def preprocess_image(pil_img: Image.Image) -> torch.Tensor:
-    """Convert to grayscale, normalize, and resize for model."""
-    if pil_img.mode != "L":
-        pil_img = pil_img.convert("L")
-
-    img_array = np.array(pil_img).astype(np.float32)
-    img_array = xrv.datasets.normalize(img_array, 255)  # normalize [-1024, 1024]
-    img_array = img_array[None, ...]  # [1, H, W]
-
-    # Apply center crop and resize
-    transform = transforms.Compose([
-        xrv.datasets.XRayCenterCrop(),
-        xrv.datasets.XRayResizer(224)
-    ])
-    img_array = transform(img_array)
-
-    tensor = torch.from_numpy(img_array).unsqueeze(0).to(DEVICE)
-    return tensor
+# --- Extended Medical Information ---
+DISEASE_INFO = {
+    "Atelectasis": {
+        "description":"Collapse of part or all of a lung, reducing oxygen exchange.",
+        "cause":"Blocked airway, lung compression, post-surgery.",
+        "recommendation":"Deep breathing exercises, possibly bronchoscopy or physiotherapy.",
+    },
+    "Cardiomegaly": {
+        "description":"Enlargement of the heart, seen as a broad silhouette.",
+        "cause":"High blood pressure, valve disease, cardiomyopathy.",
+        "recommendation":"Echocardiogram, consult a cardiologist.",
+    },
+    "Consolidation": {
+        "description":"Lung region filled with liquid instead of air.",
+        "cause":"Often from pneumonia (bacterial or viral).",
+        "recommendation":"Consult a physician; likely antibiotics and follow-up chest X-ray.",
+    },
+    "Edema": {
+        "description":"Fluid accumulation in the lungs.",
+        "cause":"Heart failure, kidney issues.",
+        "recommendation":"Treat underlying cause, may need diuretics, consult cardiology.",
+    },
+    "Effusion": {
+        "description":"Fluid buildup between lung and chest wall.",
+        "cause":"Infection, heart failure, cancer.",
+        "recommendation":"May need drainage (thoracentesis), see a pulmonologist.",
+    },
+    "Emphysema": {
+        "description":"Damage and enlargement of lung air sacs (alveoli).",
+        "cause":"Mainly smoking.",
+        "recommendation":"Quit smoking, pulmonary rehab, inhalers.",
+    },
+    "Fibrosis": {
+        "description":"Scarring of lung tissue, making breathing difficult.",
+        "cause":"Longstanding inflammation, auto-immune disease, occupational exposure.",
+        "recommendation":"Pulmonologist consult, immunosuppression/antifibrotic therapy.",
+    },
+    "Fracture": {
+        "description":"Break/crack in a bone (commonly ribs).",
+        "cause":"Trauma, accident, fall.",
+        "recommendation":"Pain management, monitor for organ injury, orthopedic consult if severe.",
+    },
+    "Infiltration": {
+        "description":"Something abnormal (cells/fluid) in the lungs.",
+        "cause":"Most often infection, sometimes inflammation or cancer.",
+        "recommendation":"See physician, further tests to clarify cause.",
+    },
+    "Mass": {
+        "description":"Lump or growth seen on X-ray.",
+        "cause":"Could be benign or malignant tumor.",
+        "recommendation":"Consult pulmonologist or oncologist, consider CT/biopsy.",
+    },
+    "Nodule": {
+        "description":"Small round or oval spot in the lung.",
+        "cause":"Old infection, benign, or early cancer.",
+        "recommendation":"May need CT and follow-up scans, discuss with doctor.",
+    },
+    "Pleural_Thickening": {
+        "description":"Thickening of chest lining.",
+        "cause":"Old infection, asbestos exposure.",
+        "recommendation":"Pulmonology follow-up; rarely needs intervention.",
+    },
+    "Pneumonia": {
+        "description":"Infection causing inflammation in the lungs.",
+        "cause":"Bacteria, viruses, or fungus.",
+        "recommendation":"Antibiotics/antivirals if needed. Seek prompt medical attention.",
+    },
+    "Pneumothorax": {
+        "description":"Collapsed lung (air leaks into chest cavity).",
+        "cause":"Trauma, rupture, sometimes spontaneous.",
+        "recommendation":"May need emergency care to remove air; consult ER.",
+    },
+    # Add more if you want to extend—see LABELS for all possible findings
+}
 
-# --- Medical Recommendations ---
+# --- Recommendations for top-line advice ---
 ADVICE = {
     "Pneumonia": "Possible infection. Recommend antibiotics and pulmonology consult.",
     "Cardiomegaly": "Enlarged heart. Recommend echocardiography and cardiologist review.",
@@ -47,27 +102,54 @@ ADVICE = {
 def get_advice(label):
     return ADVICE.get(label, "Please consult a radiologist for further evaluation.")
 
+def get_disease_info(label):
+    d = DISEASE_INFO.get(label)
+    if d:
+        return (
+            f"<b>{label}</b>: {d['description']}<br>"
+            f"<b>Possible Causes:</b> {d['cause']}<br>"
+            f"<b>Recommendation:</b> {d['recommendation']}"
+        )
+    return f"<b>{label}</b>: No extra info available. Please consult a radiologist."
+
+# --- Image Preprocessing ---
+def preprocess_image(pil_img: Image.Image) -> torch.Tensor:
+    """Convert to grayscale, normalize, and resize for model."""
+    if pil_img.mode != "L":
+        pil_img = pil_img.convert("L")
+    img_array = np.array(pil_img).astype(np.float32)
+    img_array = xrv.datasets.normalize(img_array, 255)
+    img_array = img_array[None, ...]  # [1, H, W]
+    img_array = xrv.datasets.XRayCenterCrop()(img_array)
+    img_array = xrv.datasets.XRayResizer(224)(img_array)
+    tensor = torch.from_numpy(img_array).unsqueeze(0).to(DEVICE)
+    return tensor
+
 # --- X-ray Analysis (No CAM) ---
 def analyse_xray(img: Image.Image):
     try:
         if img is None:
             return "Please upload an X-ray image.", None
-
         x = preprocess_image(img)
         with torch.no_grad():
             output = MODEL(x)
-            probs = torch.sigmoid(output)[0] * 100  # convert to percent
-
+            probs = torch.sigmoid(output)[0] * 100
         topk = torch.topk(probs, 5)
-        html = "<h3>🩺 Top 5 Predictions</h3><table border='1'><tr><th>Condition</th><th>Confidence</th></tr>"
+
+        html = "<h3>🩺 Top 5 Predictions</h3><table border='1'><tr><th>Condition</th><th>Confidence</th><th>Details</th></tr>"
         for idx in topk.indices:
-            html += f"<tr><td>{LABELS[idx]}</td><td>{probs[idx]:.1f}%</td></tr>"
-        html += "</table><br>"
+            label = LABELS[idx]
+            html += (
+                f"<tr><td>{label}</td>"
+                f"<td>{probs[idx]:.1f}%</td>"
+                f"<td>{get_disease_info(label)}</td></tr>"
+            )
+        html += "</table>"
 
         top_label = LABELS[topk.indices[0].item()]
-        html += f"<b>Recommended Action for '{top_label}':</b> {get_advice(top_label)}"
+        html += f"<br><b>Recommended Action for '{top_label}':</b> {get_advice(top_label)}"
 
-        return html, img.resize((224, 224))  # return resized image for display
+        return html, img.resize((224, 224))
     except Exception as e:
         return f"Error processing image: {str(e)}", None
 
@@ -76,41 +158,38 @@ def analyse_report(file):
     try:
         if file is None:
             return "Please upload a PDF report."
-
         doc = fitz.open(file.name)
         text = "\n".join(page.get_text() for page in doc)
         doc.close()
-
         found = []
         for label in LABELS:
             if re.search(rf"\b{label.lower()}\b", text.lower()):
                 found.append(label)
-
         if found:
             html = "<h3>📃 Findings Detected in Report:</h3><ul>"
             for label in found:
-                html += f"<li><b>{label}</b>: {get_advice(label)}</li>"
+                html += f"<li>{get_disease_info(label)}</li>"
             html += "</ul>"
         else:
             html = "<p>No known conditions detected from report text.</p>"
-
         return html
     except Exception as e:
         return f"Error processing PDF: {str(e)}"
 
 # --- Gradio UI ---
 with gr.Blocks(title="🩻 RadiologyScan AI") as demo:
-    gr.Markdown("## 🩻 RadiologyScan AI\nPerform fast AI-based analysis of Chest X-rays and medical reports")
-
+    gr.Markdown(
+        "## 🩻 RadiologyScan AI\n"
+        "Perform fast AI-based analysis of Chest X-rays and medical reports\n"
+        "<em>This tool provides informative summaries for common radiological findings. Not a substitute for a professional medical opinion.</em>"
+    )
     with gr.Tabs():
         with gr.Tab("🔍 X-ray Analysis"):
             x_input = gr.Image(label="Upload Chest X-ray", type="pil")
             x_out_html = gr.HTML()
             x_out_image = gr.Image(label="Resized X-ray (224x224)")
-
             analyze_btn = gr.Button("Analyze X-ray")
             clear_btn = gr.Button("Clear")
-
             analyze_btn.click(analyse_xray, inputs=x_input, outputs=[x_out_html, x_out_image])
             clear_btn.click(lambda: (None, "", None), None, [x_input, x_out_html, x_out_image])
 
@@ -119,7 +198,6 @@ with gr.Blocks(title="🩻 RadiologyScan AI") as demo:
             pdf_output = gr.HTML()
             analyze_pdf_btn = gr.Button("Analyze Report")
             clear_pdf_btn = gr.Button("Clear")
-
             analyze_pdf_btn.click(analyse_report, inputs=pdf_input, outputs=pdf_output)
             clear_pdf_btn.click(lambda: (None, ""), None, [pdf_input, pdf_output])