evaluation/evaluate_ldm.py

"""Score a trained latent diffusion model: realism, diversity and whether it obeys the sequence it is asked for.

Built from the tutorial's own code (its models.py, transforms, scheduler settings and inferer), so it
scores the model the tutorial trains, in the grid it trains on. Three comparisons, per MR sequence:

* Realism: Gaussian-kernel MMD against held-out real volumes, and k-NN precision and recall against training
  ones, on MedicalNet ResNet-50 features (global-average-pooled final map, the network behind the
  autoencoder's perceptual loss); see ldm_metrics.py. Two references bracket it: real against real (the
  floor) and autoencoder reconstructions of the held-out scans (the ceiling: the LDM decodes through it).
* Diversity: MS-SSIM between pairs of generated volumes of one sequence, beside the same for real pairs.
  Much higher than real means the model repeats itself.
* Conditioning: a logistic regression on every MedicalNet stage, fit on training-real volumes to tell the
  four sequences apart; its accuracy on held-out real and on reconstructions is the check on the probe,
  and on generated volumes it is the share that look like the sequence they were asked for.

python evaluate_ldm.py --app-files <ldm app_files> --autoencoder-app-files <autoencoder app_files> \
    --checkpoint <ldm FL_global_model.pt> --weights <medicalnet .pth> --images <images dir> \
    --train-frame <dataframe_train.csv> --heldout-cases <cases file> --out-dir <dir> [--steps 1000]
"""

import argparse
import importlib.util
import itertools
import json
import os
import random
import sys
import time
from pathlib import Path

import numpy as np
import torch


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--app-files", type=Path, required=True)
    parser.add_argument("--autoencoder-app-files", type=Path, required=True)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--weights", type=Path, required=True)
    parser.add_argument("--images", type=Path, required=True)
    parser.add_argument("--train-frame", type=Path, required=True)
    parser.add_argument("--heldout-cases", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--per-sequence", type=int, default=40, help="generated volumes per sequence")
    parser.add_argument("--train-real", type=int, default=160, help="training volumes per sequence (probe + floor)")
    parser.add_argument("--heldout-limit", type=int, default=0, help="first N held-out cases only (0 = all)")
    parser.add_argument("--steps", type=int, default=1000, help="reverse diffusion steps")
    parser.add_argument("--batch", type=int, default=8)
    parser.add_argument("--pairs", type=int, default=200, help="MS-SSIM pairs per sequence and set")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--generated", type=Path, help="reuse the generated volumes a previous run saved (no sampling)")
    return parser.parse_args()


def import_file(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def whole_storage(saved: torch.Tensor) -> torch.Tensor:
    """Every volume in a saved tensor's storage: a saved slice (the first runs' ``[:8]``) still holds them all."""
    t = saved.as_subclass(torch.Tensor)
    count = t.untyped_storage().nbytes() // (t[0].numel() * t.element_size())
    return torch.empty(0, dtype=t.dtype).set_(t.untyped_storage(), 0, (count, *t.shape[1:]))


def main() -> None:
    args = parse_args()
    for key, value in vars(args).items():
        if isinstance(value, Path):
            setattr(args, key, value.resolve())
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    args.out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # The LDM tutorial's flat modules first: its models.py, transforms.py and trainer.py shadow the
    # autoencoder tutorial's files of the same names.
    sys.path.insert(0, str(args.app_files))
    os.chdir(args.app_files)  # models.py reads config.json and autoencoder_config.yaml beside it
    import models
    from latent_utils import build_inferer
    from modality import one_hot_condition
    from monai.metrics import MultiScaleSSIMMetric
    from monai.networks.schedulers import DDPMScheduler
    from transforms import get_brain_mri_transforms

    medicalnet = import_file("medicalnet_perceptual", args.autoencoder_app_files / "medicalnet_perceptual.py")
    metrics = import_file("ldm_metrics", Path(__file__).resolve().with_name("ldm_metrics.py"))
    config = json.loads((args.app_files / "config.json").read_text())
    modalities = config["MODALITIES"]
    chain = get_brain_mri_transforms(is_validation=True)

    # --- Real volumes: held-out patients, and a sample of training ones ---
    folders = {p.name.removeprefix("label_").removesuffix(".nii.gz"): p.parent for p in args.images.glob("*/scans/label_*.nii.gz")}
    heldout_cases = args.heldout_cases.read_text().split()
    if args.heldout_limit:
        heldout_cases = heldout_cases[: args.heldout_limit]
    train_accessions = [line.split(",")[0] for line in args.train_frame.read_text().splitlines()[1:]]
    train_folders = [args.images / accession / "scans" for accession in train_accessions]

    def load(paths: list[Path]) -> torch.Tensor:
        return torch.stack([chain({"image": str(p)})["image"].as_tensor() for p in paths])

    heldout = {m: load([folders[c] / f"input_{m}_{c}.nii.gz" for c in heldout_cases]) for m in modalities}
    train_real = {}
    for m in modalities:
        candidates = sorted(p for f in train_folders for p in f.glob(f"input_{m}_*.nii.gz"))
        train_real[m] = load(random.sample(candidates, min(args.train_real, len(candidates))))
    print(f"real: {sum(len(v) for v in heldout.values())} held-out, {sum(len(v) for v in train_real.values())} training", flush=True)

    # --- The trained model, exactly as the tutorial builds it ---
    model = models.get_model().to(device).eval()
    blob = torch.load(args.checkpoint, map_location="cpu", weights_only=False)  # FLARE envelope; our own run
    model.load_state_dict(blob["model"] if isinstance(blob.get("model"), dict) else blob)
    scheduler = DDPMScheduler(num_train_timesteps=1000, schedule="scaled_linear_beta", clip_sample=False,
                              prediction_type="epsilon", beta_start=0.0015, beta_end=0.0195)
    scheduler.set_timesteps(num_inference_steps=args.steps)
    inferer, spatial = build_inferer(model, scheduler, config["spatial_shape"], heldout[modalities[0]][:4],
                                     device, config.get("LATENT_SCALE_FACTOR"))
    latent_shape = [model.autoencoder.latent_channels, *spatial]  # channels first, as the trainer's get_inferer

    generated, reconstructed = {}, {}
    start = time.time()
    saved = torch.load(args.generated, map_location="cpu", weights_only=False) if args.generated else {}  # our own output
    with torch.no_grad():
        for index, m in enumerate(modalities):
            if saved:
                generated[m] = whole_storage(saved[m])[: args.per_sequence]
                reconstructed[m] = torch.cat([model.autoencoder(heldout[m][i : i + args.batch].to(device))[0].float().cpu()
                                              for i in range(0, len(heldout[m]), args.batch)])
                continue
            batches = []
            for first in range(0, args.per_sequence, args.batch):
                count = min(args.batch, args.per_sequence - first)
                batches.append(inferer.sample(
                    input_noise=torch.randn([count, *latent_shape], device=device),
                    conditioning=one_hot_condition(torch.full((count,), index), len(modalities), device),
                    diffusion_model=model.diffusion_model, autoencoder_model=model.autoencoder,
                    scheduler=scheduler, verbose=False).float().cpu())
            generated[m] = torch.cat(batches)
            reconstructed[m] = torch.cat([model.autoencoder(heldout[m][i : i + args.batch].to(device))[0].float().cpu()
                                          for i in range(0, len(heldout[m]), args.batch)])
            print(f"{m}: {len(generated[m])} generated, {time.time() - start:.0f} s", flush=True)

    # --- Features ---
    backbone = medicalnet.MedicalNetPerceptualLoss(args.weights).model.to(device)

    def features(volumes: torch.Tensor) -> tuple[np.ndarray, np.ndarray]:
        """Per volume: the final stage pooled (realism), and every stage pooled (the sequence probe)."""
        final, probe_rows = [], []
        with torch.no_grad():
            for volume in volumes:
                stages = backbone(medicalnet._intensity_normalisation(volume[None]).to(device))
                pooled = [s.mean(dim=(2, 3, 4)).cpu().numpy()[0] for s in stages]
                final.append(pooled[-1])
                probe_rows.append(np.concatenate(pooled))
        return np.stack(final), np.stack(probe_rows)

    sets = {"heldout": heldout, "train_real": train_real, "reconstructed": reconstructed, "generated": generated}
    both = {name: {m: features(volumes[m]) for m in modalities} for name, volumes in sets.items()}
    feats = {name: {m: f[0] for m, f in per.items()} for name, per in both.items()}
    probe_feats = {name: {m: f[1] for m, f in per.items()} for name, per in both.items()}
    np.savez(args.out_dir / "features.npz", **{f"final/{n}/{m}": f for n, per in feats.items() for m, f in per.items()},
             **{f"probe/{n}/{m}": f for n, per in probe_feats.items() for m, f in per.items()})

    results = {"steps": args.steps, "per_sequence": args.per_sequence, "heldout_cases": len(heldout_cases),
               "latent_scale_factor": float(inferer.scale_factor), "realism": metrics.realism(feats),
               "conditioning": metrics.conditioning(probe_feats), "diversity": {}}

    # --- Diversity ---
    ms_ssim = MultiScaleSSIMMetric(spatial_dims=3, data_range=1.0, kernel_size=4)

    def unit(v: torch.Tensor) -> torch.Tensor:
        """To [0, 1] between the 1st and 99.5th percentiles: generated volumes have stray voxels below the
        background, which a min-max scaling would map to 0 instead of the background."""
        lo, hi = torch.quantile(v.flatten().float(), torch.tensor([0.01, 0.995]))
        return ((v - lo) / (hi - lo + 1e-8)).clamp(0, 1)

    for name in ("heldout", "generated"):
        results["diversity"][name] = {}
        for m in modalities:
            volumes = sets[name][m]
            pairs = list(itertools.combinations(range(len(volumes)), 2))
            pairs = random.sample(pairs, min(args.pairs, len(pairs)))
            scores = [float(ms_ssim(unit(volumes[i])[None], unit(volumes[j])[None]).item()) for i, j in pairs]
            results["diversity"][name][m] = (float(np.mean(scores)), float(np.std(scores)))

    (args.out_dir / "ldm_scores.json").write_text(json.dumps(results, indent=1))
    if not saved:
        torch.save({m: generated[m].as_subclass(torch.Tensor).clone() for m in modalities}, args.out_dir / "generated.pt")
    print(metrics.summary(results))
    print(json.dumps(results["diversity"], indent=1))


if __name__ == "__main__":
    main()