Add E2E-VarNet reconstruction API

app.py

@@ -0,0 +1,203 @@
+"""
+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 time
+from pathlib import Path
+
+import gradio as gr
+import h5py
+import numpy as np
+import torch
+from fastapi import File, HTTPException, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+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))
+
+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()
+    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:
+        kspace = f["kspace"][0]           # first slice: [H, W] complex or [coils, H, W]
+        max_value = float(f.attrs.get("max", np.abs(kspace).max()))
+
+    mask_func = EquispacedMaskFractionFunc(
+        center_fractions=CENTER_FRACTIONS,
+        accelerations=ACCELERATIONS,
+        seed=MASK_SEED,
+    )
+    transform = T.VarNetDataTransform(mask_func=mask_func)
+
+    if kspace.ndim == 2:
+        kspace_np = kspace[np.newaxis, ..., np.newaxis]              # [1, H, W, 1] complex
+        kspace_np = np.concatenate([kspace_np.real, kspace_np.imag], axis=-1)  # [1, H, W, 2]
+    else:
+        kspace_np = kspace
+
+    sample = transform(kspace_np, None, None, None, max_value)
+    return sample.masked_kspace, sample.mask, max_value
+
+
+def _run_model(model: VarNet, mk: torch.Tensor, mask: torch.Tensor, mv: float) -> torch.Tensor:
+    """One forward pass → [H, W] cropped reconstruction."""
+    mk_   = mk   if mk.dim()   == 5 else mk.unsqueeze(0).unsqueeze(0)
+    mask_ = mask.bool() if mask.dim() == 5 else mask.bool().unsqueeze(0).unsqueeze(0)
+    with torch.no_grad():
+        out = model(mk_, mask_, torch.tensor(mv, device=DEVICE))
+    return center_crop(out.squeeze(0), CROP_SIZE)
+
+
+def _infer(raw: bytes) -> dict:
+    """Full inference: two checkpoints → mean recon + uncertainty + metrics."""
+    t0 = time.perf_counter()
+
+    mk, mask, max_value = _load_slice(raw)
+
+    img_best  = _run_model(MODEL_BEST,  mk, mask, max_value)
+    img_final = _run_model(MODEL_FINAL, mk, mask, max_value)
+
+    # 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)).squeeze(0), CROP_SIZE
+    )
+    def _w(t): return t.unsqueeze(0).unsqueeze(0).numpy()
+    ssim_val = float(fastmri_ssim(_w(mean_img), _w(zf_img), maxval=max_value))
+    psnr_val = float(fastmri_psnr(_w(mean_img), _w(zf_img), maxval=max_value))
+
+    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×]."
+    ),
+)
+
+# ── Mount REST API on the Gradio FastAPI app ──────────────────────────────────
+# These endpoints are called by the React frontend — same format as the local FastAPI server.
+
+demo.app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@demo.app.get("/health")
+def health():
+    return {"status": "ok", "model_loaded": True, "device": "cpu"}
+
+@demo.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}")
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+gradio>=4.40.0
+torch
+h5py
+numpy
+Pillow
+huggingface_hub
+git+https://github.com/facebookresearch/fastMRI.git@91f2df4711adbb6d643df1810f234e4abcf5881b