| """Same held-out cases through several autoencoders, compared where the SSIM scores are: the original scan. |
| |
| Each model runs in the geometry it was trained with; its reconstruction is de-normalised and mapped back |
| onto the scan's own 1 mm grid with MONAI's inverse transforms (as compare_autoencoders.py scores it). One |
| row per case; for each plane the original, then each model, cut through the tumour centre and shown in |
| one intensity window per row. Each model's foreground SSIM at 1 mm is printed on its axial panel. |
| |
| python render_same_cases.py <autoencoder app_files> <images dir> <out.png> CASE:SEQUENCE[,CASE:SEQUENCE...] \ |
| NAME:GEOMETRY:CHECKPOINT:CONFIG_YAML [...] |
| """ |
|
|
| import os |
| import sys |
| from pathlib import Path |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import monai.transforms as mt |
| import numpy as np |
| import torch |
| import yaml |
| from monai.data import MetaTensor |
| from monai.networks.nets import AutoencoderKL |
|
|
| app, images, out = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3]) |
| cases = [tuple(item.split(":")) for item in sys.argv[4].split(",")] |
| specs = [spec.split(":") for spec in sys.argv[5:]] |
| sys.path.insert(0, str(app)) |
| from trainer import foreground_ssim |
|
|
| device = torch.device(os.environ.get("DEVICE", "cpu")) |
| PLANES = (("axial", 2, False), ("coronal", 1, False), ("sagittal", 0, True)) |
| to_ras = mt.Orientation(axcodes="RAS") |
|
|
|
|
| def geometry_chain(geometry: str, document: dict) -> mt.Compose: |
| shape = document["training"]["spatial_shape"] |
| steps = [mt.LoadImaged(keys=["image"], image_only=True), mt.EnsureChannelFirstd(keys=["image"], channel_dim="no_channel"), |
| mt.Orientationd(keys=["image"], axcodes="RAS")] |
| if geometry == "squash": |
| steps.append(mt.Resized(keys=["image"], spatial_size=shape)) |
| else: |
| steps += [mt.Spacingd(keys=["image"], pixdim=document["training"]["voxel_spacing"], mode="bilinear"), |
| mt.ResizeWithPadOrCropd(keys=["image"], spatial_size=shape)] |
| return mt.Compose(steps) |
|
|
|
|
| def load_model(checkpoint: str, document: dict) -> AutoencoderKL: |
| model = AutoencoderKL(**document["autoencoder"]).to(device).eval() |
| state = torch.load(checkpoint, map_location="cpu", weights_only=True) |
| model.load_state_dict({k.removeprefix("autoencoder."): v for k, v in state.items()}) |
| return model |
|
|
|
|
| def back_to_original(path: Path, chain: mt.Compose, model: AutoencoderKL) -> torch.Tensor: |
| data = chain({"image": str(path)}) |
| grid = data["image"] |
| mean, std = grid.mean(), grid.std() |
| with torch.no_grad(): |
| reconstruction = model(((grid - mean) / std).unsqueeze(0).to(device))[0].float().cpu() |
| data["image"] = grid.clone().copy_(reconstruction[0] * std + mean) |
| return torch.as_tensor(chain.inverse(data)["image"]).float() |
|
|
|
|
| def cut(volume: np.ndarray, axis: int, index: int, flip_lr: bool) -> np.ndarray: |
| plane = np.rot90(np.take(volume, index, axis=axis)) |
| return np.fliplr(plane) if flip_lr else plane |
|
|
|
|
| models = [] |
| for name, geometry, checkpoint, config in specs: |
| document = yaml.safe_load(Path(config).read_text()) |
| models.append((name, geometry_chain(geometry, document), load_model(checkpoint, document))) |
|
|
| columns = ["original"] + [name for name, _, _ in models] |
| figure, axes = plt.subplots(len(cases), len(PLANES) * len(columns), |
| figsize=(1.9 * len(PLANES) * len(columns), 2.1 * len(cases)), squeeze=False) |
| figure.patch.set_facecolor("black") |
| for row, (case, sequence) in enumerate(cases): |
| folder = next(p.parent for p in images.glob(f"*/scans/label_{case}.nii.gz")) |
| path = folder / f"input_{sequence}_{case}.nii.gz" |
| truth = mt.LoadImage(image_only=True)(str(path))[None] |
| label = mt.LoadImage(image_only=True)(str(folder / f"label_{case}.nii.gz"))[None] |
| volumes, scores = [to_ras(truth)[0].numpy()], [None] |
| for name, chain, model in models: |
| back = back_to_original(path, chain, model) |
| scores.append(foreground_ssim((back * (truth > 0))[None].float(), torch.as_tensor(truth)[None].float())) |
| volumes.append(to_ras(MetaTensor(back, affine=truth.affine))[0].numpy()) |
| print(f"{case} {sequence} {name}: SSIM at 1 mm {scores[-1]:.3f}", flush=True) |
| tumour = to_ras(label)[0].numpy() > 0 |
| centre = [int(round(c.mean())) for c in np.nonzero(tumour)] |
| brain = volumes[0] > 0 |
| vmax = float(np.percentile(volumes[0][brain], 99.5)) |
| for p, (view, axis, flip_lr) in enumerate(PLANES): |
| mask = cut(brain, axis, centre[axis], flip_lr) |
| rows_, cols_ = np.nonzero(mask) |
| box = (slice(max(rows_.min() - 4, 0), rows_.max() + 5), slice(max(cols_.min() - 4, 0), cols_.max() + 5)) |
| for c, (title, volume, score) in enumerate(zip(columns, volumes, scores)): |
| handle = axes[row][p * len(columns) + c] |
| handle.imshow(cut(volume, axis, centre[axis], flip_lr)[box], cmap="gray", vmin=0, vmax=vmax) |
| handle.set_xticks([]) |
| handle.set_yticks([]) |
| for spine in handle.spines.values(): |
| spine.set_visible(False) |
| if row == 0: |
| handle.set_title(f"{title}\n{view}", color="0.85", fontsize=10) |
| if p == 0 and score is not None: |
| handle.text(0.03, 0.03, f"SSIM {score:.3f}", transform=handle.transAxes, color="yellow", fontsize=9, |
| bbox={"facecolor": "black", "edgecolor": "none", "alpha": 0.75, "pad": 2}) |
| if p == 0 and c == 0: |
| handle.set_ylabel(f"{case}\n{sequence}", color="0.85", fontsize=10) |
| figure.tight_layout() |
| figure.savefig(out, facecolor=figure.get_facecolor(), dpi=110) |
| print("wrote", out) |
|
|
