Update app.py

app.py

@@ -11,38 +11,23 @@ from torchvision import transforms
 from torchcam.methods import SmoothGradCAMpp
 from torchcam.utils import overlay_mask
 import re
-import logging
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# ---------------- MODEL SETUP ---------------- #
+# --- Model Setup ---
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 MODEL = xrv.models.get_model("densenet121-res224-all").to(DEVICE)
 LABELS = MODEL.pathologies
-
-# Grad-CAM extractor (single-channel input)
 cam_extractor = SmoothGradCAMpp(MODEL, input_shape=(1, 224, 224))
 
-# ---------------- IMAGE HANDLING ---------------- #
+# --- Preprocessing ---
 def preprocess_image(pil_img: Image.Image):
-    """Convert to grayscale, normalize, resize for model"""
+    """Convert image 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)
-
-    # Add channel dimension
-    img_array = img_array[None, ...]  # Shape: [1, H, W]
-
-    transform = transforms.Compose([
-        xrv.datasets.XRayCenterCrop(),
-        xrv.datasets.XRayResizer(224)
-    ])
-    img_array = transform(img_array)
-
-    # Convert to tensor
+    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)
     tensor.requires_grad_(True)
     return tensor
@@ -57,56 +42,46 @@ def get_medical_advice(label):
     }
     return advice_dict.get(label, "Please consult a radiologist for further evaluation.")
 
+# --- X-ray Analysis ---
 def analyse_xray(img: Image.Image):
     try:
         if img is None:
             return "Please upload an X-ray image.", None
-
-        MODEL.train()  # required for CAM to calculate gradients
+        MODEL.train()  # Enable gradients for CAM
         x = preprocess_image(img)
         output = MODEL(x)
         probs = torch.sigmoid(output)[0] * 100
-
-        # Top 5 predictions
         topk = torch.topk(probs, 5)
-
         html = "<h3>🩻 Top 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
-        cam = cam_extractor(topk.indices[0].item(), output)[0]
-        img_vis = img.convert("RGB").resize((224, 224))
-        heat_img = overlay_mask(img_vis, Image.fromarray((cam.cpu().numpy() * 255).astype(np.uint8)), alpha=0.5)
-
+        # 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
-
     except Exception as e:
-        logger.error(e)
         return f"Error processing image: {str(e)}", None
 
-# ---------------- PDF HANDLING ---------------- #
+# --- PDF Report Analysis ---
 def analyse_report(file):
     try:
         if file is None:
             return "Please upload a PDF report."
-
-        # Use file.name instead of .read()
         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>"
             for label in found:
@@ -114,14 +89,11 @@ def analyse_report(file):
             html += "</ul>"
         else:
             html = "<p>No specific conditions found in the report.</p>"
-
         return html
-
     except Exception as e:
-        logger.error(e)
         return f"Error processing PDF: {str(e)}"
 
-# ---------------- UI ---------------- #
+# --- 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")
 
@@ -130,20 +102,16 @@ with gr.Blocks(title="RadiologyScan AI", theme=gr.themes.Soft()) as demo:
             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])