sandy45/ChestViT-Explainable-XRay-AI
0
1"""
2app/gradio_app.py
3------------------
4Gradio demo dashboard for ChestViT — Explainable Chest X-Ray Analysis.
5
6Features:
7 ┌────────────────────────────────────────────────────────────────┐
8 │ Upload X-Ray │ Attention Rollout Heatmap Overlay │
9 ├────────────────────────────────────────────────────────────────┤
10 │ CLAHE Preview │ Disease Probability Bar Chart (14 classes) │
11 └────────────────────────────────────────────────────────────────┘
12
13 + Top diagnoses summary text
14 + Model info sidebar
15 + MLflow metrics link
16
17Run locally:
18 python app/gradio_app.py
19
20Requirements:
21 - Trained model checkpoint at checkpoints/best_model.pt
22 - OR set DEMO_MODE=1 to run with random weights for UI preview
23"""
24
25import os
26import sys
27import time
28from pathlib import Path
29from huggingface_hub import hf_hub_download
30import torch
31
32# Fix Windows console encoding (cp1252 can't handle emoji/Unicode)
33if sys.platform == "win32":
34 sys.stdout.reconfigure(encoding="utf-8", errors="replace")
35 sys.stderr.reconfigure(encoding="utf-8", errors="replace")
36
37import numpy as np
38import torch
39import gradio as gr
40from explainability.gradcam import generate_gradcam_heatmap
41import numpy as np
42import cv2
43import matplotlib
44matplotlib.use("Agg") # Non-interactive backend for server use
45import matplotlib.pyplot as plt
46import matplotlib.patches as mpatches
47from PIL import Image
48
49
50
51from config_loader import load_config
52from data.preprocessing import load_and_preprocess_raw, get_val_transforms, denormalize, apply_clahe
53from models.vit_model import ChestViT, load_checkpoint
54from explainability.attention_rollout import explain_prediction, rollout_to_heatmap
55from data.dataset import DISEASE_LABELS
56
57# ── Config ────────────────────────────────────────────────────────────────────
58cfg = load_config()
59DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
60
61
62
63# ── CSS Styling ───────────────────────────────────────────────────────────────
64CUSTOM_CSS = """
65:root {
66 --primary: #6366f1;
67 --primary-dark: #4f46e5;
68 --surface: #1e1e2e;
69 --surface-2: #2a2a3e;
70 --text: #e2e8f0;
71 --text-muted: #94a3b8;
72 --accent: #22d3ee;
73 --danger: #ef4444;
74 --warning: #f97316;
75 --success: #22c55e;
76 --border: rgba(99, 102, 241, 0.25);
77}
78
79body, .gradio-container {
80 background: #0f0f1a !important;
81 font-family: 'Inter', 'Segoe UI', sans-serif;
82}
83
84.gr-form, .gr-panel {
85 background: var(--surface) !important;
86 border: 1px solid var(--border) !important;
87 border-radius: 16px !important;
88}
89
90.gr-button-primary {
91 background: linear-gradient(135deg, var(--primary), var(--primary-dark)) !important;
92 border: none !important;
93 border-radius: 10px !important;
94 font-weight: 600 !important;
95 letter-spacing: 0.5px !important;
96 box-shadow: 0 4px 15px rgba(99, 102, 241, 0.4) !important;
97 transition: all 0.2s ease !important;
98}
99.gr-button-primary:hover {
100 transform: translateY(-1px) !important;
101 box-shadow: 0 6px 20px rgba(99, 102, 241, 0.6) !important;
102}
103
104label, .label-wrap span {
105 color: var(--text-muted) !important;
106 font-size: 0.85rem !important;
107 font-weight: 500 !important;
108 text-transform: uppercase !important;
109 letter-spacing: 0.5px !important;
110}
111
112h1, h2, h3 { color: var(--text) !important; }
113
114.header-title {
115 font-size: 2.2rem;
116 font-weight: 800;
117 background: linear-gradient(135deg, #6366f1, #22d3ee, #22c55e);
118 -webkit-background-clip: text;
119 -webkit-text-fill-color: transparent;
120 text-align: center;
121 margin-bottom: 0.5rem;
122}
123.header-sub {
124 color: var(--text-muted);
125 text-align: center;
126 font-size: 0.95rem;
127 margin-bottom: 1.5rem;
128}
129.stat-card {
130 background: var(--surface-2);
131 border: 1px solid var(--border);
132 border-radius: 12px;
133 padding: 12px 16px;
134 margin: 4px;
135 text-align: center;
136}
137"""
138
139HEADER_HTML = """
140<div style="text-align:center; padding: 20px 0 10px 0;">
141 <div style="font-size:2.4rem; font-weight:800; background:linear-gradient(135deg,#6366f1,#22d3ee,#22c55e);
142 -webkit-background-clip:text; -webkit-text-fill-color:transparent;">
143 🫁 ChestViT — Explainable X-Ray AI
144 </div>
145 <div style="color:#94a3b8; font-size:0.95rem; margin-top:8px;">
146 ViT-Base-16 · 14-Disease Multi-Label Classification · Attention Rollout Explainability
147 </div>
148 <div style="display:flex; justify-content:center; gap:16px; margin-top:14px; flex-wrap:wrap;">
149 <span style="background:#1e1e2e; border:1px solid rgba(99,102,241,0.3); border-radius:8px;
150 padding:6px 14px; color:#a5b4fc; font-size:0.82rem; font-weight:600;">
151 🤖 google/vit-base-patch16-224-in21k
152 </span>
153 <span style="background:#1e1e2e; border:1px solid rgba(34,211,238,0.3); border-radius:8px;
154 padding:6px 14px; color:#67e8f9; font-size:0.82rem; font-weight:600;">
155 📊 NIH ChestX-ray14 Dataset
156 </span>
157 <span style="background:#1e1e2e; border:1px solid rgba(34,197,94,0.3); border-radius:8px;
158 padding:6px 14px; color:#86efac; font-size:0.82rem; font-weight:600;">
159 🔥 Attention Rollout XAI
160 </span>
161 </div>
162</div>
163"""
164
165FOOTER_HTML = """
166<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;">
167 ⚠️ <strong>Research / Educational Use Only.</strong>
168 This tool is NOT a medical device and should NOT be used for clinical diagnosis.
169 Always consult a qualified radiologist.
170</div>
171"""
172
173
174# ── Model Loading ─────────────────────────────────────────────────────────────
175
176def load_model() -> ChestViT:
177 """Load trained model from Hugging Face Model Hub."""
178
179 ckpt_path = hf_hub_download(
180 repo_id="sandy45/ChestViT-ViTBase-NIH14",
181 filename="best_model.pt"
182 )
183
184 model = load_checkpoint(ckpt_path, DEVICE)
185
186 model.to(DEVICE)
187 model.eval()
188
189 print("🔥 Loaded trained model from Hugging Face Hub")
190
191 return model
192
193
194# Load model once at startup
195print(f"\nLoading model on {DEVICE}...")
196MODEL = load_model()
197VAL_TRANSFORM = get_val_transforms(cfg.dataset.image_size)
198print("Model ready.\n")
199
200
201# ── Inference Pipeline ────────────────────────────────────────────────────────
202
203def preprocess_uploaded_image(pil_image: Image.Image) -> tuple[np.ndarray, torch.Tensor]:
204 """
205 Convert a PIL image (from Gradio upload) to:
206 1. CLAHE-enhanced numpy array for display
207 2. Normalized tensor for model input
208
209 Returns:
210 (clahe_rgb, input_tensor) where:
211 clahe_rgb: (H, W, 3) uint8 numpy array
212 input_tensor: (1, 3, 224, 224) float32 tensor
213 """
214 # Convert to numpy
215 img_np = np.array(pil_image.convert("RGB"))
216
217 # To grayscale → CLAHE → back to RGB
218 gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
219 if gray.dtype == np.uint16:
220 gray = (gray / 256).astype(np.uint8)
221 clahe_gray = apply_clahe(gray, clip_limit=2.0, tile_size=8)
222 clahe_gray = cv2.resize(clahe_gray, (224, 224), interpolation=cv2.INTER_AREA)
223 clahe_rgb = cv2.cvtColor(clahe_gray, cv2.COLOR_GRAY2RGB)
224
225 # Normalize for model
226 augmented = VAL_TRANSFORM(image=clahe_rgb)
227 tensor = augmented["image"].unsqueeze(0) # (1, 3, 224, 224)
228
229 return clahe_rgb, tensor
230
231
232def make_probability_figure(probs: np.ndarray, threshold: float = 0.5) -> plt.Figure:
233 """
234 Create a beautiful dark-themed horizontal bar chart of disease probabilities.
235 """
236 sorted_idx = np.argsort(probs) # ascending for bottom-to-top barh
237 sorted_probs = probs[sorted_idx]
238 sorted_names = [DISEASE_LABELS[i] for i in sorted_idx]
239
240 fig, ax = plt.subplots(figsize=(7, 6))
241 fig.patch.set_facecolor("#0f0f1a")
242 ax.set_facecolor("#1a1a2e")
243
244 # Color code by probability
245 colors = []
246 for p in sorted_probs:
247 if p >= 0.7:
248 colors.append("#ef4444") # Red — high confidence positive
249 elif p >= 0.5:
250 colors.append("#f97316") # Orange — positive
251 elif p >= 0.3:
252 colors.append("#eab308") # Yellow — uncertain
253 else:
254 colors.append("#3b82f6") # Blue — likely negative
255
256 bars = ax.barh(range(len(sorted_names)), sorted_probs, color=colors,
257 edgecolor="none", height=0.65)
258
259 # Threshold line
260 ax.axvline(x=threshold, color="#a855f7", linestyle="--", linewidth=1.5,
261 alpha=0.8, label=f"Threshold ({threshold})")
262
263 # Labels
264 ax.set_yticks(range(len(sorted_names)))
265 ax.set_yticklabels(sorted_names, color="#e2e8f0", fontsize=9.5)
266 ax.set_xlabel("Probability", color="#94a3b8", fontsize=10)
267 ax.set_title("Disease Probability Scores", color="white",
268 fontsize=12, fontweight="bold", pad=12)
269 ax.set_xlim(0, 1.0)
270 ax.tick_params(axis="x", colors="#94a3b8", labelsize=9)
271
272 # Value labels on bars
273 for bar, p in zip(bars, sorted_probs):
274 ax.text(min(p + 0.02, 0.95), bar.get_y() + bar.get_height() / 2,
275 f"{p:.3f}", va="center", color="white", fontsize=8.5, fontweight="bold")
276
277 # Legend
278 patches = [
279 mpatches.Patch(color="#ef4444", label="High confidence (≥0.7)"),
280 mpatches.Patch(color="#f97316", label="Positive (≥0.5)"),
281 mpatches.Patch(color="#eab308", label="Uncertain (0.3–0.5)"),
282 mpatches.Patch(color="#3b82f6", label="Likely negative (<0.3)"),
283 ]
284 ax.legend(handles=patches, loc="lower right", fontsize=7.5,
285 facecolor="#0f0f1a", labelcolor="white", framealpha=0.8)
286
287 for spine in ax.spines.values():
288 spine.set_edgecolor("#2d2d4e")
289
290 plt.tight_layout()
291 return fig
292
293
294def make_heatmap_figure(
295 clahe_rgb: np.ndarray,
296 rollout: np.ndarray,
297 overlay: np.ndarray,
298) -> plt.Figure:
299 """
300 3-panel figure: original | raw rollout | overlay.
301 """
302 fig, axes = plt.subplots(1, 3, figsize=(12, 4.5))
303 fig.patch.set_facecolor("#0f0f1a")
304
305 titles = ["CLAHE-Enhanced X-Ray", "Attention Rollout Map", "Heatmap Overlay"]
306 for ax, title in zip(axes, titles):
307 ax.set_facecolor("#1a1a2e")
308 ax.set_title(title, color="white", fontsize=10, fontweight="bold", pad=8)
309 ax.axis("off")
310
311 axes[0].imshow(clahe_rgb)
312 axes[0].text(5, 218, "Input", color="#94a3b8", fontsize=8,
313 va="bottom", ha="left", fontweight="bold")
314
315 rollout_display = cv2.resize(rollout, (224, 224), interpolation=cv2.INTER_CUBIC)
316 im = axes[1].imshow(rollout_display, cmap="inferno", vmin=0, vmax=1)
317 plt.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04,
318 label="Attention Weight")
319
320 axes[2].imshow(overlay)
321 axes[2].text(5, 218, "ViT Attention Rollout", color="white",
322 fontsize=7.5, va="bottom", ha="left",
323 bbox=dict(boxstyle="round,pad=2", facecolor="#0f0f1a", alpha=0.7))
324
325 plt.tight_layout(pad=1.5)
326 return fig
327
328
329def analyze_xray(
330 pil_image: Image.Image,
331 head_fusion: str,
332 discard_ratio: float,
333 threshold: float,
334 explainability_method: str = "Attention Rollout",
335 target_disease: str = "None (Highest Score)",
336) -> tuple:
337 """
338 Main inference function called by Gradio.
339
340 Returns:
341 (heatmap_figure, prob_figure, diagnosis_text, status_text)
342 """
343 if pil_image is None:
344 return None, None, "⬆ Please upload a chest X-ray image.", ""
345
346 start_time = time.time()
347
348 try:
349 # Preprocess
350 clahe_rgb, input_tensor = preprocess_uploaded_image(pil_image)
351
352 # Inference
353 model_was_training = MODEL.training
354 MODEL.eval()
355
356 with torch.no_grad():
357 logits, attentions = MODEL(input_tensor.to(DEVICE), output_attentions=True)
358 probs = torch.sigmoid(logits).squeeze().cpu().numpy()
359
360 target_idx = np.argmax(probs)
361 if target_disease != "None (Highest Score)" and target_disease in DISEASE_LABELS:
362 target_idx = DISEASE_LABELS.index(target_disease)
363
364 if explainability_method == "Attention Rollout":
365 with torch.no_grad():
366 _, rollout, overlay = explain_prediction(
367 model=MODEL,
368 image_tensor=input_tensor,
369 original_image=clahe_rgb,
370 device=DEVICE,
371 disease_names=DISEASE_LABELS,
372 head_fusion=head_fusion,
373 discard_ratio=discard_ratio,
374 )
375 else: # Grad-CAM
376 MODEL.zero_grad()
377 overlay = generate_gradcam_heatmap(
378 model=MODEL,
379 image_tensor=input_tensor.to(DEVICE),
380 target_class=target_idx,
381 original_image=clahe_rgb
382 )
383 # Create a dummy rollout to satisfy the function if Grad-CAM
384 rollout = np.zeros((14, 14))
385
386 if model_was_training:
387 MODEL.train()
388
389
390 elapsed = time.time() - start_time
391
392 # Build heatmap figure
393 heatmap_fig = make_heatmap_figure(clahe_rgb, rollout, overlay)
394
395 # Build probability figure
396 prob_fig = make_probability_figure(probs, threshold=threshold)
397
398 # Build diagnosis summary text
399 positives = [
400 (DISEASE_LABELS[i], probs[i])
401 for i in range(14)
402 if probs[i] >= threshold
403 ]
404 positives.sort(key=lambda x: x[1], reverse=True)
405
406 if positives:
407 diag_lines = [f"### 🔴 Detected Findings (confidence ≥ {threshold:.0%})"]
408 for disease, prob in positives:
409 bar = "█" * int(prob * 20) + "░" * (20 - int(prob * 20))
410 diag_lines.append(f"**{disease}**: {bar} `{prob:.1%}`")
411 else:
412 diag_lines = [
413 f"### 🟢 No Findings Detected",
414 f"All disease probabilities below threshold ({threshold:.0%}).",
415 "This may indicate a normal chest X-ray.",
416 ]
417
418 diag_lines.append(f"\n---\n*Inference time: {elapsed:.2f}s · Device: {DEVICE}*")
419 diag_text = "\n\n".join(diag_lines)
420
421 status = (
422 f"✅ Analysis complete in {elapsed:.2f}s | "
423 f"Device: {str(DEVICE).upper()} | "
424 f"{"🔥 Trained ChestViT • AUROC 0.789"}"
425 )
426
427 return heatmap_fig, prob_fig, diag_text, status
428
429 except Exception as e:
430 import traceback
431 err = traceback.format_exc()
432 return None, None, f"❌ Error during analysis:\n```\n{err}\n```", "Error"
433
434
435# ── Gradio Interface ──────────────────────────────────────────────────────────
436
437def build_interface() -> gr.Blocks:
438 with gr.Blocks(
439 title="ChestViT -- Explainable Chest X-Ray AI",
440 theme=gr.themes.Base(
441 primary_hue="indigo",
442 secondary_hue="cyan",
443 neutral_hue="slate",
444 font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif"],
445 ),
446 css=CUSTOM_CSS,
447 ) as demo:
448
449 # Header
450 gr.HTML(HEADER_HTML)
451
452 with gr.Row():
453 # ── Left Column: Input ───────────────────────────────────────────
454 with gr.Column(scale=1, min_width=300):
455 gr.Markdown("### 📤 Upload Chest X-Ray")
456 image_input = gr.Image(
457 type="pil",
458 label="Chest X-Ray (PNG/JPEG/DICOM-exported PNG)",
459 height=280,
460 sources=["upload", "clipboard"],
461 )
462
463 gr.Markdown("### ⚙️ Explainability Settings")
464 with gr.Group():
465 explainability_method = gr.Radio(
466 choices=["Attention Rollout", "Grad-CAM"],
467 value="Attention Rollout",
468 label="Explainability Method",
469 info="Choose between Transformer-native Attention Rollout or Grad-CAM"
470 )
471 target_disease = gr.Dropdown(
472 choices=["None (Highest Score)"] + DISEASE_LABELS,
473 value="None (Highest Score)",
474 label="Target Disease for Grad-CAM",
475 info="Forces Grad-CAM to explain this specific disease"
476 )
477 head_fusion = gr.Radio(
478 choices=["mean", "max", "min"],
479 value="mean",
480 label="Attention Head Fusion",
481 info="[Attention Rollout] How to combine 12 attention heads into one map",
482 )
483 discard_ratio = gr.Slider(
484 minimum=0.0, maximum=0.99, value=0.9, step=0.05,
485 label="Low-Attention Discard Ratio",
486 info="[Attention Rollout] Zeroes out lowest-attention patches (noise reduction)",
487 )
488 threshold = gr.Slider(
489 minimum=0.1, maximum=0.9, value=0.5, step=0.05,
490 label="Prediction Threshold",
491 info="Sigmoid probability cutoff for positive prediction",
492 )
493
494 analyze_btn = gr.Button(
495 "🔬 Analyze X-Ray",
496 variant="primary",
497 size="lg",
498 )
499 status_text = gr.Textbox(
500 label="Status",
501 interactive=False,
502 show_label=True,
503 max_lines=2,
504 )
505
506 # Sample images info
507 gr.Markdown(
508 """
509 > **💡 Tips**
510 > - Use frontal (PA or AP) chest X-ray images
511 > - PNG or JPEG format accepted
512 > - Best results with 1024×1024 pixel images
513 > - Works with exported DICOM screenshots
514 """
515 )
516
517 # ── Right Column: Output ─────────────────────────────────────────
518 with gr.Column(scale=2, min_width=600):
519 gr.Markdown("### 🔥 Attention Rollout Visualization")
520 heatmap_output = gr.Plot(
521 label="Attention Rollout Analysis",
522 show_label=False,
523 )
524
525 gr.Markdown("### 📊 Disease Probability Scores")
526 prob_output = gr.Plot(
527 label="Disease Probabilities",
528 show_label=False,
529 )
530
531 gr.Markdown("### 🩺 Diagnosis Summary")
532 diagnosis_output = gr.Markdown(
533 value="*Upload an X-ray and click Analyze to see results.*"
534 )
535
536 # ── How It Works ──────────────────────────────────────────────────────
537 with gr.Accordion("📖 How It Works", open=False):
538 gr.Markdown("""
539 ## Architecture
540
541 | Component | Details |
542 |---|---|
543 | **Model** | ViT-Base-16 (google/vit-base-patch16-224-in21k) |
544 | **Pre-training** | ImageNet-21k (14M images, 21K classes) |
545 | **Fine-tuning** | NIH ChestX-ray14 (112,120 frontal X-rays) |
546 | **Task** | Multi-label classification — 14 simultaneous disease predictions |
547 | **Loss** | Weighted Binary Cross-Entropy (handles severe class imbalance) |
548 | **Preprocessing** | CLAHE contrast enhancement → Albumentations augmentation |
549 | **Explainability** | Attention Rollout (Abnar & Zuidema, 2020) |
550
551 ## Attention Rollout Algorithm
552
553 Standard Grad-CAM doesn't work well with pure Vision Transformers because
554 ViTs don't have intermediate spatial feature maps like CNNs.
555
556 **Attention Rollout** instead:
557 1. Extracts raw attention weights from all 12 transformer layers
558 2. Averages across all 12 attention heads per layer
559 3. Adds an identity matrix (modeling residual/skip connections)
560 4. Re-normalizes each row
561 5. Multiplies all 12 matrices in sequence → propagates attention end-to-end
562 6. Reads the `[CLS]` token row → shows which 14×14 patches it attends to
563 7. Upsamples 14×14 → 224×224 and overlays as a heatmap
564
565 This shows **where in the X-ray the model is looking** when it makes each prediction.
566
567 ## NIH Dataset — 14 Disease Labels
568
569 ```
570 Atelectasis · Cardiomegaly · Effusion · Infiltration · Mass · Nodule
571 Pneumonia · Pneumothorax · Consolidation · Edema · Emphysema
572 Fibrosis · Pleural_Thickening · Hernia
573 ```
574
575 ## References
576
577 - Wang et al. (2017). *ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks*. CVPR.
578 - Dosovitskiy et al. (2021). *An Image is Worth 16x16 Words*. ICLR.
579 - Abnar & Zuidema (2020). *Quantifying Attention Flow in Transformers*. arXiv:2005.00928.
580 """)
581
582 gr.HTML(FOOTER_HTML)
583
584 # ── Event Binding ─────────────────────────────────────────────────────
585 analyze_btn.click(
586 fn=analyze_xray,
587 inputs=[image_input, head_fusion, discard_ratio, threshold, explainability_method, target_disease],
588 outputs=[heatmap_output, prob_output, diagnosis_output, status_text],
589 api_name="analyze",
590 )
591
592 # Also trigger on image upload (optional — comment out to disable auto-run)
593 # image_input.change(
594 # fn=analyze_xray,
595 # inputs=[image_input, head_fusion, discard_ratio, threshold],
596 # outputs=[heatmap_output, prob_output, diagnosis_output, status_text],
597 # )
598
599 return demo
600
601
602# ── Entry Point ───────────────────────────────────────────────────────────────
603
604if __name__ == "__main__":
605 demo = build_interface()
606 demo.launch(
607 server_port=cfg.inference.gradio_port,
608 share=cfg.inference.gradio_share,
609 show_error=True,
610 inbrowser=True,
611 )
612