app.py

"""
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()