| """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") |
|
|
| |
| |
| sys.path.insert(0, str(args.app_files)) |
| os.chdir(args.app_files) |
| 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) |
|
|
| |
| 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) |
|
|
| |
| model = models.get_model().to(device).eval() |
| blob = torch.load(args.checkpoint, map_location="cpu", weights_only=False) |
| 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] |
|
|
| generated, reconstructed = {}, {} |
| start = time.time() |
| saved = torch.load(args.generated, map_location="cpu", weights_only=False) if args.generated else {} |
| 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) |
|
|
| |
| 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": {}} |
|
|
| |
| 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() |
|
|
