Update app.py

app.py

@@ -1,262 +1,151 @@
-"""
-RadiologyScan AI – X-ray & Report analyser
-"""
-
 import os
-# Fix for PyTorch 2.6 weights_only issue
 os.environ['TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD'] = '1'
 
-import re, logging, tempfile
+import numpy as np
+import torch
 import gradio as gr
 from PIL import Image
-import torch
-import torch.nn.functional as F
-from torchvision import transforms
-import torchxrayvision as xrv
 import fitz  # PyMuPDF
+import torchxrayvision as xrv
+from torchvision import transforms
 from torchcam.methods import SmoothGradCAMpp
-from transformers import pipeline
-import numpy as np
+from torchcam.utils import overlay_mask
+import re
+import logging
 
 logging.basicConfig(level=logging.INFO)
-log = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
-# Load model - IMPORTANT: Don't call .eval() here for CAM to work
+# ---------------- MODEL SETUP ---------------- #
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 MODEL = xrv.models.get_model("densenet121-res224-all").to(DEVICE)
-# Note: We'll switch between train/eval modes as needed
 LABELS = MODEL.pathologies
 
-# Initialize CAM extractor with correct input shape for grayscale
+# Grad-CAM extractor (single-channel input)
 cam_extractor = SmoothGradCAMpp(MODEL, input_shape=(1, 224, 224))
 
-def preprocess_xray(pil_img: Image.Image) -> torch.Tensor:
-    """
-    Preprocess PIL image for TorchXRayVision model
-    """
-    # Convert to grayscale if needed
+# ---------------- IMAGE HANDLING ---------------- #
+def preprocess_image(pil_img: Image.Image):
+    """Convert to grayscale, normalize, resize for model"""
     if pil_img.mode != "L":
         pil_img = pil_img.convert("L")
-    
-    # Convert to numpy array
-    img_array = np.array(pil_img, dtype=np.float32)
-    
-    # Normalize to [-1024, 1024] range (TorchXRayVision standard)
+
+    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, ...]
-    
-    # Use TorchXRayVision transforms
+    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 with gradient enabled
-    img_tensor = torch.from_numpy(img_array).unsqueeze(0).to(DEVICE)
-    img_tensor.requires_grad_(True)  # Enable gradients for CAM
-    
-    return img_tensor
 
-def analyse_xray(img: Image.Image):
-    if img is None: 
-        return "Please upload an image.", None
+    # Convert to tensor
+    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.")
 
+def analyse_xray(img: Image.Image):
     try:
-        # Preprocess image
-        x = preprocess_xray(img)
-        
-        # Set model to train mode temporarily for gradient computation
-        MODEL.train()
-        
-        # Forward pass with gradient tracking
-        logits = MODEL(x)
-        probs = torch.sigmoid(logits)[0] * 100  # Convert to percentages
-        
-        # Get top 5 predictions
+        if img is None:
+            return "Please upload an X-ray image.", None
+
+        MODEL.train()  # required for CAM to calculate gradients
+        x = preprocess_image(img)
+        output = MODEL(x)
+        probs = torch.sigmoid(output)[0] * 100
+
+        # Top 5 predictions
         topk = torch.topk(probs, 5)
-        
-        # Generate Grad-CAM heat-map for the highest scoring condition
-        target = topk.indices[0].item()
-        
-        # Generate activation map
-        activation_map = cam_extractor(target, logits)[0]
-        
-        # Convert single channel to 3-channel for overlay
-        input_for_overlay = x.squeeze(0).cpu().detach()
-        input_for_overlay = input_for_overlay.repeat(3, 1, 1)
-        
-        # Generate heatmap
-        heatmap = cam_extractor.overlay(input_for_overlay, activation_map)
-        
-        # Set model back to eval mode
+
+        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)
+
         MODEL.eval()
-        
-        # Build HTML summary table
-        table_rows = ""
-        for i in range(len(topk.indices)):
-            idx = topk.indices[i].item()
-            prob = probs[idx].item()
-            condition = LABELS[idx]
-            table_rows += f"<tr><td>{condition}</td><td>{prob:.1f}%</td></tr>"
-        
-        top_condition = LABELS[target]
-        advice = get_medical_advice(top_condition)
-        
-        html_output = f"""
-        <div style="font-family: Arial, sans-serif;">
-            <h3>🩺 AI Analysis Results</h3>
-            <table border="1" style="border-collapse: collapse; width: 100%;">
-                <tr style="background-color: #f2f2f2;">
-                    <th style="padding: 8px; text-align: left;">Condition</th>
-                    <th style="padding: 8px; text-align: left;">Probability</th>
-                </tr>
-                {table_rows}
-            </table>
-            <br>
-            <h4>🔍 Top Finding: {top_condition}</h4>
-            <p><strong>Recommendation:</strong> {advice}</p>
-            <p><em>⚠️ This is an AI analysis tool for educational purposes only. Always consult qualified medical professionals for diagnosis and treatment.</em></p>
-        </div>
-        """
-        
-        return html_output, Image.fromarray(heatmap)
-        
+        return html, heat_img
+
     except Exception as e:
-        log.error(f"Error in X-ray analysis: {e}")
+        logger.error(e)
         return f"Error processing image: {str(e)}", None
 
-# Medical advice dictionary
-MEDICAL_ADVICE = {
-    "Atelectasis": "Lung collapse detected. Recommend pulmonology consultation and chest physiotherapy.",
-    "Cardiomegaly": "Enlarged heart detected. Recommend echocardiography and cardiology consultation.",
-    "Consolidation": "Lung consolidation detected. May indicate pneumonia or other lung disease. Seek medical attention.",
-    "Edema": "Pulmonary edema detected. Recommend urgent cardiology evaluation.",
-    "Emphysema": "Emphysema changes detected. Recommend pulmonology consultation and smoking cessation if applicable.",
-    "Fibrosis": "Lung fibrosis detected. Recommend pulmonology consultation for further evaluation.",
-    "Hernia": "Hernia detected. Recommend surgical consultation if symptomatic.",
-    "Infiltration": "Lung infiltration detected. May indicate infection or inflammation. Seek medical attention.",
-    "Mass": "Lung mass detected. Recommend urgent oncology consultation and further imaging.",
-    "Nodule": "Lung nodule detected. Recommend follow-up imaging and pulmonology consultation.",
-    "Pleural_Thickening": "Pleural thickening detected. Recommend pulmonology consultation.",
-    "Pneumonia": "Pneumonia detected. Recommend immediate antibiotic treatment and medical supervision.",
-    "Pneumothorax": "Pneumothorax (collapsed lung) detected. May require immediate medical intervention.",
-    "Effusion": "Pleural effusion detected. Recommend thoracentesis evaluation and pulmonology consultation."
-}
-
-def get_medical_advice(condition: str) -> str:
-    return MEDICAL_ADVICE.get(condition, "Consult with a radiologist or pulmonologist for proper interpretation.")
-
+# ---------------- PDF HANDLING ---------------- #
 def analyse_report(file):
-    if file is None: 
-        return "Please upload a PDF file."
-
     try:
-        # Create temporary file
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp:
-            tmp.write(file.read())
-            tmp_path = tmp.name
-        
-        # Extract text from PDF
-        doc = fitz.open(tmp_path)
+        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()
-        os.unlink(tmp_path)
-        
-        # Simple pattern matching for common conditions
-        detected_conditions = []
-        for condition in LABELS:
-            if re.search(rf'\b{condition.lower()}\b', text.lower()):
-                detected_conditions.append(condition)
-        
-        if detected_conditions:
-            html_output = "<h3>📋 Report Analysis</h3>"
-            html_output += "<h4>Detected Conditions:</h4><ul>"
-            for condition in detected_conditions[:5]:  # Show top 5
-                advice = get_medical_advice(condition)
-                html_output += f"<li><strong>{condition}</strong>: {advice}</li>"
-            html_output += "</ul>"
-            html_output += "<p><em>⚠️ This analysis is for educational purposes only. Consult medical professionals for proper diagnosis.</em></p>"
-            return html_output
+
+        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:
+                html += f"<li><b>{label}</b>: {get_medical_advice(label)}</li>"
+            html += "</ul>"
         else:
-            return "<h3>📋 Report Analysis</h3><p>No specific pathological conditions detected in the report text. Please consult with a medical professional for proper interpretation.</p>"
-            
+            html = "<p>No specific conditions found in the report.</p>"
+
+        return html
+
     except Exception as e:
-        log.error(f"Error in report analysis: {e}")
+        logger.error(e)
         return f"Error processing PDF: {str(e)}"
 
-# Gradio interface
-with gr.Blocks(title="🩻 RadiologyScan AI", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # 🩻 RadiologyScan AI
-    ### AI-Powered Chest X-ray and Medical Report Analysis
-    
-    **⚠️ IMPORTANT DISCLAIMER**: This tool is for educational and research purposes only. 
-    It should NOT be used for actual medical diagnosis or treatment decisions. 
-    Always consult qualified healthcare professionals for medical advice.
-    """)
+# ---------------- 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.Tabs():
         with gr.Tab("🔍 X-ray Analysis"):
-            gr.Markdown("Upload a chest X-ray image for AI analysis")
-            
-            with gr.Row():
-                with gr.Column():
-                    img_input = gr.Image(label="Upload Chest X-ray", type="pil")
-                    
-                    with gr.Row():
-                        analyze_btn = gr.Button("🔍 Analyze X-ray", variant="primary", size="lg")
-                        clear_btn = gr.Button("🗑️ Clear", variant="secondary")
-                
-                with gr.Column():
-                    html_output = gr.HTML(label="Analysis Results")
-                    cam_output = gr.Image(label="Attention Heatmap", type="pil")
-            
-            analyze_btn.click(
-                fn=analyse_xray,
-                inputs=img_input,
-                outputs=[html_output, cam_output]
-            )
-            
-            clear_btn.click(
-                fn=lambda: (None, "", None),
-                inputs=None,
-                outputs=[img_input, html_output, cam_output]
-            )
+            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"):
-            gr.Markdown("Upload a medical report PDF for AI analysis")
-            
-            with gr.Row():
-                with gr.Column():
-                    pdf_input = gr.File(label="Upload PDF Report", file_types=[".pdf"])
-                    
-                    with gr.Row():
-                        analyze_report_btn = gr.Button("📄 Analyze Report", variant="primary", size="lg")
-                        clear_report_btn = gr.Button("🗑️ Clear", variant="secondary")
-                
-                with gr.Column():
-                    report_output = gr.HTML(label="Report Analysis")
-            
-            analyze_report_btn.click(
-                fn=analyse_report,
-                inputs=pdf_input,
-                outputs=report_output
-            )
-            
-            clear_report_btn.click(
-                fn=lambda: (None, ""),
-                inputs=None,
-                outputs=[pdf_input, report_output]
-            )
+            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])
 
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=int(os.getenv("PORT", 7860)),
-        show_error=True,
-        share=False
-    )
+    demo.launch(server_port=int(os.getenv("PORT", 7860)), show_error=True)