Update app.py

app.py

@@ -1,128 +1,76 @@
 """
-🫁 AST Chest X-Ray Lab
-Multi-Class Chest X-Ray Detection (Normal · TB · Pneumonia · COVID-19)
-with Adaptive Sparse Training & Explainable AI (Grad-CAM)
+🫁 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 io
-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
 import torch.nn as nn
-from PIL import Image
 from torchvision import models, transforms
+from PIL import Image
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+from pathlib import Path
+import io
 
 # ============================================================================
 # Model Setup
 # ============================================================================
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# 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, NUM_CLASSES)
-
-# We expect a 4-class EfficientNet checkpoint here
-checkpoint_candidates = [
-    "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}")
-            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 model_state_dict: {e}")
-
-if not loaded:
-    raise RuntimeError(
-        "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."
-    )
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Load model - Using EfficientNet-B2 (trained model architecture)
+model = models.efficientnet_b2(weights=None)
+model.classifier[1] = nn.Linear(model.classifier[1].in_features, 4)  # 4 classes
+
+try:
+    # Try loading best.pt from root directory (HuggingFace Spaces location)
+    model.load_state_dict(torch.load('best.pt', map_location=device))
+    print("✅ Multi-class model loaded successfully from best.pt!")
+except Exception as e:
+    print(f"⚠️  Error loading model from best.pt: {e}")
+    try:
+        # Fallback to checkpoints directory
+        model.load_state_dict(torch.load('checkpoints/best_multiclass.pt', map_location=device))
+        print("✅ Multi-class model loaded successfully from checkpoints/best_multiclass.pt!")
+    except Exception as e2:
+        print(f"❌ CRITICAL ERROR: Could not load model from any location!")
+        print(f"   - best.pt error: {e}")
+        print(f"   - checkpoints/best_multiclass.pt error: {e2}")
+        raise RuntimeError("Model file not found! Please ensure best.pt is uploaded to the Space.")
 
 model = model.to(device)
 model.eval()
 
-TOTAL_PARAMS = sum(p.numel() for p in model.parameters())
-TOTAL_PARAMS_M = TOTAL_PARAMS / 1e6
-
-# ============================================================================
-# Classes & Preprocessing
-# ============================================================================
-
-CLASSES = ["Normal", "Tuberculosis", "Pneumonia", "COVID-19"]
+# Classes
+CLASSES = ['Normal', 'Tuberculosis', 'Pneumonia', 'COVID-19']
 CLASS_COLORS = {
-    "Normal": "#22c55e",        # Green
-    "Tuberculosis": "#ef4444",  # Red
-    "Pneumonia": "#f97316",     # Orange
-    "COVID-19": "#a855f7",      # Purple
+    'Normal': '#2ecc71',       # Green
+    'Tuberculosis': '#e74c3c',  # Red
+    'Pneumonia': '#f39c12',     # Orange
+    'COVID-19': '#9b59b6'       # Purple
 }
 
-transform = transforms.Compose(
-    [
-        transforms.Resize(256),
-        transforms.CenterCrop(224),
-        transforms.ToTensor(),
-        transforms.Normalize(
-            [0.485, 0.456, 0.406],
-            [0.229, 0.224, 0.225],
-        ),
-    ]
-)
+# Image preprocessing
+transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
 
 # ============================================================================
 # Grad-CAM Implementation
 # ============================================================================
 
-
 class GradCAM:
     def __init__(self, model, target_layer):
         self.model = model
@@ -147,7 +95,7 @@ class GradCAM:
 
         self.model.zero_grad()
         one_hot = torch.zeros_like(output)
-        one_hot[0, target_class] = 1
+        one_hot[0][target_class] = 1
         output.backward(gradient=one_hot, retain_graph=True)
 
         if self.gradients is None:
@@ -161,623 +109,471 @@ class GradCAM:
 
         return cam, output
 
-
+# Setup Grad-CAM
 target_layer = model.features[-1]
 grad_cam = GradCAM(model, target_layer)
 
 # ============================================================================
-# Visualization Helpers
+# Prediction Functions
 # ============================================================================
 
+def predict_chest_xray(image, show_gradcam=True):
+    """
+    Predict disease class from chest X-ray with Grad-CAM visualization
+    """
+    if image is None:
+        return None, None, None, None
 
-def _figure_to_pil():
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png", dpi=150, bbox_inches="tight", facecolor="white")
-    plt.close()
-    buf.seek(0)
-    return Image.open(buf)
+    # Convert to PIL if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image).convert('RGB')
+    else:
+        image = image.convert('RGB')
+
+    # Store original for display
+    original_img = image.copy()
+
+    # Preprocess
+    input_tensor = transform(image).unsqueeze(0).to(device)
+
+    # Get prediction with Grad-CAM
+    with torch.set_grad_enabled(show_gradcam):
+        if show_gradcam:
+            cam, output = grad_cam.generate(input_tensor)
+        else:
+            output = model(input_tensor)
+            cam = None
+
+    # Get probabilities
+    probs = torch.softmax(output, dim=1)[0].cpu().detach().numpy()
+
+    # Safety check: ensure probabilities sum to ~1.0
+    prob_sum = np.sum(probs)
+    if not (0.99 <= prob_sum <= 1.01):
+        print(f"⚠️ WARNING: Probability sum is {prob_sum}, not 1.0. Model may not be loaded correctly!")
+
+    pred_class = int(output.argmax(dim=1).item())
+    pred_label = CLASSES[pred_class]
+    confidence = float(probs[pred_class]) * 100
+
+    # Create results - ensure values are between 0-100
+    results = {
+        CLASSES[i]: float(min(100.0, max(0.0, probs[i] * 100))) for i in range(len(CLASSES))
+    }
+
+    # Generate visualizations
+    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
+
+    # Create interpretation text
+    interpretation = create_interpretation(pred_label, confidence, results)
+
+    return results, original_pil, gradcam_viz, overlay_viz, interpretation
 
 def create_original_display(image, pred_label, confidence):
-    fig, ax = plt.subplots(figsize=(7, 7))
+    """Create annotated original image"""
+    fig, ax = plt.subplots(figsize=(8, 8))
     ax.imshow(image)
-    ax.axis("off")
+    ax.axis('off')
 
+    # Add prediction box
     color = CLASS_COLORS[pred_label]
-    title = f"Prediction: {pred_label} • Confidence: {confidence:.1f}%"
-    ax.set_title(
-        title,
-        fontsize=16,
-        fontweight="bold",
-        color=color,
-        pad=20,
-    )
+    title = f'Prediction: {pred_label}\nConfidence: {confidence:.1f}%'
+    ax.set_title(title, fontsize=16, fontweight='bold', color=color, pad=20)
+
     plt.tight_layout()
-    return _figure_to_pil()
 
+    # Convert to PIL
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
+    plt.close()
+    buf.seek(0)
 
-def create_gradcam_visualization(image, cam):
+    return Image.open(buf)
+
+def create_gradcam_visualization(image, cam, pred_label, confidence):
+    """Create Grad-CAM heatmap"""
+    # Resize CAM to image size
     img_array = np.array(image.resize((224, 224)))
     cam_resized = cv2.resize(cam, (224, 224))
 
+    # Create heatmap
     heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
     heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
 
-    fig, ax = plt.subplots(figsize=(7, 7))
+    fig, ax = plt.subplots(figsize=(8, 8))
     ax.imshow(heatmap)
-    ax.axis("off")
-    ax.set_title(
-        "Attention Heatmap\n(Where the model is looking)",
-        fontsize=14,
-        fontweight="bold",
-        pad=20,
-    )
+    ax.axis('off')
+    ax.set_title('Attention Heatmap\n(Areas the model focuses on)',
+                 fontsize=14, fontweight='bold', pad=20)
+
     plt.tight_layout()
-    return _figure_to_pil()
 
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
+    plt.close()
+    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))
 
+    # Create heatmap
     heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
     heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) / 255.0
 
+    # Overlay
     overlay = img_array * 0.5 + heatmap * 0.5
     overlay = np.clip(overlay, 0, 1)
 
-    fig, ax = plt.subplots(figsize=(7, 7))
+    fig, ax = plt.subplots(figsize=(8, 8))
     ax.imshow(overlay)
-    ax.axis("off")
-    ax.set_title(
-        "Explainable AI Overlay\n(Anatomy + Attention)",
-        fontsize=14,
-        fontweight="bold",
-        pad=20,
-    )
+    ax.axis('off')
+    ax.set_title('Explainable AI Visualization\n(Original + Heatmap)',
+                 fontsize=14, fontweight='bold', pad=20)
+
     plt.tight_layout()
-    return _figure_to_pil()
 
-# ============================================================================
-# Interpretation
-# ============================================================================
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', dpi=150, bbox_inches='tight', facecolor='white')
+    plt.close()
+    buf.seek(0)
 
+    return Image.open(buf)
 
-def create_interpretation(pred_label, confidence, results, audience="Clinician"):
-    header_note = {
-        "Clinician": "This view is tuned for **clinical decision support** (not a replacement for your judgement).",
-        "Researcher": "This view is tuned for **model behavior understanding** and experimental workflows.",
-        "Patient / Public": "This view is tuned for **patient-friendly language**. Always discuss results with a doctor.",
-    }.get(audience, "Use this output as a **screening aid**, not a final diagnosis.")
+def create_interpretation(pred_label, confidence, results):
+    """Create interpretation text with improved medical disclaimers"""
 
     interpretation = f"""
-## 🔬 Analysis Results ({audience} View)
-
-> {header_note}
-
-### Primary Prediction: **{pred_label}**
+## 🔬 Analysis Results
+### Prediction: **{pred_label}**
 - Confidence: **{confidence:.1f}%**
-
-### Probability Breakdown
+### Probability Breakdown:
 - 🟢 Normal: **{results['Normal']:.1f}%**
 - 🔴 Tuberculosis: **{results['Tuberculosis']:.1f}%**
 - 🟠 Pneumonia: **{results['Pneumonia']:.1f}%**
 - 🟣 COVID-19: **{results['COVID-19']:.1f}%**
-
 ---
 """
 
-    if pred_label == "Tuberculosis":
+    # Disease-specific interpretations
+    if pred_label == 'Tuberculosis':
         if confidence >= 85:
             interpretation += """
-### 🧫 TB Pattern – High Confidence
-
-The model has detected features strongly suggestive of **pulmonary tuberculosis**.
-
-**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  
+**⚠️ High Confidence TB Detection**
+The model has detected features highly consistent with tuberculosis infection.
+**CRITICAL - Immediate Actions Required:**
+1. ✅ **Immediate consultation** with a healthcare provider
+2. ✅ **Confirmatory sputum test** (AFB smear or GeneXpert MTB/RIF)
+3. ✅ **Clinical correlation** with symptoms:
+   - Persistent cough (>2 weeks)
+   - Fever, especially night sweats
+   - Unexplained weight loss
+   - Hemoptysis (coughing blood)
+4. ✅ **Isolation** and contact tracing if confirmed
+5. ✅ **Chest CT scan** if needed for further evaluation
+**⚠️ IMPORTANT**: This is a SCREENING tool, not a diagnostic tool.
+Clinical diagnosis of TB requires laboratory confirmation (sputum test).
 """
         else:
             interpretation += """
-### 🧫 TB Pattern – Possible
-
-The scan shows features that **could** be compatible with tuberculosis, but confidence is moderate.
-
-- Correlate with symptoms and risk factors  
-- Consider sputum testing where suspicion is non-trivial  
-- Follow-up imaging as clinically indicated  
+**⚠️ Possible TB Detection**
+The model has detected features suggestive of tuberculosis, but confidence is moderate.
+**Recommended Actions:**
+1. Consult healthcare provider for clinical evaluation
+2. Consider confirmatory sputum testing
+3. Evaluate clinical symptoms
+4. Follow-up imaging may be recommended
+**Note**: Moderate confidence requires professional medical evaluation.
 """
 
-    elif pred_label == "Pneumonia":
+    elif pred_label == 'Pneumonia':
         if confidence >= 85:
             interpretation += """
-### 🌫 Pneumonia Pattern – High Confidence
-
-The model has detected an opacity pattern consistent with **pneumonia**.
-
-Typical clinical correlates:
-
-- Fever, productive cough  
-- Shortness of breath  
-- Pleuritic chest pain  
-
-Use alongside clinical exam, labs, and local treatment guidelines.
+**⚠️ High Confidence Pneumonia Detection**
+The model has detected features consistent with pneumonia (bacterial or viral).
+**Recommended Actions:**
+1. ✅ **Medical evaluation** for pneumonia diagnosis
+2. ✅ **Possible confirmatory tests**:
+   - Sputum culture
+   - Blood tests (WBC count, CRP)
+   - Additional chest imaging if needed
+3. ✅ **Clinical correlation** with symptoms:
+   - Cough with sputum production
+   - Fever and chills
+   - Shortness of breath
+   - Chest pain with breathing
+4. ✅ **Treatment**: Antibiotics (bacterial) or supportive care (viral)
+**Note**: Pneumonia can present similarly to other lung diseases.
+Professional diagnosis is essential for appropriate treatment.
 """
         else:
             interpretation += """
-### 🌫 Pneumonia Pattern – Possible
-
-Findings may be compatible with pneumonia, but alternative explanations exist.
-
-- Check vital signs and auscultation  
-- Labs (WBC, CRP, cultures) may be useful  
-- Consider watchful follow-up or repeat imaging  
+**⚠️ Possible Pneumonia**
+Features suggest possible pneumonia, but further evaluation is needed.
+**Recommended Actions:**
+1. Seek medical evaluation
+2. Clinical symptom assessment
+3. Consider additional diagnostic tests
+**Note**: Requires professional medical evaluation for confirmation.
 """
 
-    elif pred_label == "COVID-19":
+    elif pred_label == 'COVID-19':
         if confidence >= 85:
             interpretation += """
-### 🦠 COVID-19 Pattern – High Confidence
-
-Distribution and appearance of opacities are compatible with **COVID-19 pneumonia**.
-
-⚠️ 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  
+**⚠️ High Confidence COVID-19 Detection**
+The model has detected features consistent with COVID-19 pneumonia.
+**URGENT - Immediate Actions:**
+1. ✅ **COVID-19 RT-PCR test** for confirmation
+2. ✅ **Isolation** to prevent transmission
+3. ✅ **Monitor oxygen saturation** (SpO2 levels)
+4. ✅ **Seek immediate medical care** if:
+   - Difficulty breathing
+   - SpO2 < 94%
+   - Persistent chest pain
+   - Confusion or inability to stay awake
+5. ✅ **Contact tracing** if positive
+**Clinical Symptoms to Monitor:**
+- Fever, cough, shortness of breath
+- Loss of taste/smell
+- Fatigue, body aches
+- Gastrointestinal symptoms
+**⚠️ IMPORTANT**: Imaging findings alone cannot confirm COVID-19.
+RT-PCR or antigen testing is required for diagnosis.
 """
         else:
             interpretation += """
-### 🦠 COVID-19 Pattern – Possible
-
-Some features may overlap with COVID-19, but there is substantial uncertainty.
-
-- Testing (RT-PCR / antigen) is essential  
-- Integrate exposure history and symptoms  
+**⚠️ Possible COVID-19**
+Features suggest possible COVID-19, but confirmation testing is essential.
+**Recommended Actions:**
+1. Get RT-PCR or rapid antigen test
+2. Self-isolate pending test results
+3. Monitor symptoms
+4. Seek medical care if symptoms worsen
+**Note**: COVID-19 diagnosis requires laboratory confirmation.
 """
 
     else:  # Normal
         if confidence >= 85:
             interpretation += """
-### ✅ No Major Abnormality Detected
-
-The model did **not** detect strong features of TB, pneumonia, or COVID-19.
-
-Important caveats:
-
-- Early disease or small lesions may be missed  
-- Non-infective conditions (e.g., cancer, ILD) are **not** specifically evaluated  
-- Persistent symptoms still warrant clinical review  
+**✅ High Confidence Normal Result**
+The model has not detected significant abnormalities consistent with TB, pneumonia, or COVID-19.
+**Interpretation:**
+- Chest X-ray appears within normal limits
+- No features of active tuberculosis detected
+- No signs of pneumonia or COVID-19
+**Important Notes:**
+- This does NOT rule out all lung diseases
+- Early-stage diseases may not show on X-ray
+- If you have symptoms, seek medical evaluation
+- Regular health screenings are recommended
+**When to still see a doctor:**
+- Persistent cough, fever, or respiratory symptoms
+- Unexplained weight loss or night sweats
+- Shortness of breath or chest pain
+- Known exposure to TB or COVID-19
 """
         else:
             interpretation += """
-### ℹ️ Likely Normal, But Low Confidence
-
-The scan leans towards **normal**, but the model is not highly confident.
-
-- Consider follow-up or additional tests if symptoms persist  
+**⚠️ Likely Normal, Low Confidence**
+The model suggests a normal chest X-ray, but confidence is not high.
+**Recommended Actions:**
+1. If symptomatic, seek medical evaluation
+2. Consider repeat imaging if concerns persist
+3. Clinical correlation is important
+**Note**: Low confidence results should be reviewed by healthcare professionals.
 """
 
+    # Add 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  
-- Always integrate:
-  - Clinical history and examination  
-  - Laboratory tests (sputum, PCR, cultures, etc.)  
-  - Expert radiologist review  
-
-**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.
+### Model Capabilities:
+- ✅ Trained on 4 disease classes: Normal, TB, Pneumonia, COVID-19
+- ✅ Can distinguish between different lung diseases
+- ✅ ~95-97% accuracy in validation testing
+- ✅ Powered by Adaptive Sparse Training (89% energy efficient)
+### Important Limitations:
+- ⚠️ This is a **SCREENING tool**, not a diagnostic device
+- ⚠️ **NOT FDA-approved** for clinical diagnosis
+- ⚠️ Cannot detect: lung cancer, pulmonary fibrosis, bronchiectasis, other rare diseases
+- ⚠️ Cannot replace: professional radiologist review
+- ⚠️ Cannot confirm: laboratory diagnosis (sputum tests, PCR, cultures)
+### Clinical Use Guidelines:
+1. ✅ Use as a **preliminary screening** tool only
+2. ✅ ALL positive results require **confirmatory laboratory testing**
+3. ✅ ALL cases require **clinical correlation** with symptoms and history
+4. ✅ Expert radiologist review is recommended for clinical decisions
+5. ✅ Do NOT initiate treatment based solely on AI predictions
+### Diagnostic Gold Standards:
+- **TB**: Sputum AFB smear/culture, GeneXpert MTB/RIF, TB-PCR
+- **Pneumonia**: Clinical diagnosis + sputum culture + blood tests
+- **COVID-19**: RT-PCR, rapid antigen test
+**When in doubt, always consult a qualified healthcare professional.**
 ---
-🫁 **Powered by Adaptive Sparse Training (AST)**  
-Energy-efficient deep learning – designed to make chest X-ray screening more accessible.
-
-**Links**
-
-- GitHub: https://github.com/oluwafemidiakhoa/Tuberculosis  
-- Hugging Face Space: https://huggingface.co/spaces/mgbam/Tuberculosis  
+🫁 **Powered by Adaptive Sparse Training**
+Energy-efficient AI for accessible healthcare
+**Learn more:**
+- GitHub: https://github.com/oluwafemidiakhoa/Tuberculosis
+- Research: Sample-based Adaptive Sparse Training for deep learning
 """
-    return interpretation
-
-# ============================================================================
-# Prediction Pipeline
-# ============================================================================
-
-
-def predict_chest_xray(image, show_gradcam=True, audience="Clinician"):
-    """
-    Main inference function used by Gradio.
-    Returns:
-        - dict of class probabilities
-        - annotated original
-        - grad-cam heatmap
-        - overlay
-        - full markdown report
-        - short textual snapshot
-    """
-    if image is None:
-        msg = "👋 Upload a chest X-ray (PNG/JPG) and click **Analyze** to generate a full AI report."
-        return {}, None, None, None, msg, "Awaiting image upload…"
-
-    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].cpu().detach().numpy()
-    prob_sum = float(np.sum(probs))
-
-    if not (0.99 <= prob_sum <= 1.01):
-        print(f"⚠️ WARNING: Probability sum is {prob_sum}, not ≈1.0 – check model weights.")
-
-    pred_class = int(output.argmax(dim=1).item())
-    pred_label = CLASSES[pred_class]
-    confidence = float(probs[pred_class]) * 100.0
-
-    results = {
-        CLASSES[i]: float(min(100.0, max(0.0, probs[i] * 100.0)))
-        for i in range(len(CLASSES))
-    }
-
-    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 • Prob. sum: {prob_sum:.3f}"
-
-    return results, original_pil, gradcam_viz, overlay_viz, interpretation, snapshot
+    return interpretation
 
 # ============================================================================
-# WOW UI / UX – Gradio App
+# Gradio Interface
 # ============================================================================
 
+# Custom CSS
 custom_css = """
-:root {
-    --primary: #6366f1;
-    --primary-soft: rgba(99, 102, 241, 0.12);
-    --accent: #ec4899;
-}
-
-.gradio-container {
-    font-family: system-ui, -apple-system, BlinkMacSystemFont, "Inter", sans-serif;
-    background: radial-gradient(circle at top left, #111827 0, #020617 50%, #020617 100%);
-    color: #e5e7eb;
-}
-
-#hero {
-    padding: 24px 24px 8px 24px;
-    border-radius: 24px;
-    background: linear-gradient(120deg, rgba(99,102,241,0.18), rgba(236,72,153,0.14));
-    border: 1px solid rgba(148, 163, 184, 0.4);
-    box-shadow: 0 24px 60px rgba(15,23,42,0.85);
-    backdrop-filter: blur(18px);
-}
-
-.hero-title {
-    font-size: 2.4rem;
-    font-weight: 800;
-    letter-spacing: 0.04em;
-    color: #f9fafb;
-    margin-bottom: 6px;
-}
-
-.hero-subtitle {
-    font-size: 0.98rem;
-    color: #e5e7eb;
+#main-container {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    padding: 20px;
 }
-
-.hero-chip-row {
-    display: flex;
-    flex-wrap: wrap;
-    gap: 8px;
-    margin-top: 14px;
-}
-
-.hero-chip {
-    padding: 4px 10px;
-    border-radius: 999px;
-    font-size: 0.78rem;
-    background: rgba(15,23,42,0.8);
-    border: 1px solid rgba(148,163,184,0.5);
-    display: inline-flex;
-    align-items: center;
-    gap: 6px;
-    color: #e5e7eb;
-}
-
-.pulse-dot {
-    width: 8px;
-    height: 8px;
-    border-radius: 999px;
-    background: #22c55e;
-    box-shadow: 0 0 0 0 rgba(34,197,94,0.7);
-    animation: pulse 1.4s infinite;
-}
-
-@keyframes pulse {
-  0% { box-shadow: 0 0 0 0 rgba(34,197,94,0.7); }
-  70% { box-shadow: 0 0 0 10px rgba(34,197,94,0); }
-  100% { box-shadow: 0 0 0 0 rgba(34,197,94,0); }
-}
-
-.glass-card {
-    background: rgba(15,23,42,0.82);
-    border-radius: 18px;
-    border: 1px solid rgba(148,163,184,0.4);
-    box-shadow: 0 18px 40px rgba(15,23,42,0.85);
-    padding: 18px;
-    backdrop-filter: blur(16px);
+#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);
 }
-
-.glass-card-light {
-    background: rgba(15,23,42,0.65);
-    border-radius: 18px;
-    border: 1px solid rgba(148,163,184,0.3);
-    box-shadow: 0 12px 24px rgba(15,23,42,0.85);
-    padding: 16px;
-    backdrop-filter: blur(12px);
+#subtitle {
+    text-align: center;
+    color: #f0f0f0;
+    font-size: 1.2em;
+    margin-bottom: 20px;
 }
-
-.stat-pill {
-    padding: 10px 12px;
-    border-radius: 14px;
-    background: rgba(15,23,42,0.9);
-    border: 1px solid rgba(148,163,184,0.5);
-    font-size: 0.78rem;
-    display: flex;
-    flex-direction: column;
-    gap: 2px;
+#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);
 }
-
-.stat-pill-label {
-    color: #9ca3af;
-    text-transform: uppercase;
-    font-size: 0.68rem;
+.gradio-container {
+    font-family: 'Inter', sans-serif;
 }
-
-.stat-pill-value {
-    color: #e5e7eb;
-    font-weight: 600;
+#upload-box {
+    border: 3px dashed #667eea;
+    border-radius: 15px;
+    padding: 20px;
+    background: rgba(255,255,255,0.95);
 }
-
-.dropzone-image img {
-    border-radius: 16px !important;
+#results-box {
+    background: white;
+    border-radius: 15px;
+    padding: 20px;
+    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
 }
-
-.output-image img {
-    border-radius: 16px !important;
+.output-image {
+    border-radius: 10px;
+    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
 }
-
 footer {
     text-align: center;
-    margin-top: 24px;
-    color: #9ca3af;
-    font-size: 0.78rem;
+    margin-top: 30px;
+    color: white;
+    font-size: 0.9em;
 }
 """
 
-theme = gr.themes.Soft(
-    primary_hue="indigo",
-    secondary_hue="pink",
-    neutral_hue="slate",
-).set(
-    button_primary_background_fill="linear-gradient(135deg,#4f46e5,#ec4899)",
-    button_primary_background_fill_hover="linear-gradient(135deg,#6366f1,#f97316)",
-)
-
-with gr.Blocks(css=custom_css, theme=theme) as demo:
-    # HERO
-    gr.HTML(
-        f"""
-        <div id="hero">
-            <div style="display:flex;justify-content:space-between;gap:16px;align-items:flex-start;">
-                <div>
-                    <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> – Normal · Tuberculosis · Pneumonia · COVID-19.
-                    </div>
-                    <div class="hero-chip-row">
-                        <div class="hero-chip">
-                            <span class="pulse-dot"></span>
-                            Live Inference
-                        </div>
-                        <div class="hero-chip">
-                            EfficientNet-B0 · ~{TOTAL_PARAMS_M:.1f}M params
-                        </div>
-                        <div class="hero-chip">
-                            95–97% validation accuracy · ~89% energy savings
-                        </div>
-                        <div class="hero-chip">
-                            {MODEL_LOAD_INFO}
-                        </div>
-                    </div>
-                </div>
-                <div style="min-width:210px;display:flex;flex-direction:column;gap:8px;">
-                    <div class="stat-pill">
-                        <div class="stat-pill-label">Device</div>
-                        <div class="stat-pill-value">{device.type.upper()}</div>
-                    </div>
-                    <div class="stat-pill">
-                        <div class="stat-pill-label">Task</div>
-                        <div class="stat-pill-value">Normal · TB · Pneumonia · COVID-19</div>
-                    </div>
-                </div>
+# Create interface
+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
+                <br><br>
+                <b>Detects:</b> Normal • Tuberculosis • Pneumonia • COVID-19
             </div>
         </div>
-        """
-    )
-
-    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")
+    """)
 
+    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="Drop a chest X-ray here",
-                elem_classes=["dropzone-image"],
+                label="Upload X-Ray Image",
+                elem_classes="output-image"
             )
 
-            with gr.Row():
-                show_gradcam = gr.Checkbox(
-                    value=True,
-                    label="Explainable AI (Grad-CAM)",
-                    info="Highlight regions that drive the prediction",
-                )
-                audience_select = gr.Radio(
-                    ["Clinician", "Researcher", "Patient / Public"],
-                    value="Clinician",
-                    label="Report Style",
-                )
-
-            with gr.Row():
-                analyze_btn = gr.Button("🔬 Analyze X-Ray", variant="primary", scale=3)
-                clear_btn = gr.Button("🧹 Reset", variant="secondary")
-
-            gr.Markdown(
-                """
-**Tips**
+            show_gradcam = gr.Checkbox(
+                value=True,
+                label="Enable Grad-CAM Visualization (Explainable AI)",
+                info="Shows which areas the model focuses on"
+            )
 
-- Use frontal (PA/AP) chest X-rays in PNG / JPG format  
-- This tool is best used as a **triage / screening assistant**  
-- For noisy or rotated images, consider preprocessing before upload  
-"""
+            analyze_btn = gr.Button(
+                "🔬 Analyze X-Ray",
+                variant="primary",
+                size="lg"
             )
 
-        # RIGHT: RESULTS PANEL
-        with gr.Column(scale=2, elem_classes="glass-card-light"):
-            gr.Markdown("### 2️⃣ AI Dashboard")
+            gr.Markdown("""
+            ### 📋 Supported Images:
+            - Chest X-rays (PA or AP view)
+            - PNG, JPG, JPEG formats
+            - Grayscale or RGB
+            ### ⚡ What's New:
+            - ✅ **Improved Specificity**: Can distinguish TB from Pneumonia
+            - ✅ **4 Disease Classes**: Normal, TB, Pneumonia, COVID-19
+            - ✅ **Fewer False Positives**: <5% on pneumonia cases
+            - ✅ **Same Energy Efficiency**: 89% savings with AST
+            """)
+
+        with gr.Column(scale=2, elem_id="results-box"):
+            gr.Markdown("## 📊 Analysis Results")
+
+            # Results display
+            with gr.Row():
+                prob_output = gr.Label(
+                    label="Prediction Confidence",
+                    num_top_classes=4
+                )
 
             with gr.Tabs():
-                with gr.Tab("Snapshot"):
-                    snapshot_output = gr.Markdown(
-                        "No scan analyzed yet. Upload an X-ray to get started."
-                    )
-                    prob_output = gr.Label(
-                        label="Prediction Confidence (All Classes)",
-                        num_top_classes=4,
-                    )
-
-                with gr.Tab("Visual Explanations"):
-                    with gr.Row():
-                        original_output = gr.Image(
-                            label="Annotated X-ray",
-                            elem_classes=["output-image"],
-                        )
-                        gradcam_output = gr.Image(
-                            label="Attention Heatmap",
-                            elem_classes=["output-image"],
-                        )
-                    overlay_output = gr.Image(
-                        label="Explainable Overlay",
-                        elem_classes=["output-image"],
+                with gr.Tab("Original"):
+                    original_output = gr.Image(
+                        label="Annotated X-Ray",
+                        elem_classes="output-image"
                     )
 
-                with gr.Tab("Full Report"):
-                    interpretation_output = gr.Markdown(
-                        "The full clinical / research report will appear here after inference."
+                with gr.Tab("Grad-CAM Heatmap"):
+                    gradcam_output = gr.Image(
+                        label="Attention Heatmap",
+                        elem_classes="output-image"
                     )
 
-                with gr.Tab("Model Card"):
-                    gr.Markdown(
-                        f"""
-### 🧠 Model Card – AST Chest X-Ray
-
-- **Backbone**: EfficientNet-B0  
-- **Classes**: Normal, Tuberculosis, Pneumonia, COVID-19  
-- **Optimization**: Sample-based Adaptive Sparse Training (AST)  
-- **Energy Profile**: ~89% training energy reduction vs dense baseline  
-
-**Goals**
-
-1. Provide **fast, explainable triage** support for TB & pneumonia  
-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.
-"""
+                with gr.Tab("Overlay"):
+                    overlay_output = gr.Image(
+                        label="Explainable AI Visualization",
+                        elem_classes="output-image"
                     )
 
-    gr.Markdown("---")
-
-    gr.HTML(
-        """
-        <footer>
-            <p>
-                <b>AST Chest X-Ray Lab</b> · Normal · TB · Pneumonia · COVID-19 · Explainable AI<br/>
-                Built for research, education, and early-stage screening support.
-            </p>
-            <p style="margin-top:6px;">
-                ⚠️ <b>MEDICAL DISCLAIMER:</b> This tool is not FDA-approved and cannot replace a clinician
-                or radiologist. All decisions must be made by qualified healthcare professionals.
-            </p>
-        </footer>
-        """
-    )
-
-    # Wiring
-    analyze_btn.click(
-        fn=predict_chest_xray,
-        inputs=[image_input, show_gradcam, audience_select],
-        outputs=[
-            prob_output,
-            original_output,
-            gradcam_output,
-            overlay_output,
-            interpretation_output,
-            snapshot_output,
-        ],
-    )
-
-    clear_btn.click(
-        fn=lambda: (
-            {},
-            None,
-            None,
-            None,
-            "Awaiting image upload…",
-            "Awaiting image upload…",
-        ),
-        inputs=None,
-        outputs=[
-            prob_output,
-            original_output,
-            gradcam_output,
-            overlay_output,
-            interpretation_output,
-            snapshot_output,
-        ],
-    )
+            interpretation_output = gr.Markdown(
+                label="Clinical Interpretation"
+            )
 
-    # Example X-rays (optional – comment out if you don't have these files)
-    gr.Markdown("### 🔍 Try Example X-rays")
+    # Example images
+    gr.Markdown("## 📁 Example X-Rays")
     gr.Examples(
         examples=[
             ["examples/normal.png"],
@@ -786,16 +582,44 @@ with gr.Blocks(css=custom_css, theme=theme) as demo:
             ["examples/covid.png"],
         ],
         inputs=image_input,
+        label="Click to load example"
     )
 
-# ============================================================================
-# Launch
-# ============================================================================
+    # Connect components
+    analyze_btn.click(
+        fn=predict_chest_xray,
+        inputs=[image_input, show_gradcam],
+        outputs=[prob_output, original_output, gradcam_output, overlay_output, interpretation_output]
+    )
+
+    # Footer
+    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 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,
+        show_error=True
     )
+