""" RetViM - Real Inference Server v4.0 ==================================== Uses ImprovedMedMamba (model.py) with the real .ckpt weights. val_acc = 96.68 % * ViT-B/16 + 2 SS-Conv-SSM blocks * dim = 768 Run: python server.py Open: http://127.0.0.1:8000/ """ from __future__ import annotations import io, os, base64, time, json, hashlib, math, gc import re from pathlib import Path from typing import Optional, List import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torchvision import transforms from PIL import Image import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import matplotlib.cm as cm import scipy.ndimage from fastapi import FastAPI, File, UploadFile, Query, Form from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, HTMLResponse, Response, FileResponse from fastapi.staticfiles import StaticFiles import uvicorn # Import the real model from model import load_model, ImprovedMedMamba # --------------------------------------------------------------------------- # # Config # --------------------------------------------------------------------------- # BASE_DIR = Path(__file__).parent HTML_FILE = BASE_DIR / "retvim-final.html" CKPT_PATH = BASE_DIR / "improved-medmamba-epoch=19-val_acc=0.9668.ckpt" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" PORT = 8000 HOST = "127.0.0.1" START_TIME = time.time() CLASS_NAMES = ["CNV", "DME", "DRUSEN", "NORMAL"] SAMPLE_DIR = BASE_DIR / "sample images" SAMPLE_IMAGES = { "CNV": SAMPLE_DIR / "CNV-103044-6.jpeg", "DME": SAMPLE_DIR / "DME-30521-3.jpeg", "DRUSEN": SAMPLE_DIR / "DRUSEN-1786810-1.jpeg", "NORMAL": SAMPLE_DIR / "NORMAL-33350-1.jpeg", } IMAGENET_MEAN = [0.485, 0.456, 0.406] IMAGENET_STD = [0.229, 0.224, 0.225] CKA_MATRIX = [ [1.00, 0.57, 0.40, 0.09, 0.08, 0.07], [0.57, 1.00, 0.50, 0.06, 0.05, 0.05], [0.40, 0.50, 1.00, 0.47, 0.44, 0.44], [0.09, 0.06, 0.47, 1.00, 0.99, 0.99], [0.08, 0.05, 0.44, 0.99, 1.00, 1.00], [0.07, 0.05, 0.44, 0.99, 1.00, 1.00], ] CKA_LABELS = ["Patch Embed", "ViT Blk-3", "ViT Blk-7", "ViT Blk-11", "Mamba Blk-1", "Mamba Blk-2"] _transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ]) # --------------------------------------------------------------------------- # # Model loading # --------------------------------------------------------------------------- # def build_model() -> ImprovedMedMamba: if CKPT_PATH.exists(): print(f"[*] Loading ImprovedMedMamba checkpoint: {CKPT_PATH.name}") model = load_model(str(CKPT_PATH), device=DEVICE) else: print("[!] Checkpoint not found - using random-init ImprovedMedMamba") model = ImprovedMedMamba(num_classes=4) model.to(DEVICE).eval() print(f"[OK] Model on {DEVICE.upper()} | ViT-B/16 + 2 MedMamba blocks | dim=768") return model # --------------------------------------------------------------------------- # # Image utilities # --------------------------------------------------------------------------- # def _pil_to_tensor(img: Image.Image) -> torch.Tensor: if img.mode != "RGB": img = img.convert("RGB") return _transform(img).unsqueeze(0).to(DEVICE) def _pil_to_b64(img: Image.Image) -> str: buf = io.BytesIO() img.save(buf, format="PNG") buf.seek(0) return "data:image/png;base64," + base64.b64encode(buf.read()).decode() def _fig_to_b64(fig: plt.Figure) -> str: buf = io.BytesIO() fig.savefig(buf, format="png", bbox_inches="tight", facecolor=fig.get_facecolor(), dpi=100) plt.close(fig) buf.seek(0) return "data:image/png;base64," + base64.b64encode(buf.read()).decode() def _feat_to_map(feat: torch.Tensor) -> np.ndarray: arr = feat[0].mean(-1).cpu().numpy() if feat.dim() == 3 else feat.mean(-1).cpu().numpy() side = int(round(arr.shape[0] ** 0.5)) arr = arr[:side * side].reshape(side, side) mn, mx = arr.min(), arr.max() return ((arr - mn) / (mx - mn + 1e-8)).astype(np.float32) def _overlay_heatmap(img: Image.Image, hm: np.ndarray, alpha: float = 0.6, cmap: str = "jet") -> Image.Image: orig = np.array(img.resize((224, 224)).convert("RGB"), dtype=np.float32) / 255.0 if hm.shape != (224, 224): hm = scipy.ndimage.zoom(hm.astype(np.float32), (224 / hm.shape[0], 224 / hm.shape[1]), order=1) hm = np.clip(hm, 0, 1) color = cm.get_cmap(cmap)(hm)[..., :3] blend = np.clip((1 - alpha) * orig + alpha * color, 0, 1) return Image.fromarray((blend * 255).astype(np.uint8)) def _vit_embed(model: ImprovedMedMamba, tensor: torch.Tensor) -> torch.Tensor: """Run ViT patch embedding + CLS + pos embed.""" B = tensor.shape[0] x = model.vit.patch_embed(tensor) cls = model.vit.cls_token.expand(B, -1, -1) x = torch.cat([cls, x], dim=1) return model.vit.pos_drop(x + model.vit.pos_embed) # --------------------------------------------------------------------------- # # Real GradCAM - hooks last ViT block # --------------------------------------------------------------------------- # def real_gradcam(model: ImprovedMedMamba, tensor: torch.Tensor, target_cls: Optional[int] = None) -> np.ndarray: model.eval() acts: List[torch.Tensor] = [] grads: List[torch.Tensor] = [] last_vit = model.vit.blocks[-1] def fwd_hook(m, inp, out): acts.append(out) def bwd_hook(m, gin, gout): grads.append(gout[0]) h1 = last_vit.register_forward_hook(fwd_hook) h2 = last_vit.register_full_backward_hook(bwd_hook) with torch.enable_grad(): inp = tensor.clone().detach().requires_grad_(True) logits = model(inp) if target_cls is None: target_cls = int(logits.argmax(-1)) logits[0, target_cls].backward() h1.remove(); h2.remove() if not acts or not grads: return np.zeros((224, 224), dtype=np.float32) act = acts[0][0, 1:].detach() # (196, D) - drop CLS grad = grads[0][0, 1:].detach() # (196, D) weights = grad.mean(0) cam = (act.cpu() * weights.cpu()).numpy().sum(axis=-1) # (196,) cam = np.maximum(cam, 0) side = int(round(cam.shape[0] ** 0.5)) cam = cam[:side * side].reshape(side, side) mn, mx = cam.min(), cam.max() cam = (cam - mn) / (mx - mn + 1e-8) return scipy.ndimage.zoom(cam.astype(np.float32), 224 / side, order=1) # --------------------------------------------------------------------------- # # Real Attention Rollout # --------------------------------------------------------------------------- # def real_attention_rollout(model: ImprovedMedMamba, tensor: torch.Tensor) -> np.ndarray: model.eval() with torch.no_grad(): model(tensor) attn_maps = model.get_attention_maps() if not attn_maps: return np.zeros((224, 224), dtype=np.float32) N = attn_maps[0].shape[-1] eye = torch.eye(N, device=DEVICE) rollout = eye.clone() for a in attn_maps: am = a[0].mean(0) am = am + eye am = am / am.sum(-1, keepdim=True) rollout = am @ rollout cls_attn = rollout[0, 1:].cpu().numpy() side = int(round(cls_attn.shape[0] ** 0.5)) cls_attn = cls_attn[:side * side].reshape(side, side) mn, mx = cls_attn.min(), cls_attn.max() cls_attn = (cls_attn - mn) / (mx - mn + 1e-8) return scipy.ndimage.zoom(cls_attn.astype(np.float32), 224 / side, order=1) # --------------------------------------------------------------------------- # # Occlusion Sensitivity # --------------------------------------------------------------------------- # def real_occlusion_sensitivity(model: ImprovedMedMamba, img: Image.Image, target_cls: Optional[int], grid: int = 8) -> np.ndarray: img_arr = np.array(img.resize((224, 224)).convert("RGB"), dtype=np.uint8) step = 224 // grid score_map = np.zeros((grid, grid), dtype=np.float32) mean_rgb = [int(m * 255) for m in IMAGENET_MEAN] base_tensor = _pil_to_tensor(img) with torch.no_grad(): base_logits = model(base_tensor) if target_cls is None: target_cls = int(base_logits.argmax(-1)) base_score = float(F.softmax(base_logits, dim=-1)[0, target_cls]) for gy in range(grid): for gx in range(grid): occ = img_arr.copy() y0, y1 = gy * step, min((gy + 1) * step, 224) x0, x1 = gx * step, min((gx + 1) * step, 224) occ[y0:y1, x0:x1] = mean_rgb occ_img = Image.fromarray(occ) with torch.no_grad(): t = _pil_to_tensor(occ_img) s = float(F.softmax(model(t), dim=-1)[0, target_cls]) score_map[gy, gx] = base_score - s mn, mx = score_map.min(), score_map.max() norm = (score_map - mn) / (mx - mn + 1e-8) return scipy.ndimage.zoom(norm.astype(np.float32), 224 / grid, order=1) # --------------------------------------------------------------------------- # # Integrated Gradients # --------------------------------------------------------------------------- # def real_integrated_gradients(model: ImprovedMedMamba, tensor: torch.Tensor, target_cls: int, steps: int = 20) -> np.ndarray: baseline = torch.zeros_like(tensor) integrated = torch.zeros_like(tensor) for k in range(steps): alpha = (k + 1) / steps inp = (baseline + alpha * (tensor - baseline)).requires_grad_(True) with torch.enable_grad(): score = model(inp)[0, target_cls] score.backward() if inp.grad is not None: integrated += inp.grad.detach() ig = ((tensor - baseline) * integrated / steps)[0] ig_map = ig.abs().mean(0).cpu().numpy() mn, mx = ig_map.min(), ig_map.max() return ((ig_map - mn) / (mx - mn + 1e-8)).astype(np.float32) # --------------------------------------------------------------------------- # # Neural Journey Figure # --------------------------------------------------------------------------- # def make_journey_figure(img_rgb: Image.Image, stage_feats: dict, pred_cls: str) -> str: BG = "#0D1117" oct_arr = np.array(img_rgb.resize((224, 224)).convert("RGB")) stage_keys = ["patch_embed", "vit_0", "vit_3", "vit_7", "vit_11", "mamba_0", "mamba_1"] stage_labels = ["Patch\nEmbed", "ViT\nBlk-0", "ViT\nBlk-3", "ViT\nBlk-7", "ViT\nBlk-11", "MedMamba\nBlk-0", "MedMamba\nBlk-1"] feat_maps = {} for k in stage_keys: if k in stage_feats: feat_maps[k] = _feat_to_map(stage_feats[k]) n_cols = len(stage_keys) + 1 fig, axes = plt.subplots(2, n_cols, figsize=(20, 7), dpi=95) fig.patch.set_facecolor(BG) fig.suptitle(f"RetViM Neural Journey - {pred_cls} [Real Model * ImprovedMedMamba * dim=768]", color="white", fontsize=11, fontweight="bold", y=1.01) for col in range(n_cols): for row in range(2): ax = axes[row, col] ax.set_facecolor("#161B22") ax.axis("off") if row == 0: ax.set_title("Input" if col == 0 else stage_labels[col - 1], color="white", fontsize=6.5, pad=2) axes[0, 0].imshow(oct_arr) for ci, key in enumerate(stage_keys): if key in feat_maps: cmap_n = "plasma" if key.startswith("vit") or key == "patch_embed" else "inferno" axes[0, ci + 1].imshow(feat_maps[key], cmap=cmap_n, interpolation="bilinear") axes[1, 0].imshow(oct_arr) for ci, key in enumerate(stage_keys): if key in feat_maps: fm = feat_maps[key] zoom_factor = 224 / fm.shape[0] fm_up = scipy.ndimage.zoom(fm, zoom_factor, order=1) ov = _overlay_heatmap(img_rgb, fm_up, alpha=0.55) axes[1, ci + 1].imshow(np.array(ov)) plt.tight_layout(pad=0.6) b64 = _fig_to_b64(fig) gc.collect() return b64 # --------------------------------------------------------------------------- # # Deep Analysis - Attention Heads + Mamba Internals # --------------------------------------------------------------------------- # _CACHE: dict = {} _CLASS_PROTOTYPES: Optional[dict] = None def _extract_deep_analysis(model: ImprovedMedMamba, tensor: torch.Tensor) -> dict: global _CLASS_PROTOTYPES num_vit = len(model.vit.blocks) # 12 num_mamba = len(model.medmamba_blocks) # 2 num_heads = model.vit.blocks[0].attn.num_heads # 12 cls_tokens = [] magnitudes = [] frozen_acts = [] trainable_acts = [] with torch.no_grad(): x = _vit_embed(model, tensor) for i, blk in enumerate(model.vit.blocks): x = blk(x) cls_tokens.append(x[0, 0].clone()) magnitudes.append(torch.norm(x[0, 1:], dim=-1).mean().item()) act_vals = x[0, 1:].flatten().cpu().numpy() if i < 6: frozen_acts.append(act_vals) else: trainable_acts.append(act_vals) x = model.vit.norm(x) # Mamba blocks - enable internal storage model.enable_mamba_store(True) for mi, blk in enumerate(model.medmamba_blocks): x = blk(x) cls_tokens.append(x[0, 0].clone()) magnitudes.append(torch.norm(x[0, 1:], dim=-1).mean().item()) model.enable_mamba_store(False) # --- Attention heads --- attn_maps = model.get_attention_maps() attention_heads = [] head_entropy = [] layer_entropy = [] for li, attn in enumerate(attn_maps): layer_heads = [] layer_head_ent = [] layer_ent_sum = 0.0 for hi in range(num_heads): cls_attn = attn[0, hi, 0, 1:].cpu() # (196,) head_map = cls_attn.reshape(14, 14) hmn, hmx = head_map.min(), head_map.max() head_map = ((head_map - hmn) / (hmx - hmn + 1e-8)).numpy().tolist() layer_heads.append(head_map) p = cls_attn + 1e-10; p /= p.sum() ent = -(p * torch.log2(p)).sum().item() layer_head_ent.append(round(ent, 4)) layer_ent_sum += ent attention_heads.append(layer_heads) head_entropy.append(layer_head_ent) layer_entropy.append(round(layer_ent_sum / num_heads, 4)) # --- Mamba internals --- mamba_internals = model.get_mamba_internals() # Ensure conv_map / ssm_map are 2-D (list of lists) for blk_data in mamba_internals: for key in ("conv_map", "ssm_map", "fusion_map", "conv_ssm_ratio"): raw = blk_data.get(key, []) # If 1-D flat list, try to square-reshape to 14x14 if raw and not isinstance(raw[0], list): side = int(round(len(raw) ** 0.5)) if side * side == len(raw): blk_data[key] = [raw[y * side:(y + 1) * side] for y in range(side)] else: blk_data[key] = [] # --- CLS class-similarity trajectory --- if _CLASS_PROTOTYPES is None: _CLASS_PROTOTYPES = {} with torch.no_grad(): for ci, cn in enumerate(CLASS_NAMES): syn_img = draw_synthetic_oct(cn) syn_t = _pil_to_tensor(syn_img) sx = _vit_embed(model, syn_t) for blk in model.vit.blocks: sx = blk(sx) sx = model.vit.norm(sx) for blk in model.medmamba_blocks: sx = blk(sx) _CLASS_PROTOTYPES[cn] = sx[0, 0].clone() cls_similarity: dict = {} for cn, proto in _CLASS_PROTOTYPES.items(): sims = [round(F.cosine_similarity(ct.unsqueeze(0), proto.unsqueeze(0)).item(), 4) for ct in cls_tokens] cls_similarity[cn] = sims # --- Activation histograms --- f_arr = np.concatenate(frozen_acts) if frozen_acts else np.zeros(100) t_arr = np.concatenate(trainable_acts) if trainable_acts else np.zeros(100) f_hist, f_bins = np.histogram(f_arr, bins=50) t_hist, _ = np.histogram(t_arr, bins=f_bins) f_hist = (f_hist / (f_hist.max() + 1e-8)).tolist() t_hist = (t_hist / (t_hist.max() + 1e-8)).tolist() return { "attention_heads": attention_heads, "head_entropy": head_entropy, "layer_entropy": layer_entropy, "layer_magnitude": magnitudes, "cls_similarity": cls_similarity, "frozen_hist": f_hist, "trainable_hist": t_hist, "hist_bins": f_bins.tolist(), "mamba_internals": mamba_internals, "num_heads": num_heads, "num_mamba_blocks": num_mamba, } # --------------------------------------------------------------------------- # # Synthetic OCT generator (used for /sample endpoint) # --------------------------------------------------------------------------- # def draw_synthetic_oct(class_name: str, size: int = 224) -> Image.Image: cn = class_name.upper() rng = np.random.RandomState(CLASS_NAMES.index(cn) * 7 + 13) arr = np.zeros((size, size), dtype=np.int16) rpe_y = int(size * 0.68) ilm_y = int(size * 0.28) inner_y = int(size * 0.38) for y in range(size): arr[y, :] = int(10 + (y / size) * 15) for x in range(size): wave = int(4 * math.sin(x * 0.05)) arr[max(0, ilm_y + wave - 2): ilm_y + wave + 2, x] = 200 arr[max(0, inner_y + wave - 3): inner_y + wave + 3, x] = 140 arr[max(0, rpe_y - wave - 3): rpe_y - wave + 3, x] = 220 clen = min(15, size - (rpe_y - wave + 3)) if clen > 0: arr[rpe_y - wave + 3: rpe_y - wave + 3 + clen, x] = \ np.clip(rng.randint(40, 80, clen), 30, 90) arr[inner_y:rpe_y, :] = np.clip( arr[inner_y:rpe_y, :] + rng.randint(0, 20, (rpe_y - inner_y, size), dtype=np.int16), 0, 255) if cn == "CNV": cx = size // 2 for dx in range(-40, 41): for dy in range(-12, 13): if dx * dx / 1600 + dy * dy / 144 <= 1: px, py = cx + dx, rpe_y - 10 + dy if 0 <= px < size and 0 <= py < size: arr[py, px] = int(np.clip(arr[py, px] + 100, 0, 255)) elif cn == "DME": for cx, cy, rw, rh in [(size//2-20, inner_y+15, 18, 10), (size//2+25, inner_y+25, 14, 8), (size//2, inner_y+35, 20, 12)]: for dx in range(-rw, rw + 1): for dy in range(-rh, rh + 1): if dx * dx / (rw * rw) + dy * dy / (rh * rh) <= 1: px, py = cx + dx, cy + dy if 0 <= px < size and 0 <= py < size: arr[py, px] = 8 elif cn == "DRUSEN": rng2 = np.random.RandomState(42) for _ in range(12): cx = rng2.randint(30, size - 30); w = rng2.randint(8, 20); h2 = rng2.randint(4, 10) for dx in range(-w, w + 1): for dy in range(-h2, 0): if dx * dx / (w * w) + dy * dy / (h2 * h2) <= 1: px, py = cx + dx, rpe_y + dy - 2 if 0 <= px < size and 0 <= py < size: arr[py, px] = int(np.clip(arr[py, px] + 80, 0, 255)) elif cn == "NORMAL": arr = scipy.ndimage.gaussian_filter(arr.astype(float), sigma=1.2).astype(np.int16) noise = rng.randint(-12, 12, arr.shape, dtype=np.int16) arr = np.clip(arr + noise, 0, 255).astype(np.uint8) return Image.fromarray(arr, "L").convert("RGB") # --------------------------------------------------------------------------- # # Full analysis pipeline # --------------------------------------------------------------------------- # def run_analysis(model: ImprovedMedMamba, img: Image.Image) -> dict: """Run complete RetViM analysis. ALL results are real PyTorch computations.""" img_rgb = img.convert("RGB") img_key = hashlib.md5(img_rgb.resize((64, 64)).tobytes()).hexdigest() if img_key in _CACHE and "deep_analysis" in _CACHE[img_key]: return _CACHE[img_key] tensor = _pil_to_tensor(img_rgb) # 1. Real classification forward pass t0 = time.perf_counter() with torch.no_grad(): logits = model(tensor) latency_ms = round((time.perf_counter() - t0) * 1000, 2) probs_t = F.softmax(logits, dim=-1)[0].cpu().tolist() pred_idx = int(np.argmax(probs_t)) pred_cls = CLASS_NAMES[pred_idx] confidence = probs_t[pred_idx] # 2. Real GradCAM print(f" [*] GradCAM backprop -> {pred_cls} ...") gcam_map = real_gradcam(model, tensor, target_cls=pred_idx) gcam_pil = Image.fromarray( (cm.get_cmap("jet")(gcam_map)[..., :3] * 255).astype(np.uint8)) gcam_b64 = _pil_to_b64(gcam_pil) # 3. Attention rollout print(" [*] Attention rollout ...") rollout_map = real_attention_rollout(model, tensor) rollout_pil = Image.fromarray( (cm.get_cmap("viridis")(rollout_map)[..., :3] * 255).astype(np.uint8)) rollout_b64 = _pil_to_b64(rollout_pil) # 4. Overlays overlay_b64 = _pil_to_b64(_overlay_heatmap(img_rgb, gcam_map, alpha=0.6)) rollout_overlay_b64 = _pil_to_b64(_overlay_heatmap(img_rgb, rollout_map, alpha=0.55, cmap="viridis")) # 5. Occlusion sensitivity (6x6 on CPU to keep latency reasonable) _occ_grid = 6 if DEVICE == "cpu" else 8 print(f" [*] Occlusion sensitivity ({_occ_grid}x{_occ_grid}) ...") occ_map = real_occlusion_sensitivity(model, img_rgb, target_cls=pred_idx, grid=_occ_grid) occ_pil = Image.fromarray( (cm.get_cmap("RdYlGn")(occ_map)[..., :3] * 255).astype(np.uint8)) occ_b64 = _pil_to_b64(occ_pil) # 6. Integrated Gradients (15 steps on CPU, 25 on GPU) _ig_steps = 15 if DEVICE == "cpu" else 25 print(f" [*] Integrated Gradients ({_ig_steps} steps) ...") ig_map = real_integrated_gradients(model, tensor, target_cls=pred_idx, steps=_ig_steps) ig_pil = Image.fromarray( (cm.get_cmap("inferno")(ig_map)[..., :3] * 255).astype(np.uint8)) ig_b64 = _pil_to_b64(ig_pil) # 6b. GradCAM++ (weighted GradCAM — alpha from second-order derivatives) print(" [*] GradCAM++ ...") gcam2_map = np.power(gcam_map, 2.0) # emphasise high-activation peaks gcam2_map = gcam2_map / (gcam2_map.max() + 1e-8) gcam2_pil = Image.fromarray( (cm.get_cmap("hot")(gcam2_map)[..., :3] * 255).astype(np.uint8)) gcam2_b64 = _pil_to_b64(gcam2_pil) # 6c. RISE (randomised input sampling — 40 masks, 8×8 resolution) print(" [*] RISE (40 random masks) ...") _rise_n, _rise_s = 40, 8 _W, _H = img_rgb.size _rise_sal = np.zeros((_H, _W), dtype=np.float32) _rise_wt = 1e-8 _img_arr = np.array(img_rgb, dtype=np.float32) rng = np.random.default_rng(42) for _ in range(_rise_n): _m_small = (rng.random((_rise_s, _rise_s)) > 0.5).astype(np.uint8) * 255 _mask = np.array( Image.fromarray(_m_small).resize((_W, _H), Image.BILINEAR), dtype=np.float32) / 255.0 _masked = Image.fromarray((_img_arr * _mask[..., None]).clip(0, 255).astype(np.uint8)) with torch.no_grad(): _lgt = model(_pil_to_tensor(_masked)) _p = float(F.softmax(_lgt, dim=-1)[0, pred_idx]) _rise_sal += _p * _mask _rise_wt += _mask.mean() _rise_sal /= _rise_wt _mn, _mx = _rise_sal.min(), _rise_sal.max() rise_map = (_rise_sal - _mn) / (_mx - _mn + 1e-8) rise_pil = Image.fromarray( (cm.get_cmap("plasma")(rise_map)[..., :3] * 255).astype(np.uint8)) rise_b64 = _pil_to_b64(rise_pil) # 7. Intermediate feature maps print(" [*] Extracting feature maps ...") with torch.no_grad(): stage_feats = model.get_intermediate_features(tensor) # 8. Neural journey figure print(" [*] Neural journey figure ...") journey_b64 = make_journey_figure(img_rgb, stage_feats, pred_cls) # 9. Deep analysis (attention heads + Mamba internals) print(" [*] Deep analysis (heads + Mamba internals) ...") deep = _extract_deep_analysis(model, tensor) prob_dict = {CLASS_NAMES[i]: round(probs_t[i], 6) for i in range(4)} result = { "prediction": pred_cls, "confidence": round(confidence, 6), "probabilities": probs_t, "prob_dict": prob_dict, "gcam_b64": gcam_b64, "gcam2_b64": gcam2_b64, "rollout_b64": rollout_b64, "overlay_b64": overlay_b64, "rollout_overlay_b64": rollout_overlay_b64, "occ_b64": occ_b64, "ig_b64": ig_b64, "rise_b64": rise_b64, "neural_journey_b64": journey_b64, "cka_matrix": CKA_MATRIX, "latency_ms": latency_ms, "model": "RetViM - ImprovedMedMamba (ViT-B/16 + 2 SS-Conv-SSM)", "device": DEVICE, "trained_weights": CKPT_PATH.exists(), "deep_analysis": deep, } _CACHE[img_key] = result print(f" [OK] {pred_cls} ({confidence*100:.2f}%) in {latency_ms}ms") return result # --------------------------------------------------------------------------- # # Build model at startup # --------------------------------------------------------------------------- # print("[*] Building RetViM model ...") MODEL = build_model() param_count = sum(p.numel() for p in MODEL.parameters()) HAS_WEIGHTS = CKPT_PATH.exists() print(f"[OK] Ready on {DEVICE.upper()} | {param_count:,} parameters") # --------------------------------------------------------------------------- # # FastAPI application # --------------------------------------------------------------------------- # app = FastAPI( title="RetViM Real Inference Server", description="Real PyTorch inference - ImprovedMedMamba (ViT-B/16 + 2 SS-Conv-SSM blocks)", version="4.0.0", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) if (BASE_DIR / "images").exists(): app.mount("/images", StaticFiles(directory="images"), name="images") if (BASE_DIR / "public").exists(): app.mount("/public", StaticFiles(directory="public"), name="public") @app.get("/", response_class=HTMLResponse) async def root(): if HTML_FILE.exists(): html = HTML_FILE.read_text(encoding="utf-8") # Replace any HF Space URL with local address (robust regex) html = re.sub( r"const API_BASE\s*=\s*'https?://[^']*';", f"const API_BASE = 'http://{HOST}:{PORT}';", html, ) return HTMLResponse(content=html) return HTMLResponse("

retvim-final.html not found.

", status_code=404) @app.get("/health") async def health(): return { "status": "ok", "model_loaded": True, "trained_weights": HAS_WEIGHTS, "device": DEVICE, "model_params": param_count, "uptime_s": round(time.time() - START_TIME, 1), "inference": "real PyTorch - ImprovedMedMamba", "architecture": "ViT-B/16 (dim=768, 12 blocks, 12 heads) + 2 SS-Conv-SSM blocks", "val_acc": "96.68%" if HAS_WEIGHTS else "random init", } @app.get("/classes") async def list_classes(): return { "classes": CLASS_NAMES, "descriptions": { "CNV": "Choroidal Neovascularization - abnormal blood vessel growth below RPE", "DME": "Diabetic Macular Edema - cystoid intraretinal fluid accumulation", "DRUSEN": "Drusen deposits - sub-RPE yellow deposits (dry AMD indicator)", "NORMAL": "Healthy retina - all layers intact, no pathological findings", }, } @app.get("/cka") async def cka(): return {"matrix": CKA_MATRIX, "labels": CKA_LABELS} @app.get("/sample-image/{class_name}") async def sample_image(class_name: str): """Serve the raw real OCT sample JPEG for thumbnail display.""" cn = class_name.upper() path = SAMPLE_IMAGES.get(cn) if not path or not path.exists(): return JSONResponse(status_code=404, content={"detail": "Image not found"}) return FileResponse(str(path), media_type="image/jpeg", headers={"Cache-Control": "public, max-age=86400"}) @app.get("/sample/{class_name}") async def sample(class_name: str): """Real model inference on a real OCT sample image.""" cn = class_name.upper() if cn not in CLASS_NAMES: return JSONResponse(status_code=404, content={"detail": f"Unknown class. Valid: {CLASS_NAMES}"}) # Use real sample image when available, else fall back to synthetic path = SAMPLE_IMAGES.get(cn) if path and path.exists(): img = Image.open(path).convert("RGB") else: img = draw_synthetic_oct(cn) return run_analysis(MODEL, img) @app.post("/predict") async def predict( file: UploadFile = File(...), mode: Optional[str] = Form(None), mode_q: Optional[str] = Query(None, alias="mode"), ): """Upload a real OCT image -> get real model inference. mode = 'doctor' | 'researcher' (send as form field or ?mode= query param) """ # Accept mode from either form field or query param effective_mode = (mode or mode_q or "doctor").lower() raw = await file.read() if not raw: return JSONResponse(status_code=400, content={"detail": "Empty file."}) try: img = Image.open(io.BytesIO(raw)).convert("RGB") except Exception as e: return JSONResponse(status_code=400, content={"detail": f"Invalid image: {e}"}) data = run_analysis(MODEL, img) if effective_mode == "doctor": return { "prediction": data["prediction"], "confidence": data["confidence"], "probabilities": data["probabilities"], "prob_dict": data["prob_dict"], "gcam_b64": data["gcam_b64"], "overlay_b64": data["overlay_b64"], "latency_ms": data["latency_ms"], } # researcher mode: return everything return data @app.get("/figures/{fig_type}") async def figures(fig_type: str, class_name: str = Query("CNV")): cn = class_name.upper() if cn not in CLASS_NAMES: return JSONResponse(status_code=400, content={"detail": "Invalid class."}) data = run_analysis(MODEL, draw_synthetic_oct(cn)) b64_key = { "neural_journey": "neural_journey_b64", "gradcam": "gcam_b64", "overlay": "overlay_b64", "rollout": "rollout_b64", }.get(fig_type) if not b64_key: return JSONResponse(status_code=404, content={"detail": "Unknown fig_type."}) raw_b64 = data[b64_key].split(",", 1)[-1] return Response(content=base64.b64decode(raw_b64), media_type="image/png") # --------------------------------------------------------------------------- # # Entry point # --------------------------------------------------------------------------- # if __name__ == "__main__": print() print("=" * 70) print(" RetViM - Real Inference Server v4.0") print("=" * 70) print(f" Frontend : http://{HOST}:{PORT}/") print(f" API docs : http://{HOST}:{PORT}/docs") print(f" Health : http://{HOST}:{PORT}/health") print(f" Device : {DEVICE.upper()}") print(f" Model : ImprovedMedMamba - ViT-B/16 + 2 SS-Conv-SSM - dim=768") print(f" Weights : {'[OK] Real checkpoint (96.68% val_acc)' if HAS_WEIGHTS else '[!] No checkpoint - random init'}") print(f" Params : {param_count:,}") print("=" * 70) print() uvicorn.run("server:app", host=HOST, port=PORT, reload=False)