Update app.py

app.py

@@ -1,10 +1,7 @@
 """
-🫁 Multi-Class Chest X-Ray Detection with Adaptive Sparse Training
-WOW UI/UX Edition – 4 Disease Classes
-
-- Normal, Tuberculosis, Pneumonia, COVID-19
-- Grad-CAM (Explainable AI)
-- Energy-efficient Adaptive Sparse Training
+🫁 AST Chest X-Ray Lab
+Multi-Class Chest X-Ray Detection (Normal · TB · Pneumonia · COVID-19)
+with Adaptive Sparse Training & Explainable AI (Grad-CAM)
 """
 
 import io
@@ -12,6 +9,8 @@ from pathlib import Path
 
 import cv2
 import gradio as gr
+import matplotlib
+matplotlib.use("Agg")  # safe backend for servers
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
@@ -25,36 +24,68 @@ from torchvision import models, transforms
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# EfficientNet backbone
+# EfficientNet backbone – we want 4 output classes
+NUM_CLASSES = 4
 model = models.efficientnet_b0(weights=None)
-model.classifier[1] = nn.Linear(model.classifier[1].in_features, 4)  # 4 classes
+model.classifier[1] = nn.Linear(model.classifier[1].in_features, NUM_CLASSES)
 
-# Try a few reasonable checkpoint locations
+# We expect a 4-class EfficientNet checkpoint here
 checkpoint_candidates = [
-    "best.pt",
-    "checkpoints/best.pt",          # <-- your current file
-    "checkpoints/lasttb.pt",        # optional fallback
+    "checkpoints/best.pt",   # main location from your screenshot
+    "best.pt",               # optional fallback in repo root
 ]
 
 MODEL_LOAD_INFO = ""
 loaded = False
 
+
+def extract_state_dict(ckpt):
+    """
+    Handle both:
+    - plain state_dict (just parameter tensors)
+    - training checkpoints with keys like 'model_state_dict', 'state_dict', etc.
+    """
+    if isinstance(ckpt, dict):
+        for key in ["model_state_dict", "state_dict", "model"]:
+            if key in ckpt and isinstance(ckpt[key], dict):
+                return ckpt[key]
+    return ckpt  # already a raw state dict
+
+
 for ckpt_path in checkpoint_candidates:
     if Path(ckpt_path).is_file():
         try:
             print(f"🔍 Trying to load weights from: {ckpt_path}")
-            state_dict = torch.load(ckpt_path, map_location=device)
-            model.load_state_dict(state_dict)
-            MODEL_LOAD_INFO = f"✅ Model loaded from **{ckpt_path}** on **{device.type.upper()}**."
+            raw_ckpt = torch.load(ckpt_path, map_location=device)
+            state_dict = extract_state_dict(raw_ckpt)
+
+            # Check classifier size to ensure it's truly 4-class
+            if "classifier.1.weight" in state_dict:
+                out_features = state_dict["classifier.1.weight"].shape[0]
+                if out_features != NUM_CLASSES:
+                    raise ValueError(
+                        f"Checkpoint at {ckpt_path} has {out_features} output "
+                        f"classes, but this app expects {NUM_CLASSES}."
+                    )
+
+            # Load strict – we want the full EfficientNet weights
+            model.load_state_dict(state_dict, strict=True)
+
+            MODEL_LOAD_INFO = (
+                f"✅ Model loaded from **{ckpt_path}** on **{device.type.upper()}**."
+            )
             loaded = True
             break
         except Exception as e:
-            print(f"⚠️ Found {ckpt_path} but failed to load: {e}")
+            print(f"⚠️ Found {ckpt_path} but failed to load model_state_dict: {e}")
 
 if not loaded:
     raise RuntimeError(
-        "Model file not found or could not be loaded. "
-        "Please upload 'checkpoints/best.pt' (or 'best.pt' in the repo root)."
+        "Model file not found or could not be loaded.\n"
+        "Expected a 4-class EfficientNet checkpoint at 'checkpoints/best.pt' "
+        "or 'best.pt' that was saved with model.state_dict().\n"
+        "If you saved a training checkpoint, make sure it has a "
+        "'model_state_dict' key with the 4-class EfficientNet weights."
     )
 
 model = model.to(device)
@@ -69,10 +100,10 @@ TOTAL_PARAMS_M = TOTAL_PARAMS / 1e6
 
 CLASSES = ["Normal", "Tuberculosis", "Pneumonia", "COVID-19"]
 CLASS_COLORS = {
-    "Normal": "#2ecc71",        # Green
-    "Tuberculosis": "#e74c3c",  # Red
-    "Pneumonia": "#f39c12",     # Orange
-    "COVID-19": "#9b59b6",      # Purple
+    "Normal": "#22c55e",        # Green
+    "Tuberculosis": "#ef4444",  # Red
+    "Pneumonia": "#f97316",     # Orange
+    "COVID-19": "#a855f7",      # Purple
 }
 
 transform = transforms.Compose(
@@ -207,7 +238,6 @@ def create_overlay_visualization(image, cam):
     plt.tight_layout()
     return _figure_to_pil()
 
-
 # ============================================================================
 # Interpretation
 # ============================================================================
@@ -237,7 +267,6 @@ def create_interpretation(pred_label, confidence, results, audience="Clinician")
 ---
 """
 
-    # Disease-specific sections (same logic, slightly formatted)
     if pred_label == "Tuberculosis":
         if confidence >= 85:
             interpretation += """
@@ -245,18 +274,12 @@ def create_interpretation(pred_label, confidence, results, audience="Clinician")
 
 The model has detected features strongly suggestive of **pulmonary tuberculosis**.
 
-**Recommended Clinical Pathway:**
-1. ✅ Immediate medical review by a clinician or chest physician  
-2. ✅ **Sputum testing** (AFB smear, GeneXpert MTB/RIF, or TB-PCR)  
-3. ✅ Correlate with symptoms:
-   - Persistent cough > 2 weeks  
-   - Weight loss, night sweats  
-   - Fever, fatigue  
-   - Hemoptysis (coughing blood)  
-4. ✅ Consider CT scan or additional imaging if uncertainty remains  
-5. ✅ Infection control and contact tracing if TB is confirmed  
-
-> This tool helps *flag* suspicious cases. TB diagnosis still requires **laboratory confirmation**.
+**Suggested Clinical Pathway**
+1. Immediate review by a clinician or chest physician  
+2. Sputum testing (AFB smear, GeneXpert MTB/RIF, or TB-PCR)  
+3. Correlate with symptoms: chronic cough, weight loss, night sweats, fever, hemoptysis  
+4. Consider CT or further imaging if uncertainty remains  
+5. Infection control and contact tracing if TB is confirmed  
 """
         else:
             interpretation += """
@@ -264,10 +287,9 @@ The model has detected features strongly suggestive of **pulmonary tuberculosis*
 
 The scan shows features that **could** be compatible with tuberculosis, but confidence is moderate.
 
-**Suggested Actions:**
-- Clinical review and detailed history  
-- Consider sputum testing if symptoms or risk factors are present  
-- Follow-up imaging where clinically indicated  
+- Correlate with symptoms and risk factors  
+- Consider sputum testing where suspicion is non-trivial  
+- Follow-up imaging as clinically indicated  
 """
 
     elif pred_label == "Pneumonia":
@@ -277,15 +299,13 @@ The scan shows features that **could** be compatible with tuberculosis, but conf
 
 The model has detected an opacity pattern consistent with **pneumonia**.
 
-**Typical Clinical Correlates:**
+Typical clinical correlates:
+
 - Fever, productive cough  
 - Shortness of breath  
 - Pleuritic chest pain  
 
-**Next Steps (for clinicians):**
-- Correlate with fever, auscultation, and lab results  
-- Consider antibiotics for bacterial pneumonia as per local guidelines  
-- Repeat imaging if clinical evolution is atypical  
+Use alongside clinical exam, labs, and local treatment guidelines.
 """
         else:
             interpretation += """
@@ -293,10 +313,9 @@ The model has detected an opacity pattern consistent with **pneumonia**.
 
 Findings may be compatible with pneumonia, but alternative explanations exist.
 
-**Recommended:**
-- Clinical evaluation (vital signs, exam)  
-- Consider labs (WBC, CRP, cultures)  
-- Watchful follow-up or repeat imaging as appropriate  
+- Check vital signs and auscultation  
+- Labs (WBC, CRP, cultures) may be useful  
+- Consider watchful follow-up or repeat imaging  
 """
 
     elif pred_label == "COVID-19":
@@ -306,27 +325,20 @@ Findings may be compatible with pneumonia, but alternative explanations exist.
 
 Distribution and appearance of opacities are compatible with **COVID-19 pneumonia**.
 
-**Critical Points:**
-- Imaging is **not** diagnostic by itself  
-- **RT-PCR / rapid antigen testing** is mandatory for confirmation  
-
-**If clinically suspected:**
-- Isolate per local infection-control policies  
-- Monitor SpO₂ and respiratory status  
-- Escalate care if:
-  - SpO₂ < 94% on room air  
-  - Increasing work of breathing  
-  - Hemodynamic instability  
+⚠️ Imaging alone is **not diagnostic**.
+
+- Confirm with RT-PCR or validated antigen testing  
+- Follow local isolation and infection-control policies  
+- Monitor SpO₂ and respiratory status; escalate care if deterioration occurs  
 """
         else:
             interpretation += """
 ### 🦠 COVID-19 Pattern – Possible
 
-Some features may overlap with COVID-19, but there is **significant uncertainty**.
+Some features may overlap with COVID-19, but there is substantial uncertainty.
 
-**Do not rely on imaging alone.**
-- Obtain RT-PCR / rapid antigen testing  
-- Use clinical context and epidemiology to guide decisions  
+- Testing (RT-PCR / antigen) is essential  
+- Integrate exposure history and symptoms  
 """
 
     else:  # Normal
@@ -334,12 +346,13 @@ Some features may overlap with COVID-19, but there is **significant uncertainty*
             interpretation += """
 ### ✅ No Major Abnormality Detected
 
-The model did **not** detect features suggestive of TB, pneumonia, or COVID-19.
+The model did **not** detect strong features of TB, pneumonia, or COVID-19.
+
+Important caveats:
 
-**Important Caveats:**
 - Early disease or small lesions may be missed  
 - Non-infective conditions (e.g., cancer, ILD) are **not** specifically evaluated  
-- If symptoms are present, further workup may still be required  
+- Persistent symptoms still warrant clinical review  
 """
         else:
             interpretation += """
@@ -347,43 +360,38 @@ The model did **not** detect features suggestive of TB, pneumonia, or COVID-19.
 
 The scan leans towards **normal**, but the model is not highly confident.
 
-**If symptoms persist:**
-- Consider follow-up imaging  
-- Seek a clinician’s interpretation  
+- Consider follow-up or additional tests if symptoms persist  
 """
 
-    # Universal disclaimer
     interpretation += """
 ---
 ## ⚠️ CRITICAL MEDICAL DISCLAIMER
 
 - This AI model is a **screening / decision-support tool only**  
-- It is **not FDA-approved** and **must not** be used as a stand-alone diagnostic device  
+- It is **not FDA-approved** and must **not** be used as a stand-alone diagnostic device  
 - Always integrate:
   - Clinical history and examination  
-  - Laboratory tests (e.g., sputum, PCR, cultures)  
+  - Laboratory tests (sputum, PCR, cultures, etc.)  
   - Expert radiologist review  
 
-**Gold Standards:**
+**Gold Standards**
+
 - TB: Sputum AFB / culture, GeneXpert MTB/RIF, TB-PCR  
 - Pneumonia: Clinical diagnosis + labs / microbiology  
 - COVID-19: RT-PCR or validated antigen tests  
 
 When in doubt, consult a qualified healthcare professional.
-"""
-
-    interpretation += """
 ---
 🫁 **Powered by Adaptive Sparse Training (AST)**  
-Energy-efficient deep learning – designed to make advanced chest X-ray screening more accessible.
+Energy-efficient deep learning – designed to make chest X-ray screening more accessible.
+
+**Links**
 
-**Links:**
 - GitHub: https://github.com/oluwafemidiakhoa/Tuberculosis  
 - Hugging Face Space: https://huggingface.co/spaces/mgbam/Tuberculosis  
 """
     return interpretation
 
-
 # ============================================================================
 # Prediction Pipeline
 # ============================================================================
@@ -412,7 +420,6 @@ def predict_chest_xray(image, show_gradcam=True, audience="Clinician"):
     original_img = image.copy()
     input_tensor = transform(image).unsqueeze(0).to(device)
 
-    # Inference with optional Grad-CAM
     with torch.set_grad_enabled(show_gradcam):
         if show_gradcam:
             cam, output = grad_cam.generate(input_tensor)
@@ -435,18 +442,15 @@ def predict_chest_xray(image, show_gradcam=True, audience="Clinician"):
         for i in range(len(CLASSES))
     }
 
-    # Visuals
     original_pil = create_original_display(original_img, pred_label, confidence)
     gradcam_viz = create_gradcam_visualization(original_img, cam) if cam is not None else None
     overlay_viz = create_overlay_visualization(original_img, cam) if cam is not None else None
 
     interpretation = create_interpretation(pred_label, confidence, results, audience=audience)
-
-    snapshot = f"**{pred_label}** · {confidence:.1f}% confidence • Sum of probabilities: {prob_sum:.3f}"
+    snapshot = f"**{pred_label}** · {confidence:.1f}% confidence • Prob. sum: {prob_sum:.3f}"
 
     return results, original_pil, gradcam_viz, overlay_viz, interpretation, snapshot
 
-
 # ============================================================================
 # WOW UI / UX – Gradio App
 # ============================================================================
@@ -595,8 +599,7 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
                     <div class="hero-title">🫁 AST Chest X-Ray Lab</div>
                     <div class="hero-subtitle">
                         Multi-class chest X-ray analysis with <b>Explainable AI</b> and
-                        <b>Adaptive Sparse Training</b>.
-                        Designed for TB, Pneumonia, COVID-19 and Normal scans.
+                        <b>Adaptive Sparse Training</b> – Normal · Tuberculosis · Pneumonia · COVID-19.
                     </div>
                     <div class="hero-chip-row">
                         <div class="hero-chip">
@@ -604,7 +607,7 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
                             Live Inference
                         </div>
                         <div class="hero-chip">
-                            4-class EfficientNet · ~{TOTAL_PARAMS_M:.1f}M params
+                            EfficientNet-B0 · ~{TOTAL_PARAMS_M:.1f}M params
                         </div>
                         <div class="hero-chip">
                             95–97% validation accuracy · ~89% energy savings
@@ -620,8 +623,8 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
                         <div class="stat-pill-value">{device.type.upper()}</div>
                     </div>
                     <div class="stat-pill">
-                        <div class="stat-pill-label">Model</div>
-                        <div class="stat-pill-value">EfficientNet-B0 · 4-way classifier</div>
+                        <div class="stat-pill-label">Task</div>
+                        <div class="stat-pill-value">Normal · TB · Pneumonia · COVID-19</div>
                     </div>
                 </div>
             </div>
@@ -632,9 +635,7 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
     gr.Markdown(" ")
 
     with gr.Row(equal_height=True):
-        # ----------------------------------
         # LEFT: INPUT PANEL
-        # ----------------------------------
         with gr.Column(scale=1, elem_classes="glass-card"):
             gr.Markdown("### 1️⃣ Upload & Configure")
 
@@ -666,13 +667,11 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
 
 - Use frontal (PA/AP) chest X-rays in PNG / JPG format  
 - This tool is best used as a **triage / screening assistant**  
-- For noisy images or rotated scans, consider preprocessing before upload  
+- For noisy or rotated images, consider preprocessing before upload  
 """
             )
 
-        # ----------------------------------
         # RIGHT: RESULTS PANEL
-        # ----------------------------------
         with gr.Column(scale=2, elem_classes="glass-card-light"):
             gr.Markdown("### 2️⃣ AI Dashboard")
 
@@ -712,15 +711,15 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
 ### 🧠 Model Card – AST Chest X-Ray
 
 - **Backbone**: EfficientNet-B0  
-- **Task**: 4-way classification (Normal, Tuberculosis, Pneumonia, COVID-19)  
+- **Classes**: Normal, Tuberculosis, Pneumonia, COVID-19  
 - **Optimization**: Sample-based Adaptive Sparse Training (AST)  
 - **Energy Profile**: ~89% training energy reduction vs dense baseline  
 
-**Design Goals**
+**Goals**
 
 1. Provide **fast, explainable triage** support for TB & pneumonia  
-2. Maintain **high specificity**, especially differentiating TB from pneumonia  
-3. Be lightweight enough for **deployment in resource-constrained settings**  
+2. Maintain high specificity, especially for TB vs pneumonia  
+3. Be lightweight enough for deployment in **resource-constrained settings**  
 
 > This model is a research prototype. Do **not** use it as a stand-alone clinical device.
 """
@@ -743,9 +742,7 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
         """
     )
 
-    # ----------------------------------------------------------------------
     # Wiring
-    # ----------------------------------------------------------------------
     analyze_btn.click(
         fn=predict_chest_xray,
         inputs=[image_input, show_gradcam, audience_select],
@@ -760,7 +757,14 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
     )
 
     clear_btn.click(
-        fn=lambda: ({}, None, None, None, "Awaiting image upload…", "Awaiting image upload…"),
+        fn=lambda: (
+            {},
+            None,
+            None,
+            None,
+            "Awaiting image upload…",
+            "Awaiting image upload…",
+        ),
         inputs=None,
         outputs=[
             prob_output,
@@ -772,7 +776,7 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
         ],
     )
 
-    # Example X-rays section (optional – remove if you don't have these paths)
+    # Example X-rays (optional – comment out if you don't have these files)
     gr.Markdown("### 🔍 Try Example X-rays")
     gr.Examples(
         examples=[