Create app.py

app.py

@@ -0,0 +1,723 @@
+"""
+🫁 Multi-Class Chest X-Ray Detection with Adaptive Sparse Training
+Advanced Gradio Interface - 4 Disease Classes
+
+Features:
+- Real-time detection: Normal, TB, Pneumonia, COVID-19
+- Grad-CAM visualization (explainable AI)
+- Improved specificity - distinguishes TB from pneumonia
+- Confidence scores with visual indicators
+- Clinical interpretation and recommendations
+- Mobile-responsive design
+"""
+
+import os
+from pathlib import Path
+import io
+
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+
+# ============================================================================
+# Device
+# ============================================================================
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"✅ Using device: {device}")
+
+# ============================================================================
+# Model Setup & Robust Loader
+# ============================================================================
+
+NUM_CLASSES = 4
+CLASSES = ["Normal", "Tuberculosis", "Pneumonia", "COVID-19"]
+CLASS_COLORS = {
+    "Normal": "#2ecc71",        # Green
+    "Tuberculosis": "#e74c3c",  # Red
+    "Pneumonia": "#f39c12",     # Orange
+    "COVID-19": "#9b59b6",      # Purple
+}
+
+
+def build_base_model(num_classes: int = NUM_CLASSES) -> nn.Module:
+    """
+    Build the base EfficientNet-B2 model with a 4-class classifier head.
+    This matches the architecture used during training.
+    """
+    # 🔒 Do NOT change this to efficientnet_b0 – your checkpoint is B2 (1408 features)
+    model = models.efficientnet_b2(weights=None)
+    in_features = model.classifier[1].in_features
+    model.classifier[1] = nn.Linear(in_features, num_classes)
+    return model
+
+
+def load_trained_model() -> nn.Module:
+    """
+    Load EfficientNet-B2 4-class checkpoint from:
+    - 'best.pt' OR
+    - 'checkpoints/best.pt'
+
+    Supports both:
+    - Plain state_dict
+    - Training checkpoint with 'model_state_dict' or 'state_dict' keys
+    """
+    model = build_base_model().to(device)
+
+    search_paths = [
+        Path("best.pt"),
+        Path("checkpoints/best.pt"),
+    ]
+
+    ckpt_path = None
+    for p in search_paths:
+        if p.exists():
+            ckpt_path = p
+            break
+
+    if ckpt_path is None:
+        raise RuntimeError(
+            "❌ Could not find model checkpoint.\n"
+            "Expected 'best.pt' in the project root OR 'checkpoints/best.pt'.\n"
+            "Please upload your 4-class EfficientNet-B2 weights as 'best.pt' or 'checkpoints/best.pt'."
+        )
+
+    print(f"🔍 Loading weights from: {ckpt_path}")
+
+    ckpt = torch.load(ckpt_path, map_location=device)
+
+    # Try to extract the actual state_dict
+    if isinstance(ckpt, dict):
+        if "model_state_dict" in ckpt:
+            state_dict = ckpt["model_state_dict"]
+        elif "state_dict" in ckpt:
+            state_dict = ckpt["state_dict"]
+        else:
+            # Assume it's already a plain state_dict
+            state_dict = ckpt
+    else:
+        # Definitely just a state_dict
+        state_dict = ckpt
+
+    # Now load strictly – if this fails, the checkpoint truly doesn't match the architecture
+    try:
+        missing, unexpected = model.load_state_dict(state_dict, strict=True)
+        if missing or unexpected:
+            # This branch rarely happens with strict=True, but keep for clarity
+            print(f"⚠️ Missing keys in state_dict: {missing}")
+            print(f"⚠️ Unexpected keys in state_dict: {unexpected}")
+    except RuntimeError as e:
+        # Most common cause: trying to load B2 checkpoint into B0/B1 or wrong architecture
+        raise RuntimeError(
+            f"❌ Failed to load weights from {ckpt_path}.\n"
+            "Most likely cause: the checkpoint was trained with a different architecture.\n"
+            "This app expects an EfficientNet-B2 checkpoint with 4 output classes.\n\n"
+            f"PyTorch error:\n{e}"
+        )
+
+    print("✅ Model weights loaded successfully!")
+    model.eval()
+    return model
+
+
+model = load_trained_model()
+
+# ============================================================================
+# Preprocessing
+# ============================================================================
+
+transform = transforms.Compose(
+    [
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            [0.485, 0.456, 0.406],  # ImageNet mean
+            [0.229, 0.224, 0.225],  # ImageNet std
+        ),
+    ]
+)
+
+# ============================================================================
+# Grad-CAM Implementation
+# ============================================================================
+
+
+class GradCAM:
+    def __init__(self, model: nn.Module, target_layer: nn.Module):
+        self.model = model
+        self.target_layer = target_layer
+        self.gradients = None
+        self.activations = None
+
+        def save_activation(module, input, output):
+            self.activations = output.detach()
+
+        def save_gradient(module, grad_input, grad_output):
+            # grad_output is a tuple; take the gradient wrt output activations
+            self.gradients = grad_output[0].detach()
+
+        target_layer.register_forward_hook(save_activation)
+        target_layer.register_full_backward_hook(save_gradient)
+
+    def generate(self, input_image: torch.Tensor, target_class=None):
+        """
+        Generate CAM for a single image batch (1, C, H, W).
+        """
+        output = self.model(input_image)
+
+        if target_class is None:
+            target_class = output.argmax(dim=1)
+
+        self.model.zero_grad()
+        one_hot = torch.zeros_like(output)
+        one_hot[0][target_class] = 1
+        output.backward(gradient=one_hot, retain_graph=True)
+
+        if self.gradients is None or self.activations is None:
+            return None, output
+
+        # Global average pooling over H, W
+        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
+        cam = (weights * self.activations).sum(dim=1, keepdim=True)
+        cam = torch.relu(cam)
+        cam = cam.squeeze().cpu().numpy()
+        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
+
+        return cam, output
+
+
+# Target the last feature block for Grad-CAM
+target_layer = model.features[-1]
+grad_cam = GradCAM(model, target_layer)
+
+# ============================================================================
+# Visualization Helpers
+# ============================================================================
+
+
+def create_original_display(image, pred_label, confidence):
+    """Create annotated original image"""
+    fig, ax = plt.subplots(figsize=(8, 8))
+    ax.imshow(image)
+    ax.axis("off")
+
+    color = CLASS_COLORS[pred_label]
+    title = f"Prediction: {pred_label}\nConfidence: {confidence:.1f}%"
+    ax.set_title(title, fontsize=16, fontweight="bold", color=color, pad=20)
+
+    plt.tight_layout()
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png", dpi=150, bbox_inches="tight", facecolor="white")
+    plt.close(fig)
+    buf.seek(0)
+    return Image.open(buf)
+
+
+def create_gradcam_visualization(image, cam, pred_label, confidence):
+    """Create Grad-CAM heatmap"""
+    img_array = np.array(image.resize((224, 224)))
+    cam_resized = cv2.resize(cam, (224, 224))
+
+    heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
+    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+
+    fig, ax = plt.subplots(figsize=(8, 8))
+    ax.imshow(heatmap)
+    ax.axis("off")
+    ax.set_title(
+        "Attention Heatmap\n(Areas the model focuses on)",
+        fontsize=14,
+        fontweight="bold",
+        pad=20,
+    )
+
+    plt.tight_layout()
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png", dpi=150, bbox_inches="tight", facecolor="white")
+    plt.close(fig)
+    buf.seek(0)
+    return Image.open(buf)
+
+
+def create_overlay_visualization(image, cam):
+    """Create overlay of image and heatmap"""
+    img_array = np.array(image.resize((224, 224))) / 255.0
+    cam_resized = cv2.resize(cam, (224, 224))
+
+    heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
+    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) / 255.0
+
+    overlay = img_array * 0.5 + heatmap * 0.5
+    overlay = np.clip(overlay, 0, 1)
+
+    fig, ax = plt.subplots(figsize=(8, 8))
+    ax.imshow(overlay)
+    ax.axis("off")
+    ax.set_title(
+        "Explainable AI Visualization\n(Original + Heatmap)",
+        fontsize=14,
+        fontweight="bold",
+        pad=20,
+    )
+
+    plt.tight_layout()
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png", dpi=150, bbox_inches="tight", facecolor="white")
+    plt.close(fig)
+    buf.seek(0)
+    return Image.open(buf)
+
+
+def create_interpretation(pred_label, confidence, results):
+    """Create interpretation text with medical-style narrative + disclaimers"""
+
+    interpretation = f"""
+## 🔬 Analysis Results
+
+### Prediction: **{pred_label}**
+- Confidence: **{confidence:.1f}%**
+
+### Probability Breakdown:
+- 🟢 Normal: **{results['Normal']:.1f}%**
+- 🔴 Tuberculosis: **{results['Tuberculosis']:.1f}%**
+- 🟠 Pneumonia: **{results['Pneumonia']:.1f}%**
+- 🟣 COVID-19: **{results['COVID-19']:.1f}%**
+
+---
+"""
+
+    # Disease-specific sections
+    if pred_label == "Tuberculosis":
+        if confidence >= 85:
+            interpretation += """
+**⚠️ High Confidence TB Detection**
+
+The model has detected features highly consistent with pulmonary tuberculosis.
+
+**CRITICAL – Suggested next clinical steps (for clinicians):**
+1. **Immediate clinical review** of the patient (history + physical exam)
+2. **Confirmatory tests**:
+   - Sputum smear microscopy and/or GeneXpert MTB/RIF
+   - TB culture where available
+3. **Correlate with symptoms**:
+   - Cough > 2 weeks
+   - Fever, night sweats
+   - Weight loss, hemoptysis
+4. **Consider isolation** and contact tracing if active TB is suspected
+5. **Additional imaging** (e.g., CT chest) if diagnosis remains uncertain
+
+> This tool is **screening-only** and cannot replace microbiological confirmation.
+"""
+        else:
+            interpretation += """
+**⚠️ Possible Tuberculosis**
+
+There are radiographic features that *may* be compatible with TB, but the model's confidence is moderate.
+
+**Recommended actions (for clinicians):**
+1. Perform focused clinical assessment
+2. Consider sputum testing (smear / GeneXpert)
+3. Review prior imaging for evolution of disease
+4. Use this result as a **second reader**, not definitive evidence
+
+> Moderate probability predictions always require clinical judgment.
+"""
+
+    elif pred_label == "Pneumonia":
+        if confidence >= 85:
+            interpretation += """
+**⚠️ High Confidence Pneumonia Detection**
+
+Findings are strongly suggestive of pneumonia (bacterial or viral).
+
+**Suggested steps:**
+1. Clinical evaluation for pneumonia severity
+2. Laboratory assessment:
+   - CBC, CRP/ESR
+   - Blood cultures if severely unwell
+3. Consider empiric antibiotics (if bacterial suspected) per local guidelines
+4. Repeat imaging if no improvement or worsening
+
+> Classic pneumonia patterns can overlap with other diseases – interpretation must remain clinical.
+"""
+        else:
+            interpretation += """
+**⚠️ Possible Pneumonia**
+
+The X-ray may show early or subtle changes of pneumonia.
+
+**Suggested steps:**
+1. Correlate with respiratory symptoms (cough, fever, dyspnea)
+2. Consider repeat imaging in 24–72 hours if clinically indicated
+3. Use this AI opinion as supportive, not definitive
+"""
+
+    elif pred_label == "COVID-19":
+        if confidence >= 85:
+            interpretation += """
+**⚠️ High Confidence COVID-19 Pattern**
+
+Pattern is compatible with COVID-19 pneumonia.
+
+**Suggested next steps:**
+1. **Confirmatory testing** with RT-PCR or validated antigen test
+2. **Infection control**:
+   - Isolation according to institutional policy
+   - Appropriate PPE for staff
+3. **Clinical monitoring**:
+   - Oxygen saturation (SpO₂)
+   - Respiratory rate, hemodynamics
+4. **Escalation** if:
+   - SpO₂ < 94%
+   - Increased work of breathing
+   - Hemodynamic instability
+
+> Radiology alone cannot confirm COVID-19 – virological testing is mandatory.
+"""
+        else:
+            interpretation += """
+**⚠️ Possible COVID-19**
+
+Some features overlap with COVID-19, but the model is not highly confident.
+
+**Suggested steps:**
+1. Test with RT-PCR or validated antigen assay
+2. Assess epidemiologic risk and exposure history
+3. Follow local protocols for isolation and monitoring
+"""
+
+    else:  # Normal
+        if confidence >= 85:
+            interpretation += """
+**✅ High Confidence “Normal” Chest X-Ray (for the 4 modeled diseases)**
+
+Within the limits of this model:
+- No strong evidence of **TB**, **pneumonia**, or **COVID-19** is detected.
+- Lung fields appear within normal limits on this projection.
+
+**Important caveats:**
+- A “normal” AI result does **not** exclude all lung disease.
+- Early or subtle TB/pneumonia/COVID-19 may still be radiographically occult.
+- Other conditions (PE, asthma, COPD, malignancy, etc.) are **outside the scope** of this model.
+
+Clinical review remains essential, especially if symptoms persist.
+"""
+        else:
+            interpretation += """
+**⚠️ Likely Normal, but with Lower Confidence**
+
+The model leans towards a normal study, but with limited confidence.
+
+**Suggested steps:**
+1. If the patient is symptomatic, clinical evaluation is still required.
+2. Consider repeat imaging if symptoms evolve.
+3. Use this output as an adjunct, not reassurance in isolation.
+"""
+
+    # Global disclaimer and technical note
+    interpretation += """
+---
+## ⚠️ CRITICAL MEDICAL DISCLAIMER
+
+### What this model *can* do:
+- ✅ Screen for 4 specific classes: **Normal**, **Tuberculosis**, **Pneumonia**, **COVID-19**
+- ✅ Provide **explainable heatmaps** (Grad-CAM) to highlight regions of interest
+- ✅ Offer **probabilistic support** to human readers
+- ✅ Leverage **Adaptive Sparse Training (AST)** for ~89% energy savings vs dense baselines
+
+### What this model *cannot* do:
+- ❌ It is **not** FDA/EMA-approved – research / educational use only
+- ❌ It does **not** replace radiologists, pulmonologists, or infectious disease specialists
+- ❌ It does **not** detect many other thoracic pathologies (e.g., cancer, fibrosis, PE)
+- ❌ It does **not** provide a microbiological diagnosis
+
+### Clinical usage guidance:
+1. Use as a **second reader** or screening tool.
+2. Always **correlate with clinical history, examination, and lab tests**.
+3. Never start, stop, or change treatment **solely** based on this AI prediction.
+4. Follow your local and international guidelines for TB, pneumonia, and COVID-19 management.
+
+### Diagnostic gold standards:
+- **TB**: Sputum AFB, culture, GeneXpert MTB/RIF, TB-PCR
+- **Pneumonia**: Clinical + imaging + microbiology
+- **COVID-19**: RT-PCR / validated antigen testing
+
+> When in doubt, a qualified healthcare professional’s judgment takes absolute precedence.
+
+---
+🫁 **Powered by Adaptive Sparse Training (AST)**
+Energy-efficient AI for accessible lung disease screening.
+
+**Project links:**
+- GitHub: https://github.com/oluwafemidiakhoa/Tuberculosis
+- Hugging Face Space: https://huggingface.co/spaces/mgbam/Tuberculosis
+"""
+
+    return interpretation
+
+
+# ============================================================================
+# Prediction Function
+# ============================================================================
+
+
+def predict_chest_xray(image, show_gradcam=True):
+    """
+    Main prediction function used by Gradio.
+    Returns:
+      - dict of class probabilities
+      - Annotated original image
+      - Grad-CAM heatmap
+      - Overlay image
+      - Markdown interpretation
+    """
+    if image is None:
+        return None, None, None, None, "Please upload a chest X-ray image."
+
+    # Ensure PIL RGB
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image).convert("RGB")
+    else:
+        image = image.convert("RGB")
+
+    original_img = image.copy()
+
+    input_tensor = transform(image).unsqueeze(0).to(device)
+
+    with torch.set_grad_enabled(show_gradcam):
+        if show_gradcam:
+            cam, output = grad_cam.generate(input_tensor)
+        else:
+            output = model(input_tensor)
+            cam = None
+
+    probs = torch.softmax(output, dim=1)[0].detach().cpu().numpy()
+    prob_sum = float(np.sum(probs))
+
+    if not (0.98 <= prob_sum <= 1.02):
+        print(f"⚠️ Probability sum = {prob_sum:.4f} (should be ~1.0). Check model/weights.")
+
+    pred_idx = int(output.argmax(dim=1).item())
+    pred_label = CLASSES[pred_idx]
+    confidence = float(probs[pred_idx]) * 100.0
+
+    results = {
+        CLASSES[i]: float(np.clip(probs[i] * 100.0, 0.0, 100.0))
+        for i in range(len(CLASSES))
+    }
+
+    original_pil = create_original_display(original_img, pred_label, confidence)
+
+    if cam is not None and show_gradcam:
+        gradcam_viz = create_gradcam_visualization(
+            original_img, cam, pred_label, confidence
+        )
+        overlay_viz = create_overlay_visualization(original_img, cam)
+    else:
+        gradcam_viz = None
+        overlay_viz = None
+
+    interpretation = create_interpretation(pred_label, confidence, results)
+
+    return results, original_pil, gradcam_viz, overlay_viz, interpretation
+
+
+# ============================================================================
+# Gradio Interface
+# ============================================================================
+
+custom_css = """
+#main-container {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    padding: 20px;
+}
+#title {
+    text-align: center;
+    color: white;
+    font-size: 2.5em;
+    font-weight: bold;
+    margin-bottom: 10px;
+    text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
+}
+#subtitle {
+    text-align: center;
+    color: #f0f0f0;
+    font-size: 1.2em;
+    margin-bottom: 20px;
+}
+#stats {
+    text-align: center;
+    color: #fff;
+    font-size: 0.95em;
+    margin-bottom: 30px;
+    padding: 15px;
+    background: rgba(255,255,255,0.1);
+    border-radius: 10px;
+    backdrop-filter: blur(10px);
+}
+.gradio-container {
+    font-family: 'Inter', sans-serif;
+}
+#upload-box {
+    border: 3px dashed #667eea;
+    border-radius: 15px;
+    padding: 20px;
+    background: rgba(255,255,255,0.95);
+}
+#results-box {
+    background: white;
+    border-radius: 15px;
+    padding: 20px;
+    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
+}
+.output-image {
+    border-radius: 10px;
+    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+}
+footer {
+    text-align: center;
+    margin-top: 30px;
+    color: white;
+    font-size: 0.9em;
+}
+"""
+
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    gr.HTML(
+        """
+        <div id="main-container">
+            <div id="title">🫁 Multi-Class Chest X-Ray Detection AI</div>
+            <div id="subtitle">Advanced chest X-ray analysis with Explainable AI</div>
+            <div id="stats">
+                <b>95–97% Accuracy</b> across 4 disease classes |
+                <b>89% Energy Efficient</b> |
+                Powered by Adaptive Sparse Training (AST)
+                <br><br>
+                <b>Detects:</b> Normal • Tuberculosis • Pneumonia • COVID-19
+            </div>
+        </div>
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1, elem_id="upload-box"):
+            gr.Markdown("## 📤 Upload Chest X-Ray")
+            image_input = gr.Image(
+                type="pil",
+                label="Upload X-Ray Image",
+                elem_classes="output-image",
+            )
+
+            show_gradcam = gr.Checkbox(
+                value=True,
+                label="Enable Grad-CAM Visualization (Explainable AI)",
+                info="Shows which areas the model focuses on",
+            )
+
+            analyze_btn = gr.Button("🔬 Analyze X-Ray", variant="primary", size="lg")
+
+            gr.Markdown(
+                """
+            ### 📋 Supported Images:
+            - Chest X-rays (PA or AP view)
+            - PNG, JPG, JPEG formats
+            - Grayscale or RGB
+
+            ### ⚡ Model Highlights:
+            - ✅ **Improved Specificity**: Better separation of TB vs Pneumonia
+            - ✅ **4 Disease Classes**: Normal, TB, Pneumonia, COVID-19
+            - ✅ **Energy-Aware**: ~89% energy savings with AST
+            - ✅ **Explainable**: Grad-CAM heatmaps for clinical teams
+            """
+            )
+
+        with gr.Column(scale=2, elem_id="results-box"):
+            gr.Markdown("## 📊 Analysis Results")
+
+            prob_output = gr.Label(
+                label="Prediction Confidence", num_top_classes=4
+            )
+
+            with gr.Tabs():
+                with gr.Tab("Original"):
+                    original_output = gr.Image(
+                        label="Annotated X-Ray", elem_classes="output-image"
+                    )
+
+                with gr.Tab("Grad-CAM Heatmap"):
+                    gradcam_output = gr.Image(
+                        label="Attention Heatmap", elem_classes="output-image"
+                    )
+
+                with gr.Tab("Overlay"):
+                    overlay_output = gr.Image(
+                        label="Explainable AI Visualization",
+                        elem_classes="output-image",
+                    )
+
+            interpretation_output = gr.Markdown(label="Clinical Interpretation")
+
+    gr.Markdown("## 📁 Example X-Rays (Demo Only)")
+    gr.Examples(
+        examples=[
+            ["examples/normal.png"],
+            ["examples/tb.png"],
+            ["examples/pneumonia.png"],
+            ["examples/covid.png"],
+        ],
+        inputs=image_input,
+        label="Click an example to load",
+    )
+
+    analyze_btn.click(
+        fn=predict_chest_xray,
+        inputs=[image_input, show_gradcam],
+        outputs=[
+            prob_output,
+            original_output,
+            gradcam_output,
+            overlay_output,
+            interpretation_output,
+        ],
+    )
+
+    gr.HTML(
+        """
+        <footer>
+            <p>
+                <b>🫁 Multi-Class Chest X-Ray Detection with AST</b><br>
+                Trained on Normal, Tuberculosis, Pneumonia, and COVID-19 cases<br>
+                95–97% Accuracy | 89% Energy Savings | Explainable AI<br><br>
+                <a href="https://github.com/oluwafemidiakhoa/Tuberculosis" target="_blank" style="color: white;">
+                    📂 GitHub Repository
+                </a> |
+                <a href="https://huggingface.co/spaces/mgbam/Tuberculosis" target="_blank" style="color: white;">
+                    🤗 Hugging Face Space
+                </a>
+            </p>
+            <p style="font-size: 0.8em; margin-top: 15px;">
+                ⚠️ <b>MEDICAL DISCLAIMER</b>: This is a screening / research tool, not a diagnostic device.
+                All predictions require professional medical evaluation and laboratory confirmation.
+                Not FDA-approved for clinical use.
+            </p>
+        </footer>
+        """
+    )
+
+# ============================================================================
+# Launch
+# ============================================================================
+
+if __name__ == "__main__":
+    demo.launch(
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+    )