#!/usr/bin/env python3 """ RetinaSense v3.0 — Grad-CAM Explainability Pipeline ==================================================== Implements: 1. ViTGradCAM : Gradient-weighted Class Activation Maps for ViT backbone 2. OODDetector : Mahalanobis-distance out-of-distribution detection 3. predict_with_gradcam : Full inference pipeline (preprocess → OOD → CAM → calibrate) 4. Batch evaluation on 20 test images (4 per class) 5. Disease-specific heatmap validation against known anatomical regions 6. Clinical output report (GRADCAM_REPORT.md) Usage: python gradcam_v3.py """ import os import sys import json import warnings 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 datetime import datetime import time warnings.filterwarnings('ignore') import torch import torch.nn as nn import torch.nn.functional as F from torchvision import transforms import timm # ================================================================ # CONFIGURATION # ================================================================ BASE_DIR = '/teamspace/studios/this_studio' OUTPUT_DIR = os.path.join(BASE_DIR, 'outputs_v3') GRADCAM_DIR = os.path.join(OUTPUT_DIR, 'gradcam') os.makedirs(GRADCAM_DIR, exist_ok=True) MODEL_PATH = os.path.join(OUTPUT_DIR, 'best_model.pth') THRESHOLDS_PATH = os.path.join(OUTPUT_DIR, 'thresholds.json') TEMPERATURE_PATH = os.path.join(OUTPUT_DIR, 'temperature.json') TEST_CSV = os.path.join(BASE_DIR, 'data', 'test_split.csv') NORM_STATS_PATH = os.path.join(BASE_DIR, 'data', 'fundus_norm_stats.json') CLASS_NAMES = ['Normal', 'Diabetes/DR', 'Glaucoma', 'Cataract', 'AMD'] NUM_CLASSES = 5 IMG_SIZE = 224 DROPOUT = 0.3 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Anatomical regions expected for each disease class EXPECTED_REGIONS = { 0: 'low uniform activation (Normal)', 1: 'scattered periphery and macula (DR)', 2: 'optic disc (Glaucoma)', 3: 'diffuse lens opacity (Cataract)', 4: 'macula/centre-temporal (AMD)', } print('=' * 65) print(' RetinaSense v3.0 — Grad-CAM Explainability Pipeline') print('=' * 65) print(f' Device : {DEVICE}') if torch.cuda.is_available(): print(f' GPU : {torch.cuda.get_device_name(0)}') print(f' Output : {GRADCAM_DIR}') print('=' * 65) # ================================================================ # LOAD NORMALISATION STATS # ================================================================ if os.path.exists(NORM_STATS_PATH): with open(NORM_STATS_PATH) as f: norm_stats = json.load(f) NORM_MEAN = norm_stats['mean_rgb'] NORM_STD = norm_stats['std_rgb'] print(f' Fundus norm stats: mean={[round(v,4) for v in NORM_MEAN]}, std={[round(v,4) for v in NORM_STD]}') else: NORM_MEAN = [0.485, 0.456, 0.406] NORM_STD = [0.229, 0.224, 0.225] print(' Using ImageNet normalisation fallback') # ================================================================ # MODEL ARCHITECTURE (mirrors retinasense_v3.py exactly) # ================================================================ class MultiTaskViT(nn.Module): """ViT-Base-Patch16-224 with disease + severity heads.""" def __init__(self, n_disease=NUM_CLASSES, n_severity=5, drop=DROPOUT): super().__init__() self.backbone = timm.create_model( 'vit_base_patch16_224', pretrained=False, num_classes=0 ) feat = 768 # CLS token dimension self.drop = nn.Dropout(drop) self.disease_head = nn.Sequential( nn.Linear(feat, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(0.3), nn.Linear(512, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2), nn.Linear(256, n_disease), ) self.severity_head = nn.Sequential( nn.Linear(feat, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, n_severity), ) def forward(self, x): f = self.backbone(x) # (B, 768) — CLS token features f = self.drop(f) return self.disease_head(f), self.severity_head(f) def get_features(self, x): """Return raw CLS token features (before heads and dropout).""" return self.backbone(x) # (B, 768) def forward_with_tokens(self, x): """Return (disease_logits, full_token_sequence (B,197,768)).""" tokens = self.backbone.forward_features(x) # (B, 197, 768) cls_feat = tokens[:, 0, :] cls_feat_d = self.drop(cls_feat) d_out = self.disease_head(cls_feat_d) return d_out, tokens # ================================================================ # LOAD MODEL # ================================================================ print('\nLoading model...') model = MultiTaskViT().to(DEVICE) ckpt = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False) model.load_state_dict(ckpt['model_state_dict']) model.eval() print(f' Loaded: {MODEL_PATH}') print(f' Checkpoint epoch: {ckpt.get("epoch", "?") + 1} val_acc={ckpt.get("val_acc", 0):.2f}%') # Load thresholds and temperature with open(THRESHOLDS_PATH) as f: thr_data = json.load(f) THRESHOLDS = thr_data['thresholds'] with open(TEMPERATURE_PATH) as f: temp_data = json.load(f) TEMPERATURE = temp_data['temperature'] print(f' Temperature T = {TEMPERATURE:.4f}') print(f' Thresholds = {[round(t,3) for t in THRESHOLDS]}') # ================================================================ # IMAGE PREPROCESSING # ================================================================ def ben_graham(path, sz=IMG_SIZE, sigma=10): """Ben Graham high-frequency fundus enhancement (APTOS-style).""" img = cv2.imread(path) if img is None: img = np.array(Image.open(path).convert('RGB')) img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, (sz, sz)) img = cv2.addWeighted(img, 4, cv2.GaussianBlur(img, (0, 0), sigma), -4, 128) mask = np.zeros(img.shape[:2], dtype=np.uint8) cv2.circle(mask, (sz // 2, sz // 2), int(sz * 0.48), 255, -1) return cv2.bitwise_and(img, img, mask=mask) def clahe_preprocess(path, sz=IMG_SIZE): """CLAHE-based contrast enhancement (ODIR-style).""" img = cv2.imread(path) if img is None: img = np.array(Image.open(path).convert('RGB')) img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) img = cv2.resize(img, (sz, sz)) lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB) clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) lab[:, :, 0] = clahe.apply(lab[:, :, 0]) img = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR) return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) def load_and_preprocess(image_path, dataset='auto'): """ Load image and apply domain-conditional preprocessing. Returns: img_np : numpy (224,224,3) uint8 preprocessed img_orig : numpy (224,224,3) uint8 original (for overlay) """ # Normalise path: handle relative paths from CSV (e.g. "aptos/..." or "./aptos/...") # If the path is already absolute and exists, use it directly. # Otherwise resolve relative to BASE_DIR, stripping any leading ./ or .// first. if not os.path.isabs(image_path): # Strip any leading './' or '../' patterns to get a clean relative path clean = image_path while clean.startswith('./') or clean.startswith('.//'): clean = clean[2:] if clean.startswith('./') else clean[3:] image_path = os.path.join(BASE_DIR, clean)thinl # Auto-detect domain if dataset == 'auto': if 'aptos' in image_path.lower() or 'gaussian' in image_path.lower(): dataset = 'APTOS' else: dataset = 'ODIR' # Load original (unprocessed, for overlay) raw = cv2.imread(image_path) if raw is None: raw = np.array(Image.open(image_path).convert('RGB')) else: raw = cv2.cvtColor(raw, cv2.COLOR_BGR2RGB) img_orig = cv2.resize(raw, (IMG_SIZE, IMG_SIZE)) # Apply preprocessing if dataset == 'APTOS': img_np = ben_graham(image_path) else: img_np = clahe_preprocess(image_path) return img_np, img_orig def preprocess_to_tensor(img_np): """Convert preprocessed numpy image to normalised tensor (1, 3, 224, 224).""" transform = transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), transforms.Normalize(NORM_MEAN, NORM_STD), ]) return transform(img_np).unsqueeze(0) # ================================================================ # ViT GRAD-CAM # ================================================================ class ViTAttentionRollout: """ Attention Rollout for Vision Transformer (Abnar & Zuidema, 2020). WHY this works better than Grad-CAM for ViT: - ViT uses CLS token pooling: gradients flow ONLY through CLS token (index 0) - All 196 patch token gradients at block output = zero → Grad-CAM fails - Attention Rollout instead traces how information flows from image patches to the CLS token across ALL 12 transformer layers - Accounts for residual connections by adding identity to each attention map - Produces spatially meaningful maps that highlight actual disease regions Algorithm: 1. Collect attention maps A_l from all 12 blocks: shape (B, H, N, N) 2. Average over H heads: A_l → (B, N, N) 3. Add identity: A_l = A_l + I (accounts for residual connection) 4. Row-normalize: A_l = A_l / row_sum 5. Matrix-multiply all layers: Rollout = A_0 @ A_1 @ ... @ A_11 6. Take CLS row, patch tokens only: Rollout[0, 1:] → (196,) 7. Reshape 14×14 → bilinear upsample → 224×224 """ def __init__(self, model, discard_ratio=0.97): self.model = model self.discard_ratio = discard_ratio # zero out weakest attention weights self._attention_maps = [] self._hooks = [] # Disable fused attention for explicit weight access for block in model.backbone.blocks: block.attn.fused_attn = False # Register forward hooks on ALL transformer blocks for block in model.backbone.blocks: h = block.attn.register_forward_hook(self._attn_hook) self._hooks.append(h) def _attn_hook(self, module, input, output): """Capture softmax attention weights from each block.""" x = input[0] B, N, C = x.shape with torch.no_grad(): qkv = module.qkv(x).reshape(B, N, 3, module.num_heads, module.head_dim).permute(2, 0, 3, 1, 4) q, k, _ = qkv.unbind(0) q, k = module.q_norm(q), module.k_norm(k) attn = (q * module.scale @ k.transpose(-2, -1)).softmax(dim=-1) self._attention_maps.append(attn.detach().cpu()) # (B, H, N, N) def generate(self, image_tensor, class_idx=None): """ Generate attention rollout heatmap. Returns: heatmap : np.ndarray (224, 224) float32 [0, 1] High values = regions most important for prediction predicted_label : int confidence : float (raw softmax) """ self.model.eval() self._attention_maps = [] with torch.no_grad(): image_tensor = image_tensor.to(DEVICE) d_out, _ = self.model(image_tensor) probs = torch.softmax(d_out, dim=1) predicted_label = int(probs.argmax(dim=1).item()) confidence = float(probs[0, predicted_label].item()) if class_idx is None: class_idx = predicted_label # --- Attention Rollout computation --- # Stack all layer attentions: list of (1, H, N, N) → (L, H, N, N) attn_stack = torch.stack(self._attention_maps, dim=0) # (L, 1, H, N, N) attn_stack = attn_stack[:, 0] # (L, H, N, N), batch=1 # Average over heads attn_mean = attn_stack.mean(dim=1) # (L, N, N) # Optional: discard weakest connections (sharpens the map) if self.discard_ratio > 0: flat = attn_mean.reshape(attn_mean.shape[0], -1) thresh = torch.quantile(flat, self.discard_ratio, dim=1, keepdim=True) thresh = thresh.unsqueeze(-1) # broadcast over N,N attn_mean = torch.where(attn_mean >= thresh, attn_mean, torch.zeros_like(attn_mean)) # Add identity matrix for residual connection, then row-normalize I = torch.eye(attn_mean.shape[-1]).unsqueeze(0) # (1, N, N) attn_aug = attn_mean + I attn_aug = attn_aug / attn_aug.sum(dim=-1, keepdim=True).clamp(min=1e-8) # Matrix-multiply across all layers rollout = attn_aug[0] for l in range(1, len(attn_aug)): rollout = rollout @ attn_aug[l] # CLS token's attention to all patch tokens (skip CLS at index 0) cls_attention = rollout[0, 1:] # (196,) # Reshape and upsample spatial = cls_attention.numpy().reshape(14, 14).astype(np.float32) spatial = cv2.resize(spatial, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_LINEAR) # Normalize to [0, 1] s_min, s_max = spatial.min(), spatial.max() if s_max - s_min > 1e-8: spatial = (spatial - s_min) / (s_max - s_min) else: spatial = np.zeros_like(spatial) # Power-curve stretch: boosts mid-range attention values for visual clarity # gamma < 1 brightens the map; 0.4 gives strong contrast enhancement spatial = np.power(spatial, 0.4) return spatial.astype(np.float32), predicted_label, confidence def overlay(self, original_image_np, heatmap, alpha=0.5): """ Blend attention rollout heatmap onto original fundus image. Uses INFERNO colormap (dark=low, bright=high) — better for medical images. Args: original_image_np : (224, 224, 3) uint8 RGB heatmap : (224, 224) float32 [0, 1] alpha : heatmap opacity (0.5 gives good visibility) Returns: overlay : (224, 224, 3) uint8 RGB """ # Apply JET colormap heatmap_uint8 = (heatmap * 255).astype(np.uint8) colormap = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET) colormap_rgb = cv2.cvtColor(colormap, cv2.COLOR_BGR2RGB) # Apply circular mask to ignore black borders (fundus images are circular) h, w = heatmap.shape cy, cx = h // 2, w // 2 radius = min(h, w) // 2 - 5 mask = np.zeros((h, w), dtype=np.float32) cv2.circle(mask, (cx, cy), radius, 1.0, -1) mask = cv2.GaussianBlur(mask, (21, 21), 0) # Blend only inside the retinal circle orig = original_image_np.astype(np.float32) cmap = colormap_rgb.astype(np.float32) blended = orig.copy() for c in range(3): blended[:, :, c] = ( orig[:, :, c] * (1 - alpha * mask) + cmap[:, :, c] * (alpha * mask) ) return np.clip(blended, 0, 255).astype(np.uint8) def remove_hooks(self): """Clean up all registered hooks.""" for h in self._hooks: h.remove() self._hooks = [] # Keep old name as alias for backward compatibility ViTGradCAM = ViTAttentionRollout # ================================================================ # OOD DETECTION (Mahalanobis Distance) # ================================================================ class OODDetector: """ Out-of-Distribution detector using class-conditional Mahalanobis distance. Fit on training-set CLS token features; at inference, computes the minimum Mahalanobis distance from the test feature to the nearest class centroid. High distance = likely OOD. """ def __init__(self, threshold_percentile=97.5): self.class_means = None # (num_classes, feat_dim) self.cov_inv = None # (feat_dim, feat_dim) self.ood_threshold = None self.threshold_percentile = threshold_percentile self.is_fitted = False def fit(self, model, dataloader, device, max_batches=60): """ Extract CLS token features for all samples, compute class-conditional means and shared inverse covariance matrix. """ print(' OODDetector.fit: extracting features...') all_features = [] all_labels = [] model.eval() with torch.no_grad(): for i, batch in enumerate(dataloader): if i >= max_batches: break imgs, d_lbl, _ = batch imgs = imgs.to(device) feats = model.get_features(imgs) # (B, 768) all_features.append(feats.cpu().numpy()) all_labels.append(d_lbl.numpy()) features = np.concatenate(all_features, axis=0) # (N, 768) labels = np.concatenate(all_labels, axis=0) # (N,) num_classes = NUM_CLASSES feat_dim = features.shape[1] # Class-conditional means self.class_means = np.zeros((num_classes, feat_dim), dtype=np.float64) for c in range(num_classes): mask = labels == c if mask.sum() > 0: self.class_means[c] = features[mask].mean(axis=0) # Shared (pooled) covariance matrix cov = np.zeros((feat_dim, feat_dim), dtype=np.float64) total = 0 for c in range(num_classes): mask = labels == c if mask.sum() < 2: continue diff = features[mask] - self.class_means[c] cov += diff.T @ diff total += mask.sum() cov /= max(total - num_classes, 1) # Regularise for numerical stability (add small diagonal) cov += np.eye(feat_dim) * 1e-4 # Pseudo-inverse via SVD (numerically stable for high-dim) try: self.cov_inv = np.linalg.pinv(cov) except np.linalg.LinAlgError: self.cov_inv = np.eye(feat_dim) # Compute train-set Mahalanobis distances to set threshold train_dists = [] for feat in features: d = self._mahal_min_dist(feat) train_dists.append(d) self.ood_threshold = float(np.percentile(train_dists, self.threshold_percentile)) self.is_fitted = True print(f' OOD threshold ({self.threshold_percentile}th pct): {self.ood_threshold:.4f}') print(f' Features extracted: {len(features)} samples') def _mahal_min_dist(self, feat): """Minimum Mahalanobis distance to any class centroid.""" min_dist = float('inf') for c in range(NUM_CLASSES): diff = feat - self.class_means[c] dist = float(diff @ self.cov_inv @ diff) dist = max(dist, 0.0) # guard against floating-point negatives if dist < min_dist: min_dist = dist return np.sqrt(min_dist) def score(self, features): """ Compute OOD score for a batch of features. Args: features : np.ndarray (N, 768) or (768,) Returns: distances : np.ndarray (N,) Mahalanobis distances ood_flags : np.ndarray (N,) bool, True = likely OOD """ if not self.is_fitted: raise RuntimeError('OODDetector.fit() must be called before score()') if features.ndim == 1: features = features[np.newaxis, :] distances = np.array([self._mahal_min_dist(f) for f in features]) ood_flags = distances > self.ood_threshold return distances, ood_flags def save(self, path): np.savez(path, class_means=self.class_means, cov_inv=self.cov_inv, ood_threshold=np.array([self.ood_threshold]), threshold_percentile=np.array([self.threshold_percentile])) print(f' OOD detector saved -> {path}.npz') def load(self, path): if not path.endswith('.npz'): path = path + '.npz' data = np.load(path) self.class_means = data['class_means'] self.cov_inv = data['cov_inv'] self.ood_threshold = float(data['ood_threshold'][0]) self.threshold_percentile = float(data['threshold_percentile'][0]) self.is_fitted = True print(f' OOD detector loaded <- {path}') # ================================================================ # ATTENTION REGION ANALYSER # ================================================================ def analyse_attention_region(heatmap, disease_class): """ Check if the Grad-CAM heatmap activation pattern is consistent with the expected anatomical region for the given disease. Returns: attention_region : str describing where activation is is_consistent : bool region_scores : dict with activation energy in each zone """ h, w = heatmap.shape # (224, 224) cx, cy = w // 2, h // 2 # Define anatomical zones (approximate, relative to image centre) # Centre disc zone: circle r ~ 30px (optic disc) r_disc = int(h * 0.13) # Macula zone: circle r ~ 55px centred slightly temporal r_macula = int(h * 0.25) cx_mac = int(cx + w * 0.10) # slightly nasal offset # Build zone masks Y, X = np.ogrid[:h, :w] # Optic disc (small circle, centre of image) disc_mask = ((X - cx)**2 + (Y - cy)**2) <= r_disc**2 # Macula (larger circle, centre-temporal) macula_mask = ((X - cx_mac)**2 + (Y - cy)**2) <= r_macula**2 # Periphery: outer 30% of image peri_mask = (X < int(w * 0.15)) | (X > int(w * 0.85)) | \ (Y < int(h * 0.15)) | (Y > int(h * 0.85)) # Compute mean activation in each zone disc_score = float(heatmap[disc_mask].mean()) if disc_mask.sum() > 0 else 0.0 macula_score = float(heatmap[macula_mask].mean()) if macula_mask.sum() > 0 else 0.0 peri_score = float(heatmap[peri_mask].mean()) if peri_mask.sum() > 0 else 0.0 overall_mean = float(heatmap.mean()) region_scores = { 'optic_disc': round(disc_score, 4), 'macula': round(macula_score, 4), 'periphery': round(peri_score, 4), 'overall': round(overall_mean, 4), } # Determine dominant region label max_zone = max(region_scores, key=lambda k: region_scores[k] if k != 'overall' else -1) zone_labels = { 'optic_disc': 'optic disc (centre)', 'macula': 'macula (centre-temporal)', 'periphery': 'scattered periphery', } dominant_label = zone_labels.get(max_zone, 'diffuse') # Assess uniformity (low std = diffuse / uniform) if heatmap.std() < 0.10: dominant_label = 'diffuse (low activation)' # Check consistency with expected region consistency_map = { 0: lambda s: s['overall'] < 0.25, # Normal → low uniform 1: lambda s: s['periphery'] > 0.20 or s['macula'] > 0.25, # DR → periphery/macula 2: lambda s: s['optic_disc'] > 0.30, # Glaucoma → disc 3: lambda s: heatmap.std() < 0.15, # Cataract → diffuse 4: lambda s: s['macula'] > 0.25, # AMD → macula } check_fn = consistency_map.get(disease_class, lambda s: True) is_consistent = check_fn(region_scores) return dominant_label, is_consistent, region_scores # ================================================================ # FULL INFERENCE PIPELINE # ================================================================ def predict_with_gradcam(image_path, model, gradcam, ood_detector, thresholds, temperature, device, true_label=None, dataset='auto'): """ End-to-end inference with Grad-CAM and OOD detection. Steps: 1. Load and preprocess image (Ben Graham for APTOS, CLAHE for ODIR) 2. OOD check on ViT CLS token features 3. Generate Grad-CAM heatmap 4. Apply temperature scaling to logits 5. Apply per-class thresholds 6. Analyse attention region Returns: dict with predicted_class, confidence, gradcam_heatmap, etc. """ # 1. Preprocess img_np, img_orig = load_and_preprocess(image_path, dataset=dataset) img_tensor = preprocess_to_tensor(img_np).to(device) # 2. OOD check using raw CLS features model.eval() with torch.no_grad(): features = model.get_features(img_tensor).cpu().numpy() # (1, 768) if ood_detector.is_fitted: distances, ood_flags = ood_detector.score(features) ood_distance = float(distances[0]) ood_flag = bool(ood_flags[0]) else: ood_distance = 0.0 ood_flag = False # 3. Generate Grad-CAM (also runs forward + backward pass) heatmap, predicted_label, raw_confidence = gradcam.generate(img_tensor) # 4. Temperature-scaled calibrated probabilities # Run a clean no-grad forward pass to get stable logits for calibration model.eval() with torch.no_grad(): raw_feats = model.backbone(img_tensor) # (1, 768) raw_feats = model.drop(raw_feats) logits = model.disease_head(raw_feats).float().cpu() # (1, 5) scaled_logits = logits / temperature calibrated_probs = torch.softmax(scaled_logits, dim=1)[0].numpy() # (5,) # 5. Apply per-class thresholds above = [i for i, (p, t) in enumerate(zip(calibrated_probs, thresholds)) if p >= t] if above: final_label = int(above[np.argmax([calibrated_probs[i] for i in above])]) else: final_label = int(np.argmax(calibrated_probs)) final_confidence = float(calibrated_probs[final_label]) predicted_class = CLASS_NAMES[final_label] # 6. Heatmap overlay gradcam_overlay = gradcam.overlay(img_orig, heatmap, alpha=0.7) # 7. Attention region analysis attention_region, region_consistent, region_scores = analyse_attention_region( heatmap, final_label ) # Append disease name for clarity disease_tag = CLASS_NAMES[final_label].replace('/', '-') attention_region_full = f'{attention_region} ({disease_tag})' # 8. Review flag: low confidence OR OOD review_flag = ood_flag or final_confidence < 0.50 return { 'image_path': image_path, 'predicted_class': predicted_class, 'predicted_label': final_label, 'confidence': round(final_confidence, 4), 'raw_confidence': round(raw_confidence, 4), 'all_probabilities': [round(float(p), 4) for p in calibrated_probs], 'gradcam_heatmap': heatmap, # (224, 224) float32 'gradcam_overlay': gradcam_overlay, # (224, 224, 3) uint8 'img_orig': img_orig, # original for display 'ood_flag': ood_flag, 'ood_distance': round(ood_distance, 4), 'review_flag': review_flag, 'attention_region': attention_region_full, 'region_scores': region_scores, 'region_consistent': region_consistent, 'true_label': true_label, } # ================================================================ # BATCH EVALUATION # ================================================================ def run_batch_evaluation(model, gradcam, ood_detector, thresholds, temperature, device, n_per_class=4): """ Run inference on n_per_class images per disease class (20 total). Saves individual overlay images + summary grid. """ import pandas as pd print(f'\nRunning batch evaluation ({n_per_class} per class = {n_per_class * NUM_CLASSES} total)...') df = pd.read_csv(TEST_CSV) # Collect n_per_class unique samples per class samples = [] for label in range(NUM_CLASSES): subset = df[df['disease_label'] == label].drop_duplicates(subset='image_path') chosen = subset.head(n_per_class) for _, row in chosen.iterrows(): samples.append({ 'image_path': row['image_path'], 'true_label': int(row['disease_label']), 'dataset': str(row.get('dataset', 'auto')), }) results = [] failed = [] for i, sample in enumerate(samples): img_path = sample['image_path'] true_label = sample['true_label'] dataset = sample['dataset'] print(f' [{i+1:2d}/{len(samples)}] {CLASS_NAMES[true_label]:15s} | {os.path.basename(img_path)}', end=' ') try: result = predict_with_gradcam( img_path, model, gradcam, ood_detector, thresholds, temperature, device, true_label=true_label, dataset=dataset, ) correct = (result['predicted_label'] == true_label) flag_str = ' [OOD]' if result['ood_flag'] else '' flag_str += ' [REVIEW]' if result['review_flag'] else '' print(f'-> pred={result["predicted_class"]:15s} conf={result["confidence"]:.3f} {"OK" if correct else "WRONG"}{flag_str}') # Save overlay image save_name = f'gradcam_{i+1:02d}_true{true_label}_pred{result["predicted_label"]}_{os.path.splitext(os.path.basename(img_path))[0][:20]}.png' save_path = os.path.join(GRADCAM_DIR, save_name) fig, axes = plt.subplots(1, 3, figsize=(12, 4)) axes[0].imshow(result['img_orig']) axes[0].set_title(f'Original\nTrue: {CLASS_NAMES[true_label]}', fontsize=9) axes[0].axis('off') axes[1].imshow(result['gradcam_heatmap'], cmap='jet', vmin=0, vmax=1) axes[1].set_title('Grad-CAM Heatmap', fontsize=9) axes[1].axis('off') axes[2].imshow(result['gradcam_overlay']) flag_line = ' [OOD]' if result['ood_flag'] else '' axes[2].set_title( f'Overlay\nPred: {result["predicted_class"]} ({result["confidence"]:.2f}){flag_line}', fontsize=9, color='red' if not correct else 'green' ) axes[2].axis('off') plt.suptitle( f'Attention: {result["attention_region"]}', fontsize=8, color='gray' ) plt.tight_layout() plt.savefig(save_path, dpi=120, bbox_inches='tight') plt.close() result['save_path'] = save_path results.append(result) except Exception as e: print(f' ERROR: {e}') failed.append({'image_path': img_path, 'error': str(e)}) return results, failed # ================================================================ # SUMMARY GRID (4 rows = classes 0-4, 4 cols = samples) # ================================================================ def save_summary_grid(results): """Save a 5×4 summary grid (rows=classes, cols=samples).""" n_rows = NUM_CLASSES n_cols = 4 # Group results by true label by_class = {i: [] for i in range(NUM_CLASSES)} for r in results: tl = r.get('true_label', r['predicted_label']) by_class[tl].append(r) fig, axes = plt.subplots(n_rows, n_cols, figsize=(16, 20)) fig.patch.set_facecolor('#1a1a2e') for row_idx in range(n_rows): class_results = by_class[row_idx] for col_idx in range(n_cols): ax = axes[row_idx, col_idx] if col_idx < len(class_results): r = class_results[col_idx] ax.imshow(r['gradcam_overlay']) correct = (r['predicted_label'] == r.get('true_label', r['predicted_label'])) border_color = '#2ecc71' if correct else '#e74c3c' for spine in ax.spines.values(): spine.set_edgecolor(border_color) spine.set_linewidth(3) label_str = f'{r["predicted_class"]}\n{r["confidence"]:.2f}' if r['ood_flag']: label_str += '\n[OOD]' ax.set_title(label_str, fontsize=7, color='white', pad=2) else: ax.set_facecolor('#1a1a2e') ax.axis('off') if col_idx == 0: ax.set_ylabel(CLASS_NAMES[row_idx], rotation=0, labelpad=50, fontsize=10, color='white', fontweight='bold', va='center') plt.suptitle( 'RetinaSense v3.0 — Grad-CAM Summary Grid\n' 'Rows = True Class | Green border = Correct | Red border = Wrong', fontsize=12, color='white', y=1.01 ) plt.tight_layout() grid_path = os.path.join(GRADCAM_DIR, 'gradcam_summary_grid.png') plt.savefig(grid_path, dpi=130, bbox_inches='tight', facecolor=fig.get_facecolor()) plt.close() print(f' Summary grid saved -> {grid_path}') return grid_path # ================================================================ # DISEASE-SPECIFIC HEATMAP VALIDATION # ================================================================ def validate_heatmaps(results): """ Check per-disease whether Grad-CAM activates the expected anatomical region. Returns a validation summary dict, saves to heatmap_validation.json. """ print('\nRunning disease-specific heatmap validation...') validation = {} for cls_idx, cls_name in enumerate(CLASS_NAMES): cls_results = [r for r in results if r.get('true_label') == cls_idx] if not cls_results: validation[cls_name] = {'n_samples': 0} continue consistent_count = sum(1 for r in cls_results if r.get('region_consistent', False)) avg_scores = {k: 0.0 for k in ['optic_disc', 'macula', 'periphery', 'overall']} for r in cls_results: for k in avg_scores: avg_scores[k] += r['region_scores'].get(k, 0.0) for k in avg_scores: avg_scores[k] = round(avg_scores[k] / len(cls_results), 4) dominant_zone = max( ['optic_disc', 'macula', 'periphery'], key=lambda k: avg_scores[k] ) validation[cls_name] = { 'n_samples': len(cls_results), 'expected_region': EXPECTED_REGIONS[cls_idx], 'dominant_zone': dominant_zone, 'consistent_samples': consistent_count, 'consistency_pct': round(100 * consistent_count / len(cls_results), 1), 'avg_region_scores': avg_scores, } print(f' {cls_name:15s}: {consistent_count}/{len(cls_results)} consistent ' f'({validation[cls_name]["consistency_pct"]:.0f}%) ' f'dominant={dominant_zone}') # Save val_path = os.path.join(GRADCAM_DIR, 'heatmap_validation.json') with open(val_path, 'w') as f: json.dump(validation, f, indent=2) print(f' Validation saved -> {val_path}') return validation # ================================================================ # CLINICAL REPORT # ================================================================ def generate_clinical_report(results, validation, ood_stats, failed): """Generate GRADCAM_REPORT.md with clinical analysis.""" now = datetime.now().strftime('%Y-%m-%d %H:%M:%S') n_total = len(results) n_correct = sum(1 for r in results if r.get('predicted_label') == r.get('true_label')) n_ood = sum(1 for r in results if r.get('ood_flag')) n_review = sum(1 for r in results if r.get('review_flag')) avg_conf = np.mean([r['confidence'] for r in results]) if results else 0.0 lines = [ '# RetinaSense v3.0 — Grad-CAM Clinical Report', f'', f'**Generated**: {now} ', f'**Model**: ViT-Base-Patch16-224 (81.19% test accuracy) ', f'**Pipeline**: Grad-CAM + Mahalanobis OOD + Temperature Scaling + Per-Class Thresholds', '', '---', '', '## Executive Summary', '', f'| Metric | Value |', f'|--------|-------|', f'| Images processed | {n_total} |', f'| Correct predictions | {n_correct}/{n_total} ({100*n_correct/max(n_total,1):.1f}%) |', f'| Avg calibrated confidence | {avg_conf:.3f} |', f'| OOD flags raised | {n_ood} |', f'| Human review flags | {n_review} |', f'| Failed images | {len(failed)} |', f'| Temperature T | {TEMPERATURE:.4f} |', '', '---', '', '## Per-Sample Predictions', '', '| # | Image | True | Predicted | Confidence | OOD | Review | Attention Region |', '|---|-------|------|-----------|-----------|-----|--------|-----------------|', ] for i, r in enumerate(results): true_name = CLASS_NAMES[r['true_label']] if r.get('true_label') is not None else 'Unknown' correct_marker = 'OK' if r['predicted_label'] == r.get('true_label') else '**WRONG**' lines.append( f'| {i+1} | {os.path.basename(r["image_path"])[:25]} ' f'| {true_name} ' f'| {r["predicted_class"]} ({correct_marker}) ' f'| {r["confidence"]:.3f} ' f'| {"YES" if r["ood_flag"] else "no"} ' f'| {"YES" if r["review_flag"] else "no"} ' f'| {r["attention_region"]} |' ) lines += [ '', '---', '', '## Per-Class Attention Pattern Analysis', '', '| Disease | Expected Region | Dominant Zone | Consistency |', '|---------|----------------|---------------|-------------|', ] for cls_name, v in validation.items(): if v.get('n_samples', 0) == 0: lines.append(f'| {cls_name} | N/A | N/A | N/A (no samples) |') else: lines.append( f'| {cls_name} | {v["expected_region"]} ' f'| {v["dominant_zone"]} ' f'| {v["consistency_pct"]:.0f}% ({v["consistent_samples"]}/{v["n_samples"]}) |' ) lines += [ '', '---', '', '## OOD Detection Statistics', '', f'- **Method**: Mahalanobis distance to nearest class centroid (CLS token features)', f'- **Threshold percentile**: 97.5th percentile of training-set distances', f'- **OOD threshold**: {ood_stats.get("threshold", "N/A")}', f'- **Images flagged OOD**: {n_ood}/{n_total}', '', '### Interpretation', '', '- Mahalanobis distance measures how far a feature embedding lies from known class distributions', '- Low-quality images, extreme artefacts, or off-distribution fundus cameras may trigger OOD flags', '- All OOD-flagged images are automatically sent for human review', '', '---', '', '## Grad-CAM Heatmap Descriptions', '', '| Disease | Expected activation | Clinical significance |', '|---------|--------------------|-----------------------|', '| Normal | Low, uniform | No focal pathology — model attention diffuse |', '| Diabetes/DR | Scattered periphery + macula | Microaneurysms, exudates, NV |', '| Glaucoma | Optic disc (centre) | Structural disc changes, CDR |', '| Cataract | Diffuse lens opacity | Posterior/anterior capsule opacification |', '| AMD | Macula / centre-temporal | Drusen, RPE atrophy, CNV |', '', '---', '', '## Thresholds Applied', '', '| Class | Threshold |', '|-------|-----------|', ] for cls_name, thr in zip(CLASS_NAMES, THRESHOLDS): lines.append(f'| {cls_name} | {thr:.4f} |') lines += [ '', '---', '', '## Deployment Recommendations', '', '1. **Confidence gate**: Flag predictions below 0.50 for mandatory ophthalmologist review.', '2. **OOD gate**: Any Mahalanobis distance above threshold should trigger QC check on image quality before clinical use.', '3. **Grad-CAM review**: Clinicians should inspect heatmaps for cases where model attention does not align with expected anatomy.', '4. **Glaucoma caution**: Current dataset imbalance (46 test samples) — consider supplementing ODIR with additional glaucoma images.', '5. **Continuous monitoring**: Re-calibrate temperature and thresholds quarterly on production data.', '6. **Not for standalone diagnosis**: Grad-CAM is an explainability aid; all predictions require clinical validation.', '', '---', '', f'*Report auto-generated by RetinaSense v3.0 Grad-CAM Pipeline | {now}*', ] report_path = os.path.join(GRADCAM_DIR, 'GRADCAM_REPORT.md') with open(report_path, 'w') as f: f.write('\n'.join(lines)) print(f' Clinical report saved -> {report_path}') return report_path # ================================================================ # MAIN # ================================================================ def main(): t_start = time.time() # ---- 1. Build Grad-CAM --- print('\n[1/6] Initialising ViTGradCAM...') gradcam = ViTGradCAM(model) print(f' Method : Attention Rollout (all 12 transformer blocks)') print(f' Hooks : {len(gradcam._hooks)} attention hooks registered') print(f' fused_attn disabled for attention weight access') # ---- 2. Fit OOD Detector --- print('\n[2/6] Fitting OOD detector...') ood_path = os.path.join(OUTPUT_DIR, 'ood_detector') ood_detector = OODDetector(threshold_percentile=97.5) if os.path.exists(ood_path + '.npz'): ood_detector.load(ood_path) else: # Build a small DataLoader from training data to fit OOD detector import pandas as pd from torch.utils.data import Dataset, DataLoader from torchvision import transforms as T train_df = pd.read_csv(os.path.join(BASE_DIR, 'data', 'train_split.csv')) class SimpleDataset(Dataset): def __init__(self, df): self.df = df.reset_index(drop=True) self.transform = transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), transforms.Normalize(NORM_MEAN, NORM_STD), ]) def __len__(self): return len(self.df) def __getitem__(self, idx): row = self.df.iloc[idx] img_path = str(row['image_path']) if not os.path.isabs(img_path): clean = img_path while clean.startswith('./') or clean.startswith('.//'): clean = clean[2:] if clean.startswith('./') else clean[3:] img_path = os.path.join(BASE_DIR, clean) dataset = str(row.get('dataset', 'auto')) try: img_np, _ = load_and_preprocess(img_path, dataset=dataset) img_tensor = self.transform(img_np) except Exception: img_tensor = torch.zeros(3, IMG_SIZE, IMG_SIZE) lbl = int(row['disease_label']) return img_tensor, torch.tensor(lbl, dtype=torch.long), torch.tensor(0, dtype=torch.long) ood_ds = SimpleDataset(train_df) ood_loader = DataLoader(ood_ds, batch_size=32, shuffle=False, num_workers=4) ood_detector.fit(model, ood_loader, DEVICE, max_batches=80) ood_detector.save(ood_path) # ---- 3. Batch Evaluation --- print('\n[3/6] Batch evaluation on 20 test images...') results, failed = run_batch_evaluation( model, gradcam, ood_detector, THRESHOLDS, TEMPERATURE, DEVICE, n_per_class=4 ) # ---- 4. Summary Grid --- print('\n[4/6] Generating summary grid...') grid_path = save_summary_grid(results) # ---- 5. Heatmap Validation --- print('\n[5/6] Heatmap validation...') validation = validate_heatmaps(results) # ---- 6. Clinical Report --- print('\n[6/6] Generating clinical report...') ood_stats = {'threshold': round(ood_detector.ood_threshold, 4) if ood_detector.is_fitted else 'N/A'} report_path = generate_clinical_report(results, validation, ood_stats, failed) # ---- Cleanup --- gradcam.remove_hooks() # ================================================================ # FINAL SUMMARY # ================================================================ elapsed = time.time() - t_start n_total = len(results) n_correct = sum(1 for r in results if r.get('predicted_label') == r.get('true_label')) avg_conf = np.mean([r['confidence'] for r in results]) if results else 0.0 n_ood = sum(1 for r in results if r['ood_flag']) n_review = sum(1 for r in results if r['review_flag']) print('\n' + '=' * 65) print(' RetinaSense v3.0 — GRAD-CAM PIPELINE COMPLETE') print('=' * 65) print(f' Images processed : {n_total}') print(f' Correct predictions : {n_correct}/{n_total} ({100*n_correct/max(n_total,1):.1f}%)') print(f' Avg calibrated conf : {avg_conf:.3f}') print(f' OOD flags : {n_ood}') print(f' Review flags : {n_review}') print(f' Failed images : {len(failed)}') print(f' Elapsed time : {elapsed:.1f}s') print() print(f' Output directory : {GRADCAM_DIR}') output_files = [ 'gradcam_summary_grid.png', 'heatmap_validation.json', 'GRADCAM_REPORT.md', ] + [os.path.basename(r.get('save_path', '')) for r in results if r.get('save_path')] for fname in output_files: if fname: full = os.path.join(GRADCAM_DIR, fname) exists = os.path.exists(full) print(f' {"[OK]" if exists else "[!!]"} {fname}') print('=' * 65) return results, validation if __name__ == '__main__': main()