| """Generated #0-#2 per sequence from E5, E7, E10 and E11 beside three held-out patients, mid-axial. |
| |
| python four_runs.py <ldm app_files> <images dir> <heldout_cases.txt> <out.png> E5=<pt> E7=<pt> E10=<pt> E11=<pt> |
| """ |
|
|
| import sys |
| from pathlib import Path |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import torch |
|
|
| app_files, images, cases_file, out = map(Path, sys.argv[1:5]) |
| runs = dict(a.split("=", 1) for a in sys.argv[5:]) |
| sys.path.insert(0, str(app_files)) |
| from transforms import get_brain_mri_transforms |
|
|
| M = ["FLAIR", "T1w", "T1Gd", "T2w"] |
| N = 3 |
|
|
|
|
| def every(t: torch.Tensor) -> np.ndarray: |
| t = t.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:])).numpy()[:, 0] |
|
|
|
|
| gen = {} |
| for name, pt in runs.items(): |
| saved = torch.load(pt, map_location="cpu", weights_only=False) |
| gen[name] = {m: every(saved[m])[:N] for m in M} |
| chain = get_brain_mri_transforms(is_validation=True) |
| folders = {p.name.removeprefix("label_").removesuffix(".nii.gz"): p.parent for p in images.glob("*/scans/label_*.nii.gz")} |
| cases = cases_file.read_text().split()[:N] |
| gen["held-out real"] = { |
| m: [chain({"image": str(folders[c] / f"input_{m}_{c}.nii.gz")})["image"].as_tensor().numpy()[0] for c in cases] |
| for m in M} |
|
|
| groups = list(gen) |
| fig, axes = plt.subplots(len(M), len(groups) * N, figsize=(len(groups) * N * 1.05, len(M) * 1.25), facecolor="black") |
| for r, m in enumerate(M): |
| for g, name in enumerate(groups): |
| for i in range(N): |
| ax, v = axes[r, g * N + i], gen[name][m][i] |
| ax.imshow(np.rot90(v[:, :, v.shape[2] // 2]), cmap="gray", vmin=np.percentile(v, 1), |
| vmax=np.percentile(v, 99.5)) |
| ax.set_xticks([]), ax.set_yticks([]), ax.set_frame_on(False) |
| if r == 0 and i == 1: |
| ax.set_title(name, color="white", fontsize=11) |
| axes[r, 0].set_ylabel(m, color="white", fontsize=11) |
| plt.subplots_adjust(left=0.03, right=0.995, top=0.93, bottom=0.01, wspace=0.03, hspace=0.04) |
| plt.savefig(out, dpi=90, facecolor="black") |
| print("saved", out) |
|
|
