Update app.py

app.py

@@ -8,112 +8,120 @@ from PIL import Image
 import fitz  # PyMuPDF
 import torchxrayvision as xrv
 from torchvision import transforms
-from torchcam.methods import SmoothGradCAMpp
-from torchcam.utils import overlay_mask
 import re
 
-# --- Model Setup ---
+# --- Device & Model ---
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-MODEL = xrv.models.get_model("densenet121-res224-all").to(DEVICE)
+MODEL = xrv.models.get_model("densenet121-res224-all").to(DEVICE).eval()
 LABELS = MODEL.pathologies
-cam_extractor = SmoothGradCAMpp(MODEL, input_shape=(1, 224, 224))
 
-# --- Preprocessing ---
-def preprocess_image(pil_img: Image.Image):
-    """Convert image to grayscale, normalize and resize for model"""
+# --- 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 = xrv.datasets.normalize(img_array, 255)  # normalize [-1024, 1024]
     img_array = img_array[None, ...]  # [1, H, W]
-    img_array = xrv.datasets.XRayCenterCrop()(img_array)
-    img_array = xrv.datasets.XRayResizer(224)(img_array)
+
+    # 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)
-    tensor.requires_grad_(True)
     return tensor
 
-def get_medical_advice(label):
-    advice_dict = {
-        "Pneumonia": "Consider antibiotics. Consult a pulmonologist.",
-        "Cardiomegaly": "Recommend echocardiogram and cardiologist review.",
-        "Effusion": "Pleural fluid detected. May require thoracentesis.",
-        "Fracture": "Possible bone injury. Orthopedic consultation needed.",
-        "Edema": "Fluid in lungs. Evaluate for heart failure.",
-    }
-    return advice_dict.get(label, "Please consult a radiologist for further evaluation.")
-
-# --- X-ray Analysis ---
+# --- Medical Recommendations ---
+ADVICE = {
+    "Pneumonia": "Possible infection. Recommend antibiotics and pulmonology consult.",
+    "Cardiomegaly": "Enlarged heart. Recommend echocardiography and cardiologist review.",
+    "Effusion": "Fluid in lung space. May need thoracentesis.",
+    "Fracture": "Possible bone break. Requires orthopedic consultation.",
+    "Edema": "Pulmonary fluid overload. Evaluate for heart failure.",
+}
+
+def get_advice(label):
+    return ADVICE.get(label, "Please consult a radiologist for further evaluation.")
+
+# --- X-ray Analysis (No CAM) ---
 def analyse_xray(img: Image.Image):
     try:
         if img is None:
             return "Please upload an X-ray image.", None
-        MODEL.train()  # Enable gradients for CAM
+
         x = preprocess_image(img)
-        output = MODEL(x)
-        probs = torch.sigmoid(output)[0] * 100
+        with torch.no_grad():
+            output = MODEL(x)
+            probs = torch.sigmoid(output)[0] * 100  # convert to percent
+
         topk = torch.topk(probs, 5)
-        html = "<h3>🩻 Top 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></tr>"
         for idx in topk.indices:
             html += f"<tr><td>{LABELS[idx]}</td><td>{probs[idx]:.1f}%</td></tr>"
         html += "</table><br>"
-        top_label = LABELS[topk.indices[0]]
-        advice = get_medical_advice(top_label)
-        html += f"<b>Suggested Action for '{top_label}':</b> {advice}"
-
-        # Grad-CAM overlay
-        cam = cam_extractor(topk.indices[0].item(), output)[0]  # 2D, (224,224)
-        img_rgb = img.convert("RGB").resize((224, 224))
-        cam_img = Image.fromarray((cam.cpu().numpy() * 255).astype(np.uint8))
-        heat_img = overlay_mask(img_rgb, cam_img, alpha=0.5)
-        MODEL.eval()
-        return html, heat_img
+
+        top_label = LABELS[topk.indices[0].item()]
+        html += f"<b>Recommended Action for '{top_label}':</b> {get_advice(top_label)}"
+
+        return html, img.resize((224, 224))  # return resized image for display
     except Exception as e:
         return f"Error processing image: {str(e)}", None
 
-# --- PDF Report Analysis ---
+# --- Report PDF Analysis ---
 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>📄 Detected Conditions</h3><ul>"
+            html = "<h3>📃 Findings Detected in Report:</h3><ul>"
             for label in found:
-                html += f"<li><b>{label}</b>: {get_medical_advice(label)}</li>"
+                html += f"<li><b>{label}</b>: {get_advice(label)}</li>"
             html += "</ul>"
         else:
-            html = "<p>No specific conditions found in the report.</p>"
+            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", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## 🩻 RadiologyScan AI\nUpload an X-ray or PDF report for AI-assisted analysis")
+with gr.Blocks(title="🩻 RadiologyScan AI") as demo:
+    gr.Markdown("## 🩻 RadiologyScan AI\nPerform fast AI-based analysis of Chest X-rays and medical reports")
 
     with gr.Tabs():
         with gr.Tab("🔍 X-ray Analysis"):
-            xray_input = gr.Image(label="Upload Chest X-ray", type="pil")
-            xray_html = gr.HTML()
-            xray_cam = gr.Image(label="AI Heatmap")
-            analyse_btn = gr.Button("Analyze X-ray")
-            clear_xray = gr.Button("Clear")
-            analyse_btn.click(analyse_xray, inputs=xray_input, outputs=[xray_html, xray_cam])
-            clear_xray.click(lambda: (None, "", None), None, outputs=[xray_input, xray_html, xray_cam])
-
-        with gr.Tab("📄 Report Analysis"):
-            pdf_input = gr.File(label="Upload PDF report", file_types=[".pdf"])
-            pdf_html = gr.HTML()
-            analyse_pdf_btn = gr.Button("Analyze Report")
-            clear_pdf = gr.Button("Clear")
-            analyse_pdf_btn.click(analyse_report, inputs=pdf_input, outputs=pdf_html)
-            clear_pdf.click(lambda: (None, ""), None, outputs=[pdf_input, pdf_html])
+            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])
+
+        with gr.Tab("📄 PDF Report Analysis"):
+            pdf_input = gr.File(file_types=[".pdf"], label="Upload PDF Medical Report")
+            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])
 
 if __name__ == "__main__":
     demo.launch(server_port=int(os.getenv("PORT", 7860)), show_error=True)