""" E2E-VarNet MRI Reconstruction — HuggingFace Space Two interfaces in one: • Gradio UI at / (visual demo on HF) • REST API at /reconstruct (called by React frontend) • Health at /health """ import base64 import io import os import time from pathlib import Path import gradio as gr import spaces import h5py import numpy as np import torch from fastapi import File, HTTPException, UploadFile from PIL import Image from huggingface_hub import hf_hub_download import fastmri from fastmri.data import transforms as T from fastmri.data.subsample import EquispacedMaskFractionFunc from fastmri.evaluate import ssim as fastmri_ssim, psnr as fastmri_psnr from fastmri.models import VarNet from fastmri.data.transforms import center_crop # ── Config ──────────────────────────────────────────────────────────────────── HF_REPO = "alyrraza/e2e-varnet-mri-reconstruction" CKPT_DIR = Path("checkpoints") CENTER_FRACTIONS = [0.08] ACCELERATIONS = [4] MASK_SEED = 42 # explicit seed — avoids run-to-run non-determinism CROP_SIZE = (320, 320) FLAG_THRESHOLD = 0.018 # calibrated from research (1.2× in-domain mean) DEVICE = torch.device("cpu") # ── Download & load models ──────────────────────────────────────────────────── CKPT_DIR.mkdir(exist_ok=True) def _download(filename: str) -> None: if not (CKPT_DIR / filename).exists(): print(f"Downloading {filename} from HF Hub …") hf_hub_download( repo_id=HF_REPO, filename=filename, local_dir=str(CKPT_DIR), token=os.environ.get("HF_TOKEN"), ) def _load_varnet(path: Path) -> VarNet: model = VarNet(num_cascades=4, chans=18, sens_chans=8) state = torch.load(path, map_location=DEVICE, weights_only=False) sd = state.get("model_state_dict", state) model.load_state_dict(sd) model.eval() bad = [n for n, p in model.named_parameters() if p.isnan().any() or p.isinf().any()] if bad: print(f"[FATAL] {path.name}: NaN/Inf in weights: {bad[:5]}") else: n_params = sum(p.numel() for p in model.parameters()) print(f"[OK] {path.name}: {n_params:,} params all finite") return model print("Downloading checkpoints …") _download("best_model.pt") # T4 epoch-21 (best val SSIM) _download("checkpoint_epoch_50.pt") # T4 epoch-50 (final) MODEL_BEST = _load_varnet(CKPT_DIR / "best_model.pt") MODEL_FINAL = _load_varnet(CKPT_DIR / "checkpoint_epoch_50.pt") print("Models ready.") # ── Inference helpers ───────────────────────────────────────────────────────── def _tensor_to_b64(t: torch.Tensor) -> str: arr = t.numpy().astype(np.float32) arr = (arr - arr.min()) / (arr.max() - arr.min() + 1e-8) img = Image.fromarray((arr * 255).astype(np.uint8)) buf = io.BytesIO() img.save(buf, format="PNG") return base64.b64encode(buf.getvalue()).decode() def _load_slice(data: bytes): """Parse one slice from a fastMRI .h5 file. Returns (masked_kspace, mask, max_value).""" with h5py.File(io.BytesIO(data), "r") as f: mid = len(f["kspace"]) // 2 kspace = f["kspace"][mid] max_value = float(f.attrs.get("max", np.abs(kspace).max())) # VarNetDataTransform at pinned fastMRI commit requires attrs dict # with padding_left / padding_right / encoding_size / recon_size attrs = dict(f.attrs) W = kspace.shape[-1] H = kspace.shape[-2] if kspace.ndim >= 2 else 320 attrs.setdefault("encoding_size", W) attrs.setdefault("recon_size", (H, W)) attrs.setdefault("padding_left", 0) # padding_right must be W (not 0) when there is no right padding. # apply_mask does mask[:, :, padding_right:] = 0 — so padding_right=0 # zeroes the ENTIRE mask, making masked_kspace all-zeros → NaN model output. pr = int(attrs.get("padding_right", W)) attrs["padding_right"] = W if pr == 0 else pr mask_func = EquispacedMaskFractionFunc( center_fractions=CENTER_FRACTIONS, accelerations=ACCELERATIONS, seed=MASK_SEED, ) transform = T.VarNetDataTransform(mask_func=mask_func) # Let the transform handle complex→real split via to_tensor internally if kspace.ndim == 2: kspace_np = kspace[np.newaxis, ...] # [H,W] complex → [1,H,W] complex (1 coil) else: kspace_np = kspace # [C,H,W] complex (multi-coil) sample = transform(kspace_np, None, None, attrs, "upload.h5", 0) # Return num_low_frequencies from sample — VarNet needs it for sensitivity estimation return sample.masked_kspace, sample.mask, sample.num_low_frequencies, max_value def _run_model(model: VarNet, mk: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: """One forward pass → [H, W] cropped reconstruction.""" mk_ = mk.float().unsqueeze(0) # [1,H,W,2] → [1,1,H,W,2] mask_ = mask.bool().unsqueeze(0) # [1,1,W,1] → [1,1,1,W,1] # NormUnet.norm in fastMRI divides by std with NO epsilon: (x-mean)/std # When kspace absmax ≈ 4e-4, IFFT std is ~1e-5 — tiny enough to cause NaN # on CPU float32. Pre-scaling to O(1) avoids this; VarNet output is linear # in kspace amplitude so multiplying back exactly restores the original scale. scale = float(mk_.abs().max()) + 1e-8 with torch.no_grad(): out = model(mk_ / scale, mask_) return center_crop(out.squeeze(0) * scale, CROP_SIZE) @spaces.GPU def _infer(raw: bytes) -> dict: """Full inference: two checkpoints → mean recon + uncertainty + metrics.""" t0 = time.perf_counter() mk, mask, num_low_freq, max_value = _load_slice(raw) print(f"[DEBUG] mk shape={tuple(mk.shape)} dtype={mk.dtype} " f"absmax={float(mk.abs().max()):.4e} isnan={bool(mk.isnan().any())}") print(f"[DEBUG] mask shape={tuple(mask.shape)} " f"num_low_freq={num_low_freq} max_value={max_value:.4e}") img_best = _run_model(MODEL_BEST, mk, mask) img_final = _run_model(MODEL_FINAL, mk, mask) print(f"[DEBUG] img_best has_nan={bool(img_best.isnan().any())} " f"max={float(img_best[~img_best.isnan()].max()) if not img_best.isnan().all() else 'ALL_NAN'}") print(f"[DEBUG] img_final has_nan={bool(img_final.isnan().any())}") # Ensemble stats stack = torch.stack([img_best, img_final]) # [2, H, W] mean_img = stack.mean(0) # [H, W] std_map = stack.std(0, unbiased=False) # [H, W] scalar = float(std_map.mean()) / max_value if max_value > 0 else float(std_map.mean()) # SSIM / PSNR vs zero-filled (no ground truth at inference time) zf_img = center_crop( fastmri.complex_abs(fastmri.ifft2c(mk.float())).squeeze(0), CROP_SIZE ) def _w(t): return t.unsqueeze(0).numpy() # [H,W] → [1,H,W] as fastmri_ssim expects ssim_val = float(np.asarray(fastmri_ssim(_w(mean_img), _w(zf_img), maxval=max_value)).flat[0]) psnr_val = float(np.asarray(fastmri_psnr(_w(mean_img), _w(zf_img), maxval=max_value)).flat[0]) return { "ssim": round(ssim_val, 4), "psnr": round(psnr_val, 2), "uncertainty_scalar": round(scalar, 6), "flagged_for_review": scalar > FLAG_THRESHOLD, "reconstruction_b64": _tensor_to_b64(mean_img), "uncertainty_map_b64": _tensor_to_b64(std_map), "inference_time_ms": round((time.perf_counter() - t0) * 1000, 1), "model_variant": "T4-ensemble-K2", } # ── Gradio UI ───────────────────────────────────────────────────────────────── def gradio_predict(file): if file is None: return "Upload a file first.", None, None with open(file, "rb") as f: raw = f.read() res = _infer(raw) label = ( f"SSIM: {res['ssim']} | PSNR: {res['psnr']} dB | " f"Uncertainty: {res['uncertainty_scalar']} | " f"{'⚠ OOD flagged' if res['flagged_for_review'] else '✓ In-distribution'} | " f"{res['inference_time_ms']:.0f} ms" ) recon = Image.open(io.BytesIO(base64.b64decode(res["reconstruction_b64"]))) umap = Image.open(io.BytesIO(base64.b64decode(res["uncertainty_map_b64"]))) return label, recon, umap demo = gr.Interface( fn=gradio_predict, inputs=gr.File(label="Upload fastMRI .h5 file", file_types=[".h5"]), outputs=[ gr.Textbox(label="Metrics"), gr.Image(label="Reconstruction"), gr.Image(label="Uncertainty Map"), ], title="E2E-VarNet MRI Reconstruction", description=( "Upload a fastMRI single-coil knee `.h5` file. " "The model reconstructs the image from 25 % of k-space (4× acceleration) " "and returns a checkpoint-ensemble uncertainty score. " "OOD shift: **1.54×** higher uncertainty on brain vs knee, " "bootstrap 95 % CI [1.35×, 1.74×]." ), ) # ── Launch first, then add REST routes on the real FastAPI app ──────────────── # In Gradio 5.x, demo.app before launch is a stub — routes must be added # AFTER launch() returns the real app instance. import threading # Gradio already adds CORSMiddleware allowing all origins — no need to re-add. # Routes must be registered on real_app AFTER launch() in Gradio 5.x. real_app, _, _ = demo.launch( server_name="0.0.0.0", server_port=7860, ssr_mode=False, prevent_thread_lock=True, ) @real_app.get("/health") def health(): return {"status": "ok", "model_loaded": True, "device": "cpu"} @real_app.post("/reconstruct") async def reconstruct(file: UploadFile = File(...)): if not (file.filename or "").endswith(".h5"): raise HTTPException(400, "File must be a .h5 fastMRI file") raw = await file.read() try: return _infer(raw) except Exception as exc: raise HTTPException(422, f"Inference failed: {exc}") # Block forever so HF Spaces keeps the container alive threading.Event().wait()