| """
|
| app/gradio_app.py
|
| ------------------
|
| Gradio demo dashboard for ChestViT — Explainable Chest X-Ray Analysis.
|
|
|
| Features:
|
| ┌────────────────────────────────────────────────────────────────┐
|
| │ Upload X-Ray │ Attention Rollout Heatmap Overlay │
|
| ├────────────────────────────────────────────────────────────────┤
|
| │ CLAHE Preview │ Disease Probability Bar Chart (14 classes) │
|
| └────────────────────────────────────────────────────────────────┘
|
|
|
| + Top diagnoses summary text
|
| + Model info sidebar
|
| + MLflow metrics link
|
|
|
| Run locally:
|
| python app/gradio_app.py
|
|
|
| Requirements:
|
| - Trained model checkpoint at checkpoints/best_model.pt
|
| - OR set DEMO_MODE=1 to run with random weights for UI preview
|
| """
|
|
|
| import os
|
| import sys
|
| import time
|
| from pathlib import Path
|
| from huggingface_hub import hf_hub_download
|
| import torch
|
|
|
|
|
| if sys.platform == "win32":
|
| sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
| sys.stderr.reconfigure(encoding="utf-8", errors="replace")
|
|
|
| import numpy as np
|
| import torch
|
| import gradio as gr
|
| from explainability.gradcam import generate_gradcam_heatmap
|
| import numpy as np
|
| import cv2
|
| import matplotlib
|
| matplotlib.use("Agg")
|
| import matplotlib.pyplot as plt
|
| import matplotlib.patches as mpatches
|
| from PIL import Image
|
|
|
|
|
|
|
| from config_loader import load_config
|
| from data.preprocessing import load_and_preprocess_raw, get_val_transforms, denormalize, apply_clahe
|
| from models.vit_model import ChestViT, load_checkpoint
|
| from explainability.attention_rollout import explain_prediction, rollout_to_heatmap
|
| from data.dataset import DISEASE_LABELS
|
|
|
|
|
| cfg = load_config()
|
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
|
|
|
| CUSTOM_CSS = """
|
| :root {
|
| --primary: #6366f1;
|
| --primary-dark: #4f46e5;
|
| --surface: #1e1e2e;
|
| --surface-2: #2a2a3e;
|
| --text: #e2e8f0;
|
| --text-muted: #94a3b8;
|
| --accent: #22d3ee;
|
| --danger: #ef4444;
|
| --warning: #f97316;
|
| --success: #22c55e;
|
| --border: rgba(99, 102, 241, 0.25);
|
| }
|
|
|
| body, .gradio-container {
|
| background: #0f0f1a !important;
|
| font-family: 'Inter', 'Segoe UI', sans-serif;
|
| }
|
|
|
| .gr-form, .gr-panel {
|
| background: var(--surface) !important;
|
| border: 1px solid var(--border) !important;
|
| border-radius: 16px !important;
|
| }
|
|
|
| .gr-button-primary {
|
| background: linear-gradient(135deg, var(--primary), var(--primary-dark)) !important;
|
| border: none !important;
|
| border-radius: 10px !important;
|
| font-weight: 600 !important;
|
| letter-spacing: 0.5px !important;
|
| box-shadow: 0 4px 15px rgba(99, 102, 241, 0.4) !important;
|
| transition: all 0.2s ease !important;
|
| }
|
| .gr-button-primary:hover {
|
| transform: translateY(-1px) !important;
|
| box-shadow: 0 6px 20px rgba(99, 102, 241, 0.6) !important;
|
| }
|
|
|
| label, .label-wrap span {
|
| color: var(--text-muted) !important;
|
| font-size: 0.85rem !important;
|
| font-weight: 500 !important;
|
| text-transform: uppercase !important;
|
| letter-spacing: 0.5px !important;
|
| }
|
|
|
| h1, h2, h3 { color: var(--text) !important; }
|
|
|
| .header-title {
|
| font-size: 2.2rem;
|
| font-weight: 800;
|
| background: linear-gradient(135deg, #6366f1, #22d3ee, #22c55e);
|
| -webkit-background-clip: text;
|
| -webkit-text-fill-color: transparent;
|
| text-align: center;
|
| margin-bottom: 0.5rem;
|
| }
|
| .header-sub {
|
| color: var(--text-muted);
|
| text-align: center;
|
| font-size: 0.95rem;
|
| margin-bottom: 1.5rem;
|
| }
|
| .stat-card {
|
| background: var(--surface-2);
|
| border: 1px solid var(--border);
|
| border-radius: 12px;
|
| padding: 12px 16px;
|
| margin: 4px;
|
| text-align: center;
|
| }
|
| """
|
|
|
| HEADER_HTML = """
|
| <div style="text-align:center; padding: 20px 0 10px 0;">
|
| <div style="font-size:2.4rem; font-weight:800; background:linear-gradient(135deg,#6366f1,#22d3ee,#22c55e);
|
| -webkit-background-clip:text; -webkit-text-fill-color:transparent;">
|
| 🫁 ChestViT — Explainable X-Ray AI
|
| </div>
|
| <div style="color:#94a3b8; font-size:0.95rem; margin-top:8px;">
|
| ViT-Base-16 · 14-Disease Multi-Label Classification · Attention Rollout Explainability
|
| </div>
|
| <div style="display:flex; justify-content:center; gap:16px; margin-top:14px; flex-wrap:wrap;">
|
| <span style="background:#1e1e2e; border:1px solid rgba(99,102,241,0.3); border-radius:8px;
|
| padding:6px 14px; color:#a5b4fc; font-size:0.82rem; font-weight:600;">
|
| 🤖 google/vit-base-patch16-224-in21k
|
| </span>
|
| <span style="background:#1e1e2e; border:1px solid rgba(34,211,238,0.3); border-radius:8px;
|
| padding:6px 14px; color:#67e8f9; font-size:0.82rem; font-weight:600;">
|
| 📊 NIH ChestX-ray14 Dataset
|
| </span>
|
| <span style="background:#1e1e2e; border:1px solid rgba(34,197,94,0.3); border-radius:8px;
|
| padding:6px 14px; color:#86efac; font-size:0.82rem; font-weight:600;">
|
| 🔥 Attention Rollout XAI
|
| </span>
|
| </div>
|
| </div>
|
| """
|
|
|
| FOOTER_HTML = """
|
| <div style="text-align:center; color:#475569; font-size:0.8rem; padding:16px 0 8px 0; border-top:1px solid rgba(99,102,241,0.15); margin-top:16px;">
|
| ⚠️ <strong>Research / Educational Use Only.</strong>
|
| This tool is NOT a medical device and should NOT be used for clinical diagnosis.
|
| Always consult a qualified radiologist.
|
| </div>
|
| """
|
|
|
|
|
|
|
|
|
| def load_model() -> ChestViT:
|
| """Load trained model from Hugging Face Model Hub."""
|
|
|
| ckpt_path = hf_hub_download(
|
| repo_id="sandy45/ChestViT-ViTBase-NIH14",
|
| filename="best_model.pt"
|
| )
|
|
|
| model = load_checkpoint(ckpt_path, DEVICE)
|
|
|
| model.to(DEVICE)
|
| model.eval()
|
|
|
| print("🔥 Loaded trained model from Hugging Face Hub")
|
|
|
| return model
|
|
|
|
|
|
|
| print(f"\nLoading model on {DEVICE}...")
|
| MODEL = load_model()
|
| VAL_TRANSFORM = get_val_transforms(cfg.dataset.image_size)
|
| print("Model ready.\n")
|
|
|
|
|
|
|
|
|
| def preprocess_uploaded_image(pil_image: Image.Image) -> tuple[np.ndarray, torch.Tensor]:
|
| """
|
| Convert a PIL image (from Gradio upload) to:
|
| 1. CLAHE-enhanced numpy array for display
|
| 2. Normalized tensor for model input
|
|
|
| Returns:
|
| (clahe_rgb, input_tensor) where:
|
| clahe_rgb: (H, W, 3) uint8 numpy array
|
| input_tensor: (1, 3, 224, 224) float32 tensor
|
| """
|
|
|
| img_np = np.array(pil_image.convert("RGB"))
|
|
|
|
|
| gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
|
| if gray.dtype == np.uint16:
|
| gray = (gray / 256).astype(np.uint8)
|
| clahe_gray = apply_clahe(gray, clip_limit=2.0, tile_size=8)
|
| clahe_gray = cv2.resize(clahe_gray, (224, 224), interpolation=cv2.INTER_AREA)
|
| clahe_rgb = cv2.cvtColor(clahe_gray, cv2.COLOR_GRAY2RGB)
|
|
|
|
|
| augmented = VAL_TRANSFORM(image=clahe_rgb)
|
| tensor = augmented["image"].unsqueeze(0)
|
|
|
| return clahe_rgb, tensor
|
|
|
|
|
| def make_probability_figure(probs: np.ndarray, threshold: float = 0.5) -> plt.Figure:
|
| """
|
| Create a beautiful dark-themed horizontal bar chart of disease probabilities.
|
| """
|
| sorted_idx = np.argsort(probs)
|
| sorted_probs = probs[sorted_idx]
|
| sorted_names = [DISEASE_LABELS[i] for i in sorted_idx]
|
|
|
| fig, ax = plt.subplots(figsize=(7, 6))
|
| fig.patch.set_facecolor("#0f0f1a")
|
| ax.set_facecolor("#1a1a2e")
|
|
|
|
|
| colors = []
|
| for p in sorted_probs:
|
| if p >= 0.7:
|
| colors.append("#ef4444")
|
| elif p >= 0.5:
|
| colors.append("#f97316")
|
| elif p >= 0.3:
|
| colors.append("#eab308")
|
| else:
|
| colors.append("#3b82f6")
|
|
|
| bars = ax.barh(range(len(sorted_names)), sorted_probs, color=colors,
|
| edgecolor="none", height=0.65)
|
|
|
|
|
| ax.axvline(x=threshold, color="#a855f7", linestyle="--", linewidth=1.5,
|
| alpha=0.8, label=f"Threshold ({threshold})")
|
|
|
|
|
| ax.set_yticks(range(len(sorted_names)))
|
| ax.set_yticklabels(sorted_names, color="#e2e8f0", fontsize=9.5)
|
| ax.set_xlabel("Probability", color="#94a3b8", fontsize=10)
|
| ax.set_title("Disease Probability Scores", color="white",
|
| fontsize=12, fontweight="bold", pad=12)
|
| ax.set_xlim(0, 1.0)
|
| ax.tick_params(axis="x", colors="#94a3b8", labelsize=9)
|
|
|
|
|
| for bar, p in zip(bars, sorted_probs):
|
| ax.text(min(p + 0.02, 0.95), bar.get_y() + bar.get_height() / 2,
|
| f"{p:.3f}", va="center", color="white", fontsize=8.5, fontweight="bold")
|
|
|
|
|
| patches = [
|
| mpatches.Patch(color="#ef4444", label="High confidence (≥0.7)"),
|
| mpatches.Patch(color="#f97316", label="Positive (≥0.5)"),
|
| mpatches.Patch(color="#eab308", label="Uncertain (0.3–0.5)"),
|
| mpatches.Patch(color="#3b82f6", label="Likely negative (<0.3)"),
|
| ]
|
| ax.legend(handles=patches, loc="lower right", fontsize=7.5,
|
| facecolor="#0f0f1a", labelcolor="white", framealpha=0.8)
|
|
|
| for spine in ax.spines.values():
|
| spine.set_edgecolor("#2d2d4e")
|
|
|
| plt.tight_layout()
|
| return fig
|
|
|
|
|
| def make_heatmap_figure(
|
| clahe_rgb: np.ndarray,
|
| rollout: np.ndarray,
|
| overlay: np.ndarray,
|
| ) -> plt.Figure:
|
| """
|
| 3-panel figure: original | raw rollout | overlay.
|
| """
|
| fig, axes = plt.subplots(1, 3, figsize=(12, 4.5))
|
| fig.patch.set_facecolor("#0f0f1a")
|
|
|
| titles = ["CLAHE-Enhanced X-Ray", "Attention Rollout Map", "Heatmap Overlay"]
|
| for ax, title in zip(axes, titles):
|
| ax.set_facecolor("#1a1a2e")
|
| ax.set_title(title, color="white", fontsize=10, fontweight="bold", pad=8)
|
| ax.axis("off")
|
|
|
| axes[0].imshow(clahe_rgb)
|
| axes[0].text(5, 218, "Input", color="#94a3b8", fontsize=8,
|
| va="bottom", ha="left", fontweight="bold")
|
|
|
| rollout_display = cv2.resize(rollout, (224, 224), interpolation=cv2.INTER_CUBIC)
|
| im = axes[1].imshow(rollout_display, cmap="inferno", vmin=0, vmax=1)
|
| plt.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04,
|
| label="Attention Weight")
|
|
|
| axes[2].imshow(overlay)
|
| axes[2].text(5, 218, "ViT Attention Rollout", color="white",
|
| fontsize=7.5, va="bottom", ha="left",
|
| bbox=dict(boxstyle="round,pad=2", facecolor="#0f0f1a", alpha=0.7))
|
|
|
| plt.tight_layout(pad=1.5)
|
| return fig
|
|
|
|
|
| def analyze_xray(
|
| pil_image: Image.Image,
|
| head_fusion: str,
|
| discard_ratio: float,
|
| threshold: float,
|
| explainability_method: str = "Attention Rollout",
|
| target_disease: str = "None (Highest Score)",
|
| ) -> tuple:
|
| """
|
| Main inference function called by Gradio.
|
|
|
| Returns:
|
| (heatmap_figure, prob_figure, diagnosis_text, status_text)
|
| """
|
| if pil_image is None:
|
| return None, None, "⬆ Please upload a chest X-ray image.", ""
|
|
|
| start_time = time.time()
|
|
|
| try:
|
|
|
| clahe_rgb, input_tensor = preprocess_uploaded_image(pil_image)
|
|
|
|
|
| model_was_training = MODEL.training
|
| MODEL.eval()
|
|
|
| with torch.no_grad():
|
| logits, attentions = MODEL(input_tensor.to(DEVICE), output_attentions=True)
|
| probs = torch.sigmoid(logits).squeeze().cpu().numpy()
|
|
|
| target_idx = np.argmax(probs)
|
| if target_disease != "None (Highest Score)" and target_disease in DISEASE_LABELS:
|
| target_idx = DISEASE_LABELS.index(target_disease)
|
|
|
| if explainability_method == "Attention Rollout":
|
| with torch.no_grad():
|
| _, rollout, overlay = explain_prediction(
|
| model=MODEL,
|
| image_tensor=input_tensor,
|
| original_image=clahe_rgb,
|
| device=DEVICE,
|
| disease_names=DISEASE_LABELS,
|
| head_fusion=head_fusion,
|
| discard_ratio=discard_ratio,
|
| )
|
| else:
|
| MODEL.zero_grad()
|
| overlay = generate_gradcam_heatmap(
|
| model=MODEL,
|
| image_tensor=input_tensor.to(DEVICE),
|
| target_class=target_idx,
|
| original_image=clahe_rgb
|
| )
|
|
|
| rollout = np.zeros((14, 14))
|
|
|
| if model_was_training:
|
| MODEL.train()
|
|
|
|
|
| elapsed = time.time() - start_time
|
|
|
|
|
| heatmap_fig = make_heatmap_figure(clahe_rgb, rollout, overlay)
|
|
|
|
|
| prob_fig = make_probability_figure(probs, threshold=threshold)
|
|
|
|
|
| positives = [
|
| (DISEASE_LABELS[i], probs[i])
|
| for i in range(14)
|
| if probs[i] >= threshold
|
| ]
|
| positives.sort(key=lambda x: x[1], reverse=True)
|
|
|
| if positives:
|
| diag_lines = [f"### 🔴 Detected Findings (confidence ≥ {threshold:.0%})"]
|
| for disease, prob in positives:
|
| bar = "█" * int(prob * 20) + "░" * (20 - int(prob * 20))
|
| diag_lines.append(f"**{disease}**: {bar} `{prob:.1%}`")
|
| else:
|
| diag_lines = [
|
| f"### 🟢 No Findings Detected",
|
| f"All disease probabilities below threshold ({threshold:.0%}).",
|
| "This may indicate a normal chest X-ray.",
|
| ]
|
|
|
| diag_lines.append(f"\n---\n*Inference time: {elapsed:.2f}s · Device: {DEVICE}*")
|
| diag_text = "\n\n".join(diag_lines)
|
|
|
| status = (
|
| f"✅ Analysis complete in {elapsed:.2f}s | "
|
| f"Device: {str(DEVICE).upper()} | "
|
| f"{"🔥 Trained ChestViT • AUROC 0.789"}"
|
| )
|
|
|
| return heatmap_fig, prob_fig, diag_text, status
|
|
|
| except Exception as e:
|
| import traceback
|
| err = traceback.format_exc()
|
| return None, None, f"❌ Error during analysis:\n```\n{err}\n```", "Error"
|
|
|
|
|
|
|
|
|
| def build_interface() -> gr.Blocks:
|
| with gr.Blocks(
|
| title="ChestViT -- Explainable Chest X-Ray AI",
|
| theme=gr.themes.Base(
|
| primary_hue="indigo",
|
| secondary_hue="cyan",
|
| neutral_hue="slate",
|
| font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif"],
|
| ),
|
| css=CUSTOM_CSS,
|
| ) as demo:
|
|
|
|
|
| gr.HTML(HEADER_HTML)
|
|
|
| with gr.Row():
|
|
|
| with gr.Column(scale=1, min_width=300):
|
| gr.Markdown("### 📤 Upload Chest X-Ray")
|
| image_input = gr.Image(
|
| type="pil",
|
| label="Chest X-Ray (PNG/JPEG/DICOM-exported PNG)",
|
| height=280,
|
| sources=["upload", "clipboard"],
|
| )
|
|
|
| gr.Markdown("### ⚙️ Explainability Settings")
|
| with gr.Group():
|
| explainability_method = gr.Radio(
|
| choices=["Attention Rollout", "Grad-CAM"],
|
| value="Attention Rollout",
|
| label="Explainability Method",
|
| info="Choose between Transformer-native Attention Rollout or Grad-CAM"
|
| )
|
| target_disease = gr.Dropdown(
|
| choices=["None (Highest Score)"] + DISEASE_LABELS,
|
| value="None (Highest Score)",
|
| label="Target Disease for Grad-CAM",
|
| info="Forces Grad-CAM to explain this specific disease"
|
| )
|
| head_fusion = gr.Radio(
|
| choices=["mean", "max", "min"],
|
| value="mean",
|
| label="Attention Head Fusion",
|
| info="[Attention Rollout] How to combine 12 attention heads into one map",
|
| )
|
| discard_ratio = gr.Slider(
|
| minimum=0.0, maximum=0.99, value=0.9, step=0.05,
|
| label="Low-Attention Discard Ratio",
|
| info="[Attention Rollout] Zeroes out lowest-attention patches (noise reduction)",
|
| )
|
| threshold = gr.Slider(
|
| minimum=0.1, maximum=0.9, value=0.5, step=0.05,
|
| label="Prediction Threshold",
|
| info="Sigmoid probability cutoff for positive prediction",
|
| )
|
|
|
| analyze_btn = gr.Button(
|
| "🔬 Analyze X-Ray",
|
| variant="primary",
|
| size="lg",
|
| )
|
| status_text = gr.Textbox(
|
| label="Status",
|
| interactive=False,
|
| show_label=True,
|
| max_lines=2,
|
| )
|
|
|
|
|
| gr.Markdown(
|
| """
|
| > **💡 Tips**
|
| > - Use frontal (PA or AP) chest X-ray images
|
| > - PNG or JPEG format accepted
|
| > - Best results with 1024×1024 pixel images
|
| > - Works with exported DICOM screenshots
|
| """
|
| )
|
|
|
|
|
| with gr.Column(scale=2, min_width=600):
|
| gr.Markdown("### 🔥 Attention Rollout Visualization")
|
| heatmap_output = gr.Plot(
|
| label="Attention Rollout Analysis",
|
| show_label=False,
|
| )
|
|
|
| gr.Markdown("### 📊 Disease Probability Scores")
|
| prob_output = gr.Plot(
|
| label="Disease Probabilities",
|
| show_label=False,
|
| )
|
|
|
| gr.Markdown("### 🩺 Diagnosis Summary")
|
| diagnosis_output = gr.Markdown(
|
| value="*Upload an X-ray and click Analyze to see results.*"
|
| )
|
|
|
|
|
| with gr.Accordion("📖 How It Works", open=False):
|
| gr.Markdown("""
|
| ## Architecture
|
|
|
| | Component | Details |
|
| |---|---|
|
| | **Model** | ViT-Base-16 (google/vit-base-patch16-224-in21k) |
|
| | **Pre-training** | ImageNet-21k (14M images, 21K classes) |
|
| | **Fine-tuning** | NIH ChestX-ray14 (112,120 frontal X-rays) |
|
| | **Task** | Multi-label classification — 14 simultaneous disease predictions |
|
| | **Loss** | Weighted Binary Cross-Entropy (handles severe class imbalance) |
|
| | **Preprocessing** | CLAHE contrast enhancement → Albumentations augmentation |
|
| | **Explainability** | Attention Rollout (Abnar & Zuidema, 2020) |
|
|
|
| ## Attention Rollout Algorithm
|
|
|
| Standard Grad-CAM doesn't work well with pure Vision Transformers because
|
| ViTs don't have intermediate spatial feature maps like CNNs.
|
|
|
| **Attention Rollout** instead:
|
| 1. Extracts raw attention weights from all 12 transformer layers
|
| 2. Averages across all 12 attention heads per layer
|
| 3. Adds an identity matrix (modeling residual/skip connections)
|
| 4. Re-normalizes each row
|
| 5. Multiplies all 12 matrices in sequence → propagates attention end-to-end
|
| 6. Reads the `[CLS]` token row → shows which 14×14 patches it attends to
|
| 7. Upsamples 14×14 → 224×224 and overlays as a heatmap
|
|
|
| This shows **where in the X-ray the model is looking** when it makes each prediction.
|
|
|
| ## NIH Dataset — 14 Disease Labels
|
|
|
| ```
|
| Atelectasis · Cardiomegaly · Effusion · Infiltration · Mass · Nodule
|
| Pneumonia · Pneumothorax · Consolidation · Edema · Emphysema
|
| Fibrosis · Pleural_Thickening · Hernia
|
| ```
|
|
|
| ## References
|
|
|
| - Wang et al. (2017). *ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks*. CVPR.
|
| - Dosovitskiy et al. (2021). *An Image is Worth 16x16 Words*. ICLR.
|
| - Abnar & Zuidema (2020). *Quantifying Attention Flow in Transformers*. arXiv:2005.00928.
|
| """)
|
|
|
| gr.HTML(FOOTER_HTML)
|
|
|
|
|
| analyze_btn.click(
|
| fn=analyze_xray,
|
| inputs=[image_input, head_fusion, discard_ratio, threshold, explainability_method, target_disease],
|
| outputs=[heatmap_output, prob_output, diagnosis_output, status_text],
|
| api_name="analyze",
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| return demo
|
|
|
|
|
|
|
|
|
| if __name__ == "__main__":
|
| demo = build_interface()
|
| demo.launch(
|
| server_port=cfg.inference.gradio_port,
|
| share=cfg.inference.gradio_share,
|
| show_error=True,
|
| inbrowser=True,
|
| )
|
|
|
