Create app.py

app.py

@@ -0,0 +1,875 @@
+"""
+🫁 AST Chest X-Ray Lab
+Multi-Class Chest X-Ray Detection (Normal · TB · Pneumonia · COVID-19)
+with Adaptive Sparse Training & Explainable AI (Grad-CAM)
+
+This app is a research / screening tool – not a diagnostic device.
+"""
+
+import io
+from pathlib import Path
+
+import cv2
+import gradio as gr
+import matplotlib
+matplotlib.use("Agg")  # non-interactive 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
+
+# ============================================================================
+# Model Setup
+# ============================================================================
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+NUM_CLASSES = 4
+
+# Backbone: EfficientNet-B0 with 4-class head
+model = models.efficientnet_b0(weights=None)
+model.classifier[1] = nn.Linear(model.classifier[1].in_features, NUM_CLASSES)
+
+# Where we expect the (4-class) checkpoint to live
+checkpoint_candidates = [
+    "checkpoints/best.pt",   # main location (from your HF screenshot)
+    "best.pt",               # optional fallback in root
+]
+
+MODEL_LOAD_INFO = ""
+loaded = False
+
+
+def extract_state_dict(ckpt):
+    """
+    Handle both:
+    - plain state_dict (just param tensors)
+    - training checkpoints: keys like 'model_state_dict', 'state_dict', 'model', 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
+
+
+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)
+
+            # Sanity check: classifier head must be 4-way
+            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}."
+                    )
+
+            model.load_state_dict(state_dict, strict=True)
+
+            MODEL_LOAD_INFO = (
+                f"✅ Model loaded from <b>{ckpt_path}</b> on <b>{device.type.upper()}</b>."
+            )
+            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'.\n"
+        "If you saved a training checkpoint, make sure it contains "
+        "a 'model_state_dict' key with the 4-class EfficientNet weights."
+    )
+
+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"]
+CLASS_COLORS = {
+    "Normal": "#22c55e",        # Green
+    "Tuberculosis": "#ef4444",  # Red
+    "Pneumonia": "#f97316",     # Orange
+    "COVID-19": "#a855f7",      # 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],
+        ),
+    ]
+)
+
+# ============================================================================
+# Grad-CAM Implementation
+# ============================================================================
+
+
+class GradCAM:
+    def __init__(self, model, target_layer):
+        self.model = model
+        self.target_layer = target_layer
+        self.gradients = None
+        self.activations = None
+
+        def save_gradient(grad):
+            self.gradients = grad
+
+        def save_activation(module, input, output):
+            self.activations = output.detach()
+
+        target_layer.register_forward_hook(save_activation)
+        target_layer.register_full_backward_hook(lambda m, gi, go: save_gradient(go[0]))
+
+    def generate(self, input_image, target_class=None):
+        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:
+            return None, output
+
+        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_layer = model.features[-1]
+grad_cam = GradCAM(model, target_layer)
+
+# ============================================================================
+# Visualization Helpers
+# ============================================================================
+
+
+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)
+
+
+def create_original_display(image, pred_label, confidence):
+    fig, ax = plt.subplots(figsize=(7, 7))
+    ax.imshow(image)
+    ax.axis("off")
+
+    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,
+    )
+    plt.tight_layout()
+    return _figure_to_pil()
+
+
+def create_gradcam_visualization(image, cam):
+    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=(7, 7))
+    ax.imshow(heatmap)
+    ax.axis("off")
+    ax.set_title(
+        "Attention Heatmap\n(Where the model is focusing)",
+        fontsize=14,
+        fontweight="bold",
+        pad=20,
+    )
+    plt.tight_layout()
+    return _figure_to_pil()
+
+
+def create_overlay_visualization(image, cam):
+    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=(7, 7))
+    ax.imshow(overlay)
+    ax.axis("off")
+    ax.set_title(
+        "Explainable AI Overlay\n(Anatomy + Model Attention)",
+        fontsize=14,
+        fontweight="bold",
+        pad=20,
+    )
+    plt.tight_layout()
+    return _figure_to_pil()
+
+
+def create_probability_bar(results):
+    """Horizontal bar chart of 4-class probabilities."""
+    classes = list(results.keys())
+    values = [results[c] for c in classes]
+    y_pos = np.arange(len(classes))
+
+    fig, ax = plt.subplots(figsize=(6.4, 3.5))
+    ax.barh(y_pos, values)
+    ax.set_yticks(y_pos)
+    ax.set_yticklabels(classes)
+    ax.invert_yaxis()
+    ax.set_xlim(0, 100)
+    ax.set_xlabel("Probability (%)")
+    ax.set_title("Probability Profile by Class", fontsize=12, fontweight="bold")
+    for i, v in enumerate(values):
+        ax.text(v + 1, i, f"{v:.1f}%", va="center", fontsize=9)
+    plt.tight_layout()
+    return _figure_to_pil()
+
+# ============================================================================
+# Interpretation
+# ============================================================================
+
+
+def triage_label(pred_label, confidence):
+    """
+    Simple triage categorisation for clinicians / dashboards.
+    """
+    high = confidence >= 85
+    moderate = 65 <= confidence < 85
+
+    if pred_label == "Normal":
+        if high:
+            return "🟢 Low risk – no major abnormality detected (model view)"
+        elif moderate:
+            return "🟡 Likely normal, but low confidence – consider clinical context"
+        else:
+            return "🟡 Indeterminate – imaging looks close to normal, but model is uncertain"
+    else:
+        if high:
+            return "🔴 High risk finding – prioritise expert review"
+        elif moderate:
+            return "🟠 Possible pathology – correlate with symptoms and labs"
+        else:
+            return "🟡 Weak signal – treat as a soft flag, not a diagnosis"
+
+
+def create_interpretation(pred_label, confidence, results, audience="Clinician"):
+    header_note = {
+        "Clinician": "Optimised for **clinical decision support** – not a replacement for your judgement.",
+        "Researcher": "Optimised for **model behaviour analysis** and experimental workflows.",
+        "Patient / Public": "Optimised for **patient-friendly language**. Always discuss results with a doctor.",
+    }.get(audience, "Use this output as a **screening aid**, not a final diagnosis.")
+
+    interpretation = f"""
+## 🔬 Analysis Results ({audience} View)
+
+> {header_note}
+
+### Primary Prediction: **{pred_label}**
+- Confidence: **{confidence:.1f}%**
+- Triage comment: {triage_label(pred_label, confidence)}
+
+### Probability Breakdown
+- 🟢 Normal: **{results['Normal']:.1f}%**
+- 🔴 Tuberculosis: **{results['Tuberculosis']:.1f}%**
+- 🟠 Pneumonia: **{results['Pneumonia']:.1f}%**
+- 🟣 COVID-19: **{results['COVID-19']:.1f}%**
+
+---
+"""
+
+    # Disease-specific narrative
+    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. Prompt review by a clinician / chest physician  
+2. Sputum testing (AFB smear, GeneXpert MTB/RIF, or TB-PCR)  
+3. Correlate with symptoms:
+   - Chronic cough (>2 weeks)
+   - Weight loss, night sweats
+   - Fever, fatigue
+   - Haemoptysis  
+4. Consider CT or further imaging if discordant with clinical picture  
+5. Infection control and contact tracing as per TB guidelines  
+"""
+        else:
+            interpretation += """
+### 🧫 TB Pattern – Possible
+
+There are features that **could** be compatible with TB, but the confidence is moderate.
+
+- Review history and risk factors  
+- Consider sputum testing if suspicion is non-trivial  
+- Follow-up imaging where indicated  
+"""
+
+    elif pred_label == "Pneumonia":
+        if confidence >= 85:
+            interpretation += """
+### 🌫 Pneumonia Pattern – High Confidence
+
+The model has detected an opacity pattern consistent with **pneumonia**.
+
+Typical clinical picture may include:
+
+- Fever, productive cough  
+- Shortness of breath  
+- Pleuritic chest pain  
+
+Use in combination with examination, labs (WBC, CRP, cultures) and local treatment guidelines.
+"""
+        else:
+            interpretation += """
+### 🌫 Pneumonia Pattern – Possible
+
+Findings may be compatible with pneumonia, but alternative explanations exist.
+
+- Check vitals and respiratory exam  
+- Labs and microbiology can support or refute the impression  
+- Consider watchful follow-up or repeat imaging  
+"""
+
+    elif pred_label == "COVID-19":
+        if confidence >= 85:
+            interpretation += """
+### 🦠 COVID-19 Pattern – High Confidence
+
+The distribution and appearance of opacities are compatible with **COVID-19 pneumonia**.
+
+⚠️ Imaging alone is **not diagnostic**. Key points:
+
+- Confirmation requires RT-PCR or validated antigen testing  
+- Follow local isolation and infection-control policies  
+- Monitor SpO₂ and work of breathing; escalate care if deterioration occurs  
+"""
+        else:
+            interpretation += """
+### 🦠 COVID-19 Pattern – Possible
+
+There are features that could overlap with COVID-19, but uncertainty is substantial.
+
+- Testing (RT-PCR / antigen) is essential  
+- Integrate exposure history, symptoms, and public health guidance  
+"""
+
+    else:  # Normal
+        if confidence >= 85:
+            interpretation += """
+### ✅ No Major Abnormality Detected (Model View)
+
+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. malignancy, ILD) are **not** specifically evaluated  
+- Persistent or unexplained symptoms still require clinical review  
+"""
+        else:
+            interpretation += """
+### ℹ️ Likely Normal, But Low Confidence
+
+The scan leans towards **normal**, but the model's confidence is limited.
+
+- Consider repeat imaging, further tests, or expert review if symptoms persist  
+"""
+
+    interpretation += """
+---
+## ⚠️ CRITICAL MEDICAL DISCLAIMER
+
+- This AI system is a **screening / decision-support tool** only  
+- It is **not FDA-approved**, CE-marked, or licensed as a medical device  
+- It must **not** be used as a stand-alone diagnostic system  
+
+Always integrate:
+- Clinical history and examination  
+- Laboratory tests (e.g. sputum AFB / GeneXpert, PCR, cultures)  
+- Radiologist / specialist interpretation  
+
+**Gold Standards**
+
+- Tuberculosis: 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.
+---
+🫁 **Powered by Adaptive Sparse Training (AST)**  
+Energy-efficient deep learning to support lung health in both high-resource and low-resource settings.
+
+**Project Links**
+
+- GitHub: https://github.com/oluwafemidiakhoa/Tuberculosis  
+- Hugging Face Space: https://huggingface.co/spaces/mgbam/Tuberculosis  
+"""
+    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
+        - probability bar-chart image
+    """
+    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…", None
+
+    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
+    prob_chart = create_probability_bar(results)
+
+    interpretation = create_interpretation(pred_label, confidence, results, audience=audience)
+
+    snapshot = (
+        f"### 📝 Triage Snapshot\n\n"
+        f"- **Finding:** {pred_label}\n"
+        f"- **Model confidence:** {confidence:.1f}%\n"
+        f"- **Triage comment:** {triage_label(pred_label, confidence)}\n"
+        f"- **Probability sum (sanity check):** {prob_sum:.3f}"
+    )
+
+    return results, original_pil, gradcam_viz, overlay_viz, interpretation, snapshot, prob_chart
+
+# ============================================================================
+# WOW UI / UX – Gradio App
+# ============================================================================
+
+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.22), rgba(236,72,153,0.18));
+    border: 1px solid rgba(148, 163, 184, 0.45);
+    box-shadow: 0 24px 60px rgba(15,23,42,0.9);
+    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;
+}
+
+.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.86);
+    border-radius: 18px;
+    border: 1px solid rgba(148,163,184,0.4);
+    box-shadow: 0 18px 40px rgba(15,23,42,0.9);
+    padding: 18px;
+    backdrop-filter: blur(16px);
+}
+
+.glass-card-light {
+    background: rgba(15,23,42,0.7);
+    border-radius: 18px;
+    border: 1px solid rgba(148,163,184,0.3);
+    box-shadow: 0 12px 30px rgba(15,23,42,0.9);
+    padding: 16px;
+    backdrop-filter: blur(12px);
+}
+
+.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;
+}
+
+.stat-pill-label {
+    color: #9ca3af;
+    text-transform: uppercase;
+    font-size: 0.68rem;
+}
+
+.stat-pill-value {
+    color: #e5e7eb;
+    font-weight: 600;
+}
+
+.dropzone-image img,
+.output-image img {
+    border-radius: 16px !important;
+}
+
+footer {
+    text-align: center;
+    margin-top: 24px;
+    color: #9ca3af;
+    font-size: 0.78rem;
+}
+"""
+
+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.
+                        Designed to support clinicians, researchers, and global health teams.
+                    </div>
+                    <div class="hero-chip-row">
+                        <div class="hero-chip">
+                            <span class="pulse-dot"></span>
+                            Live Inference (Research Prototype)
+                        </div>
+                        <div class="hero-chip">
+                            EfficientNet-B0 · ~{TOTAL_PARAMS_M:.1f}M parameters
+                        </div>
+                        <div class="hero-chip">
+                            95–97% validation accuracy · ~89% training energy savings
+                        </div>
+                        <div class="hero-chip">
+                            {MODEL_LOAD_INFO}
+                        </div>
+                    </div>
+                </div>
+                <div style="min-width:220px;display:flex;flex-direction:column;gap:8px;">
+                    <div class="stat-pill">
+                        <div class="stat-pill-label">Compute</div>
+                        <div class="stat-pill-value">{device.type.upper()}</div>
+                    </div>
+                    <div class="stat-pill">
+                        <div class="stat-pill-label">Use Case</div>
+                        <div class="stat-pill-value">Triage & decision support (not diagnostic)</div>
+                    </div>
+                </div>
+            </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")
+
+            image_input = gr.Image(
+                type="pil",
+                label="Drop a chest X-ray here",
+                elem_classes=["dropzone-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(
+                """
+**Usage Notes**
+
+- Best for frontal (PA/AP) chest X-rays in PNG / JPG format  
+- Intended for **triage, education, and research**, not final diagnosis  
+- For off-axis, noisy, or portable images, interpret outputs with extra caution  
+"""
+            )
+
+        # RIGHT: RESULTS PANEL
+        with gr.Column(scale=2, elem_classes="glass-card-light"):
+            gr.Markdown("### 2️⃣ AI Dashboard")
+
+            with gr.Tabs():
+                with gr.Tab("Triage Snapshot"):
+                    snapshot_output = gr.Markdown(
+                        "No scan analysed yet. Upload an X-ray to get started."
+                    )
+                    with gr.Row():
+                        prob_output = gr.Label(
+                            label="Prediction Confidence (All Classes)",
+                            num_top_classes=4,
+                        )
+                        prob_chart_output = gr.Image(
+                            label="Probability Profile",
+                            elem_classes=["output-image"],
+                        )
+
+                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("Full Report"):
+                    interpretation_output = gr.Markdown(
+                        "The full clinical / research report will appear here after inference."
+                    )
+
+                with gr.Tab("Model Card"):
+                    gr.Markdown(
+                        f"""
+### 🧠 Model Card – AST Chest X-Ray
+
+- **Backbone**: EfficientNet-B0  
+- **Task**: 4-way classification (Normal, Tuberculosis, Pneumonia, COVID-19)  
+- **Optimisation**: Sample-based Adaptive Sparse Training (AST)  
+- **Motivation**: Energy-efficient AI for global lung health  
+
+**Intended Use**
+
+- Research and prototyping  
+- Triage decision-support in pilot settings  
+- Education (medical students, residents, data scientists)  
+
+**Non-Intended Use**
+
+- Stand-alone diagnosis  
+- Automated treatment decisions  
+- Regulatory-grade clinical deployment  
+
+> Always pair the model with local guidelines, expert radiology, and laboratory testing.
+"""
+                    )
+
+    gr.Markdown("---")
+
+    gr.HTML(
+        """
+        <footer>
+            <p>
+                <b>AST Chest X-Ray Lab</b> · Normal · TB · Pneumonia · COVID-19 · Explainable AI<br/>
+                Built to explore how energy-efficient AI can support clinicians and patients worldwide.
+            </p>
+            <p style="margin-top:6px;">
+                ⚠️ <b>MEDICAL DISCLAIMER:</b> This tool is a research prototype and is not FDA-approved,
+                CE-marked, or licensed as a medical device. All clinical decisions must be made by
+                qualified healthcare professionals.
+            </p>
+        </footer>
+        """
+    )
+
+    # Wiring – connect logic to UI
+    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,
+            prob_chart_output,
+        ],
+    )
+
+    clear_btn.click(
+        fn=lambda: (
+            {},
+            None,
+            None,
+            None,
+            "Awaiting image upload…",
+            "Awaiting image upload…",
+            None,
+        ),
+        inputs=None,
+        outputs=[
+            prob_output,
+            original_output,
+            gradcam_output,
+            overlay_output,
+            interpretation_output,
+            snapshot_output,
+            prob_chart_output,
+        ],
+    )
+
+    # Example X-rays (optional – comment out if these paths don't exist)
+    gr.Markdown("### 🔍 Try Example X-rays")
+    gr.Examples(
+        examples=[
+            ["examples/normal.png"],
+            ["examples/tb.png"],
+            ["examples/pneumonia.png"],
+            ["examples/covid.png"],
+        ],
+        inputs=image_input,
+    )
+
+# ============================================================================
+# Launch
+# ============================================================================
+
+if __name__ == "__main__":
+    demo.launch(
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+    )