| """Score extracted autoencoders on held-out scans, in their own grid and back in the original 1 mm space. |
| |
| Each model is run through the geometry it was trained with: ``squash`` (resize the whole scan to the |
| grid, the chain before resampling) or ``spacing`` (resample to millimetres, then centre crop or pad). |
| The reconstruction is de-normalised and mapped back onto the scan's own 240x240x155 grid with MONAI's |
| inverse transforms, then scored against the untouched scan — so both models answer one question: how |
| much of the original brain survives the round trip, resampling included. |
| |
| python compare_autoencoders.py <autoencoder app_files> <images dir> <held-out cases file> \ |
| NAME:GEOMETRY:CHECKPOINT:CONFIG_YAML [...] |
| |
| CHECKPOINT ``none`` skips the model: the geometry's round trip alone, the ceiling any autoencoder on |
| that grid is held under. |
| """ |
|
|
| import json |
| import os |
| import statistics |
| import sys |
| from pathlib import Path |
|
|
| import monai.transforms as mt |
| import numpy as np |
| import torch |
| import yaml |
| from monai.networks.nets import AutoencoderKL |
|
|
| app, images, cases_file = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3]) |
| sys.path.insert(0, str(app)) |
| from trainer import foreground_ssim |
|
|
| cases = Path(cases_file).read_text().split() |
| |
| device = torch.device(os.environ.get("DEVICE", "cuda")) |
|
|
|
|
| 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(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 |
|
|
|
|
| results = {} |
| per_volume = {} |
| for spec in sys.argv[4:]: |
| name, geometry, checkpoint, config = spec.split(":") |
| document = yaml.safe_load(Path(config).read_text()) |
| chain = geometry_chain(geometry, document) |
| model = None if checkpoint == "none" else load(checkpoint, document) |
| own, original, paths = [], [], [] |
| for case in cases: |
| for path in sorted(images.glob(f"*/scans/input_*_{case}.nii.gz")): |
| data = chain({"image": str(path)}) |
| grid = data["image"] |
| mean, std = grid.mean(), grid.std() |
| with torch.no_grad(), torch.autocast(device_type=device.type, enabled=device.type == "cuda"): |
| x = ((grid - mean) / std).unsqueeze(0).to(device) |
| reconstruction = x if model is None else model(x)[0] |
| reconstruction = reconstruction.float() |
| own.append(foreground_ssim(reconstruction, x.float())) |
| data["image"] = grid.clone().copy_(reconstruction[0].cpu() * std + mean) |
| |
| back = torch.as_tensor(chain.inverse(data)["image"]).unsqueeze(0).to(device).float() |
| truth = torch.as_tensor(mt.LoadImage(image_only=True)(str(path))).float()[None, None].to(device) |
| assert back.shape == truth.shape, (back.shape, truth.shape) |
| back = back * (truth > 0) |
| original.append(foreground_ssim(back, truth)) |
| paths.append(path.name) |
| results[name] = (own, original) |
| per_volume[name] = [{"volume": p, "own": o, "original": r} for p, o, r in zip(paths, own, original)] |
| print(f"{name:<6} {len(own)} volumes own-grid SSIM {statistics.mean(own):.4f} ± {statistics.stdev(own):.4f} " |
| f"original 1 mm SSIM {statistics.mean(original):.4f} ± {statistics.stdev(original):.4f}", flush=True) |
| models = [name for name, spec in zip(results, sys.argv[4:]) if not spec.split(":")[2] == "none"] |
| if len(models) == 2: |
| first, second = (np.array(results[name][1]) for name in models) |
| diff = second - first |
| print(f"{models[1]} - {models[0]} at 1 mm: mean {diff.mean():+.4f}, better on {int((diff > 0).sum())} of {len(diff)} volumes") |
| if os.environ.get("OUT"): |
| Path(os.environ["OUT"]).write_text(json.dumps(per_volume, indent=1)) |
|
|
