Update app.py

app.py

@@ -16,152 +16,264 @@ import torchxrayvision as xrv
 import fitz  # PyMuPDF
 from torchcam.methods import SmoothGradCAMpp
 from transformers import pipeline
+import numpy as np
 
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger(__name__)
 
 # Load model
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-MODEL  = xrv.models.get_model("densenet121-res224-all").to(DEVICE).eval()
+MODEL = xrv.models.get_model("densenet121-res224-all").to(DEVICE).eval()
 LABELS = MODEL.pathologies
 
-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])
-])
-
-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)
+# Correct transform for TorchXRayVision (grayscale, normalized to [-1024, 1024])
+def preprocess_xray(pil_img: Image.Image) -> torch.Tensor:
+    """
+    Preprocess PIL image for TorchXRayVision model
+    TorchXRayVision expects:
+    - Single channel grayscale image
+    - Values normalized to [-1024, 1024] range
+    - Resolution of 224x224
+    """
+    # Convert to grayscale if needed
+    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)
+    # Assume input is 8-bit (0-255), scale to [-1024, 1024]
+    img_array = xrv.datasets.normalize(img_array, 255)
+    
+    # Add channel dimension and resize
+    img_array = img_array[None, ...]  # Add channel dimension
+    
+    # Use TorchXRayVision transforms
+    transform = transforms.Compose([
+        xrv.datasets.XRayCenterCrop(),
+        xrv.datasets.XRayResizer(224)
+    ])
+    
+    img_array = transform(img_array)
+    
+    # Convert to tensor
+    img_tensor = torch.from_numpy(img_array).unsqueeze(0).to(DEVICE)
+    
+    return img_tensor
 
-# X-ray prediction with Grad-CAM
-cam_extractor = SmoothGradCAMpp(MODEL)
+# Initialize CAM extractor with correct input shape for grayscale
+cam_extractor = SmoothGradCAMpp(MODEL, input_shape=(1, 224, 224))  # Single channel
 
 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
-    topk = torch.topk(probs, 3)
-
-    # Grad-CAM heat-map
-    target = topk.indices[0].item()
-    activation_map = cam_extractor(target, logits)[0]
-    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 border="1"><tr><th>Condition</th><th>Probability</th></tr>{rows}</table>
-    <p><b>Advice:</b> {advice}</p>
-    """
+    if img is None: 
+        return "Please upload an image.", None
 
-    return html, Image.fromarray(heatmap)
-
-# Medical advice
-ADVICE = {
-    "Pneumonia": "Consult a pulmonologist; antibiotics or antivirals as indicated.",
-    "Cardiomegaly": "Recommend echocardiography; refer to cardiology.",
-    "Atelectasis": "Further imaging may be needed; consult pulmonologist.",
-    "Consolidation": "Likely infection or inflammation; seek medical attention.",
-    "Pleural_Thickening": "Monitor for progression; pulmonology consultation.",
-    "Edema": "Evaluate for heart failure; cardiology consultation.",
-    "Effusion": "Thoracentesis may be needed; pulmonology consultation.",
-    "Fracture": "Orthopaedic consultation; consider CT if uncertain.",
-}
+    try:
+        # Preprocess image for TorchXRayVision
+        x = preprocess_xray(img)
+        
+        with torch.no_grad():
+            logits = MODEL(x)
+            probs = torch.sigmoid(logits)[0] * 100  # Convert to percentages
+        
+        # Get 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]
+        
+        # Overlay heatmap on original image
+        # Convert single channel to 3-channel for overlay
+        input_for_overlay = x.squeeze(0).cpu()
+        input_for_overlay = input_for_overlay.repeat(3, 1, 1)  # Repeat single channel 3 times
+        
+        heatmap = cam_extractor.overlay(input_for_overlay, activation_map)
+        
+        # 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)
+        
+    except Exception as e:
+        log.error(f"Error in X-ray analysis: {e}")
+        return f"Error processing image: {str(e)}", None
 
-def medical_advice(label): 
-    return ADVICE.get(label, "Discuss with a radiologist for next steps.")
+# 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."
+}
 
-# PDF report summariser
-try:
-    summariser = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
-except:
-    summariser = None
-    log.warning("Could not load summarization model")
+def get_medical_advice(condition: str) -> str:
+    return MEDICAL_ADVICE.get(condition, "Consult with a radiologist or pulmonologist for proper interpretation.")
 
+# PDF report analysis (simplified - focusing on the main issue)
 def analyse_report(file):
-    if file is None: return "Please upload a PDF."
+    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)
         text = "\n".join(page.get_text() for page in doc)
         doc.close()
         os.unlink(tmp_path)
-
-        disease = regex_find_disease(text)
-        if disease:
-            advice = medical_advice(disease)
-            return f"""
-            <h3>Disease detected:</h3><p>{disease}</p>
-            <p><b>Recommendation:</b> {advice}</p>
-            """
-        elif summariser:
-            # fallback LLM summary
-            short_text = text[:4000] if len(text) > 4000 else text
-            summary = summariser(short_text, max_length=120, min_length=30, do_sample=False)
-            return f"<h3>Report summary</h3><p>{summary[0]['summary_text']}</p>"
+        
+        # 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
         else:
-            return "<h3>Report processed</h3><p>No specific conditions detected. Please consult with a medical professional for interpretation.</p>"
+            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>"
+            
     except Exception as e:
+        log.error(f"Error in report analysis: {e}")
         return f"Error processing PDF: {str(e)}"
 
-def regex_find_disease(text: str):
-    patterns = {
-        "Pneumonia": r"\b(pneumonia|lung infection)\b",
-        "Cardiomegaly": r"\b(cardiomegaly|enlarged heart)\b",
-        "Atelectasis": r"\b(atelectasis|lung collapse)\b",
-        "Consolidation": r"\b(consolidation|lung consolidation)\b",
-        "Fracture": r"\b(fracture|broken bone|break)\b",
-        "Edema": r"\b(edema|fluid buildup)\b",
-        "Effusion": r"\b(effusion|fluid collection)\b",
-    }
-    for condition, pattern in patterns.items():
-        if re.search(pattern, text, flags=re.I): 
-            return condition
-    return None
-
-# Gradio UI
-with gr.Blocks(title="🩻 RadiologyScan AI") as demo:
-    gr.Markdown("## 🩻 RadiologyScan AI – Chest X-ray & Report Analyser")
+# 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.
+    """)
 
     with gr.Tabs():
-        with gr.Tab("X-ray Analysis"):
-            in_img = gr.Image(label="Upload chest X-ray", type="pil")
-            out_html = gr.HTML()
-            out_cam = gr.Image(label="Attention Map")
+        with gr.Tab("🔍 X-ray Analysis"):
+            gr.Markdown("Upload a chest X-ray image for AI analysis")
             
             with gr.Row():
-                analyze_btn = gr.Button("Analyze X-ray", variant="primary")
-                clear_btn = gr.Button("Clear")
+                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]
+            )
             
-            analyze_btn.click(analyse_xray, inputs=in_img, outputs=[out_html, out_cam])
-            clear_btn.click(lambda: (None, "", None), inputs=None, outputs=[in_img, out_html, out_cam])
+            clear_btn.click(
+                fn=lambda: (None, "", None),
+                inputs=None,
+                outputs=[img_input, html_output, cam_output]
+            )
 
-        with gr.Tab("Report Analysis"):
-            in_pdf = gr.File(label="Upload PDF report", file_types=[".pdf"])
-            out_rep = gr.HTML()
+        with gr.Tab("📄 Report Analysis"):
+            gr.Markdown("Upload a medical report PDF for AI analysis")
             
             with gr.Row():
-                analyze_rep_btn = gr.Button("Analyze Report", variant="primary")
-                clear_rep_btn = gr.Button("Clear")
+                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_rep_btn.click(analyse_report, inputs=in_pdf, outputs=out_rep)
-            clear_rep_btn.click(lambda: (None, ""), inputs=None, outputs=[in_pdf, out_rep])
+            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]
+            )
+
+    gr.Markdown("""
+    ### 📖 How to Use
+    1. **X-ray Analysis**: Upload a chest X-ray image (JPEG, PNG) and click "Analyze X-ray"
+    2. **Report Analysis**: Upload a medical report PDF and click "Analyze Report"
+    
+    ### 🔬 Technical Details
+    - Uses TorchXRayVision pre-trained DenseNet-121 model
+    - Trained on multiple chest X-ray datasets
+    - Provides attention heatmaps for interpretability
+    - Supports 18 different pathological conditions
+    
+    ### ⚠️ Limitations
+    - For educational use only
+    - Not a substitute for professional medical diagnosis
+    - Results may vary based on image quality
+    - Always consult healthcare professionals
+    """)
 
 if __name__ == "__main__":
-    demo.launch(show_error=True, server_port=int(os.getenv("PORT", 7860)))
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=int(os.getenv("PORT", 7860)),
+        show_error=True,
+        share=False
+    )