Delete app.py

app.py

@@ -1,875 +0,0 @@
-"""
-🫁 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,
-    )