"""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 --autoencoder-app-files \ --checkpoint --weights --images \ --train-frame --heldout-cases --out-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()