From 6d889b957828ba0027ef43672ff5d8bad1f86c15 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 28 Oct 2026 16:15:10 +0000 Subject: [PATCH] feat: shift the scan inside the patch and fill what lies outside it with EM-synthesised background A model trained on sagittal spine slabs has never been shown what lies beside them. nnU-Net centres the ~8 cm slab in a 128 mm patch and pads the rest with 0, so every training window contains spine and everything beyond +-4 cm is a flat slab labelled background. On whole-body axial NAKO VIBE the Ours setup then predicts vertebrae in muscle, gluteals, fat and bowel: vertebra instance precision 0.01 with ~100 false-positive components per subject, and 60% of the false-positive voxels more than 5 cm from the spine. That is missing context, not a contrast failure, and no intensity augmentation reaches it. `RandomFOVExtensionGPU` (GEO, placed right after `RandomAffineGPU`) does, per selected sample: * shifts image and mask by the same whole-voxel offset, up to `shift` of the patch per axis, so the spine reaches the window edge or leaves it; * fills everything that is not scan -- the exact-`pad_value` padding plus the vacated area -- SynthSeg-style: EM splits the patch's own background (label 0, beyond `margin` of any label) into 3-10 intensity regions, a majority filter and a random diffeomorphic field post-process the region map, warped copies of it are composited into the fill, and each region is rendered from its own mean and std with spatially correlated texture, then blurred as SynthSeg blurs. The fill's label is 0. Row selection is pinned per sample: without a `pipeline` section the trainer forces `same_on_batch=True` onto every child, which would make `p` one coin per batch and could take the spine out of both samples of a step. The EM is a batched, weighted form of `_em_gmm_1d` with no boolean indexing and no host read. Its moment accumulation is a broadcast product rather than `resp.mT @ design`: cuBLAS takes ~7.5 ms for that reduction shape, the product ~0.04 ms. `transform_params_one-sequence-to-segment-them-all-fov.json` (ae065824) is the paper's Ours config (c69b9872) plus the transform at its defaults and nothing else. Co-Authored-By: Claude Opus 5.5 --- .gitignore | 2 + README.md | 63 +++ smauglab/configs/all_augmentations.json | 13 + ..._one-sequence-to-segment-them-all-fov.json | 241 ++++++++++++ smauglab/registry.py | 5 + smauglab/transforms/__init__.py | 1 + smauglab/transforms/gpu/context.py | 363 ++++++++++++++++++ unit_tests/helpers.py | 7 +- unit_tests/test_container_fast_path.py | 2 +- unit_tests/test_fov_extension.py | 183 +++++++++ unit_tests/test_transforms_gpu.py | 1 + 11 files changed, 877 insertions(+), 4 deletions(-) create mode 100644 smauglab/configs/transform_params_one-sequence-to-segment-them-all-fov.json create mode 100644 smauglab/transforms/gpu/context.py create mode 100644 unit_tests/test_fov_extension.py diff --git a/.gitignore b/.gitignore index 794750d..6cccbce 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,8 @@ MANIFEST # configs beside it are tracked only because they predate the blanket rule above; # a new one has to say so. !smauglab/configs/transform_params_paper-synthseg.json +# The paper's Ours plus RandomFOVExtensionGPU (README "Field of view beyond the scan"). +!smauglab/configs/transform_params_one-sequence-to-segment-them-all-fov.json *.yaml *.nii *.nii.gz diff --git a/README.md b/README.md index 4755ed2..0935d31 100644 --- a/README.md +++ b/README.md @@ -208,6 +208,68 @@ transfer augmentations (`RandomScharrGPU`, `RandomUnsharpMaskGPU`, `RandomRandCo `p=0` and `RandomSynthSegGPU` at `p=1.0` in their place. The general-enhancement half is identical between the two. +### Field of view beyond the scan + +A model trained on a narrow field of view has never been shown what lies outside it. A +sagittal spine slab is about 8 cm wide, and nnU-Net centres it in a 128 mm patch and pads +the rest with 0. So every training window contains spine, and beyond about ±4 cm there is +only a flat slab labelled background. On a whole-body axial scan the model then finds +vertebrae in muscle, gluteals, fat and bowel. For the Ours setup trained on SpineGAN in-phase, +60% of the false-positive voxels on NAKO VIBE lie more than 5 cm from the spine. The contrast +augmentations do not touch this, because the problem is context, not intensity. + +`RandomFOVExtensionGPU` shows the network windows it can only get right by rejecting +context. For each selected sample it does six things: + +1. **Shift.** Image and segmentation move by the same whole-voxel shift, up to `shift` of + the patch per axis. At `shift: 1.0` the spine can leave the window entirely. +2. **Find the fill region.** That is everything that is not scan: the padding (every + channel exactly `pad_value`) plus whatever the shift vacated. +3. **Split the background with EM.** The patch's own background is split into a random 3 + to 10 intensity regions, the same way SynthSeg's label completion splits its + background. Only background voxels count: label 0, more than `margin` voxels from any + label, and off the padding's interpolated edge. +4. **Post-process the region map.** A 3×3×3 majority filter removes the EM speckle, and a + random diffeomorphic field (SynthSeg's) deforms the map. +5. **Composite into the fill.** `layers` copies of the map, each randomly rotated, scaled, + mirrored and translated, are composited into the fill region. +6. **Render.** Each fill voxel is drawn from its region's mean and standard deviation, + measured on the patch itself. The within-region variation is spatially correlated + (σ = 1 voxel), because iid draws read as grain next to the smooth scan. The result is + then blurred as SynthSeg blurs, and its label stays 0. + +| parameter | default | meaning | +| --- | --- | --- | +| `shift` | `0.5` | maximum shift as a fraction of the patch, one value or one per axis; `0` only fills the existing padding | +| `regions` | `[3, 10]` | range for the number of EM regions, drawn per sample | +| `margin` | `4` | voxels between a label and the background used to build the fill | +| `layers` | `6` | warped copies of the region map composited into the fill | +| `pad_value` | `0.0` | the intensity that marks padding; nnU-Net pads with 0 | +| `p` | `0.5` | probability per sample | + +Things to know before using it: + +* **Selection is always per sample.** The transform keeps this even when the pipeline sets + `same_on_batch`. Without a `pipeline` section a config runs with `same_on_batch: true`, + which would make `p` one coin per batch and could take the spine out of both samples of + a step. +* **The logged training loss rises.** A window the shift has emptied of spine has no + foreground, so its Dice term is about 0 and cross-entropy does the work. The loss curve + is therefore not comparable with a run without the transform. +* **Pipeline position matters.** The transform runs right after `RandomAffineGPU` and + before every intensity transform, so the synthesised background is augmented like real + tissue. A label-to-image generator later in the pipeline (SynthSeg, PALETTE synthesis) + regenerates the fill from label 0, which overwrites it. +* **Padding is found by exact value.** Two cases follow from that: + * Background that the scanner itself zeroed (the masked air of a Dixon reconstruction) + normalises to a constant other than 0. It counts as scan. It stays where it is, and + it can become one of the EM regions. + * With `use_mask_for_norm`, everything outside the body is 0, so it gets filled. + +`smauglab/configs/transform_params_one-sequence-to-segment-them-all-fov.json` (`ae065824`) +is the paper's Ours config (`c69b9872`) plus this transform at the defaults above. Nothing +else changes. + ## Available augmentations Which augmentations exist, and which of the two backends implements each one. An empty @@ -222,6 +284,7 @@ column -- see "Run Monai training with SmaugLab augmentations" above. Regenerate | affine | GEO | `RandomAffineGPU` | — | | crop | GEO | `RandomCropTransformGPU` | — | | spatial | GEO | — | `SpatialTransform` | +| fov_extension | GEO | `RandomFOVExtensionGPU` | — | | gaussian_noise | GE | `RandomGaussianNoiseGPU` | `GaussianNoiseTransform` | | gaussian_blur | GE | `RandomGaussianBlurGPU` | `GaussianBlurTransform` | | brightness | GE | `RandomBrightnessGPU` | `MultiplicativeBrightnessTransform` | diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json index 8be9b9b..cd64ba3 100644 --- a/smauglab/configs/all_augmentations.json +++ b/smauglab/configs/all_augmentations.json @@ -35,6 +35,19 @@ 0.1 ] }, + "RandomFOVExtensionGPU": { + "keepdim": true, + "layers": 6, + "margin": 4, + "p": 0.5, + "p_batch": 1.0, + "pad_value": 0.0, + "regions": [ + 3, + 10 + ], + "shift": 0.5 + }, "RandomSynthSegGPU": { "apply_bias_field": true, "apply_intensity_augmentation": true, diff --git a/smauglab/configs/transform_params_one-sequence-to-segment-them-all-fov.json b/smauglab/configs/transform_params_one-sequence-to-segment-them-all-fov.json new file mode 100644 index 0000000..0a5d21d --- /dev/null +++ b/smauglab/configs/transform_params_one-sequence-to-segment-them-all-fov.json @@ -0,0 +1,241 @@ +{ + "GPU": { + "RandomFlipTransformGPU": { + "flip_axis": [ + 0 + ], + "same_on_batch": false, + "keepdim": true, + "p": 0.5 + }, + "RandomAffineGPU": { + "degrees": 5, + "translate": [ + 0.1, + 0.1, + 0.1 + ], + "scale": [ + 0.7, + 1.4 + ], + "shears": [ + -5, + 5, + -5, + 5, + -5, + 5 + ], + "resample": "bilinear", + "p": 0.0 + }, + "RandomFOVExtensionGPU": { + "shift": 0.5, + "regions": [ + 3, + 10 + ], + "margin": 4, + "layers": 6, + "pad_value": 0.0, + "p": 0.5 + }, + "RandomInverseGPU": { + "retain_stats": true, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "mix_prob": 0.0, + "p": 0.25 + }, + "RandomHistogramEqualizationGPU": { + "retain_stats": true, + "in_seg": 0, + "out_seg": 0, + "mix_in_out": true, + "mix_prob": 0.6, + "p": 0.25 + }, + "RandomRedistributeSegGPU": { + "in_seg": 0.25, + "retain_stats": true, + "p": 0.4 + }, + "RandomScharrGPU": { + "absolute": true, + "retain_stats": true, + "mix_prob": 0.5, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "p": 0.25 + }, + "RandomUnsharpMaskGPU": { + "sigma": 1.0, + "unsharp_amount": 1.5, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "mix_prob": 0.5, + "p": 0.25 + }, + "RandomRandConvGPU": { + "kernel_sizes": [ + 3, + 5, + 7 + ], + "retain_stats": true, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "mix_prob": 0.5, + "p": 0.2 + }, + "RandomClampGPU": { + "max_clamp_amount": 0.2, + "retain_stats": true, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "p": 0.0 + }, + "RandomGaussianNoiseGPU": { + "mean": 0.0, + "std": 1.0, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "p": 0.1 + }, + "RandomGaussianBlurGPU": { + "sigma": 1.0, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "p": 0.2 + }, + "RandomBrightnessGPU": { + "brightness_range": [ + 0.75, + 1.25 + ], + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "p": 0.15 + }, + "RandomGammaGPU": { + "gamma_range": [ + 0.7, + 1.5 + ], + "retain_stats": true, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "p": 0.3 + }, + "RandomInvGammaGPU": { + "gamma_range": [ + 0.7, + 1.5 + ], + "retain_stats": true, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "p": 0.1 + }, + "RandomContrastGPU": { + "contrast_range": [ + 0.75, + 1.25 + ], + "retain_stats": false, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "p": 0.15 + }, + "RandomLog1pGPU": { + "retain_stats": true, + "in_seg": 0, + "out_seg": 0, + "mix_in_out": false, + "p": 0.1 + }, + "RandomSqrtGPU": { + "retain_stats": true, + "in_seg": 0, + "out_seg": 0, + "mix_in_out": false, + "p": 0.1 + }, + "RandomSinGPU": { + "retain_stats": true, + "in_seg": 0, + "out_seg": 0, + "mix_in_out": false, + "p": 0.1 + }, + "RandomExpGPU": { + "retain_stats": true, + "in_seg": 0, + "out_seg": 0, + "mix_in_out": false, + "p": 0.1 + }, + "RandomSigmoidGPU": { + "retain_stats": true, + "in_seg": 0, + "out_seg": 0, + "mix_in_out": false, + "p": 0.1 + }, + "RandomLowResTransformGPU": { + "scale": [ + 0.5, + 1.0 + ], + "same_on_batch": false, + "p": 0.25 + }, + "RandomAcqTransformGPU": { + "scale": [ + 0.6, + 1.0 + ], + "same_on_batch": false, + "p": 0.0 + }, + "RandomBiasFieldGPU": { + "retain_stats": true, + "coefficients": 0.2, + "in_seg": 0.0, + "out_seg": 0.0, + "mix_in_out": true, + "p": 0.15 + }, + "ZscoreNormalizationGPU": { + "p": 0.0 + } + }, + "CPU": { + "SpatialTransform": { + "patch_center_dist_from_border": 80, + "random_crop": false, + "p_elastic_deform": 0.0, + "p_rotation": 0.2, + "p_scaling": 0.2, + "scaling": [ + 0.7, + 1.4 + ], + "p_synchronize_scaling_across_axes": 1, + "bg_style_seg_sampling": false, + "mode_seg": "nearest" + } + } +} diff --git a/smauglab/registry.py b/smauglab/registry.py index 36805bc..1a85498 100644 --- a/smauglab/registry.py +++ b/smauglab/registry.py @@ -96,6 +96,7 @@ class AugId(str, Enum): AFFINE = "affine" CROP = "crop" SPATIAL = "spatial" # nnU-Net's own SpatialTransform + FOV_EXTENSION = "fov_extension" # -- general enhancement GAUSSIAN_NOISE = "gaussian_noise" GAUSSIAN_BLUR = "gaussian_blur" @@ -200,6 +201,10 @@ def __init__(self, source: str, problems: list[str]) -> None: Backend.GPU: ( "RandomFlipTransformGPU", # FlipTransform "RandomAffineGPU", # AffineTransform + # No ladder key. After the last transform that can open padding (the affine's zero + # fill) and before everything that rewrites intensities, so the synthesised + # background is augmented exactly like the anatomy around it. + "RandomFOVExtensionGPU", "RandomSynthSegGPU", # SynthSeg "RandomPaletteGPU", # RandomPALETTETransform "PaletteSynthesisGPU", # same slot: it is the same augmentation, composed diff --git a/smauglab/transforms/__init__.py b/smauglab/transforms/__init__.py index f3a9ece..cbfd0e7 100644 --- a/smauglab/transforms/__init__.py +++ b/smauglab/transforms/__init__.py @@ -16,6 +16,7 @@ class carries an `@register(...)` decorator, so the registry is complete once th """ from smauglab.transforms.cpu import artifact, contrast, external, fromSeg, spatial # noqa: F401 +from smauglab.transforms.gpu import context as gpu_context # noqa: F401 from smauglab.transforms.gpu import contrast as gpu_contrast # noqa: F401 from smauglab.transforms.gpu import domain_transfer # noqa: F401 from smauglab.transforms.gpu import fromSeg as gpu_fromSeg diff --git a/smauglab/transforms/gpu/context.py b/smauglab/transforms/gpu/context.py new file mode 100644 index 0000000..ba41761 --- /dev/null +++ b/smauglab/transforms/gpu/context.py @@ -0,0 +1,363 @@ +"""What lies around the scan: shift the scan inside the patch, and synthesise what is outside it. + +A sagittal spine slab is ~8 cm wide. nnU-Net centres it in a 128 mm patch and pads the rest +with 0 (seg -1 -> 0), so the network sees spine in every training window and nothing but a +flat zero slab beyond +-4 cm. On a whole-body axial scan it then finds vertebrae in muscle, +gluteals, fat and bowel: on NAKO VIBE most of its false-positive volume lies more than 5 cm +from the spine. `RandomFOVExtensionGPU` shows it windows where the spine sits at the edge or +is absent, surrounded by tissue that is labelled background. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import Any + +import torch +import torch.nn.functional as F +from kornia.augmentation._3d.base import RigidAffineAugmentationBase3D +from kornia.augmentation.random_generator.base import RandomGeneratorBase, UniformDistribution +from kornia.augmentation.utils import _adapted_rsampling +from torch import Tensor + +from smauglab.registry import AugId, AugType, Backend, register +from smauglab.transforms.gpu.base import segmentation_from +from smauglab.transforms.gpu.contrast import _foreground +from smauglab.transforms.synthseg.functional import _quadratic_log_likelihood_terms, gaussian_blur_3d, random_svf_field, warp_volume + +# How the region map is carried into the empty part of the patch. Constants rather than +# parameters: they shape a texture that is always labelled background, and the screen that +# introduced the transform varied `shift` and `p`, not these. +_SCALE_RANGE = (0.75, 1.33) # log-uniform, ~nnU-Net's 0.7-1.4 +_EM_ITERS = 20 # SynthSeg's em_n_iters default +_EM_SAMPLES = 100_000 # SynthSeg's em_max_fit_voxels default +_SEAM = 1 # voxels of interpolated padding edge never used as a source +_CLOSE_ITERS = 8 # max-pool steps that close what the warped layers missed +_TEXTURE_SIGMA = 1.0 # correlation length of the within-region texture, in voxels +_RENDER_SIGMA = 0.5 # SynthSeg's blur at the atlas resolution, in voxels +_RENDER_BLUR_RANGE = 1.15 # SynthSeg's DynamicGaussianBlur jitter + + +def _dilate(mask: Tensor, radius: int) -> Tensor: + """Binary dilation of a `[N, 1, D, H, W]` bool mask by a cube of half-width `radius`. + + Three 1-D max-pools rather than one dense `(2r+1)^3` pool: the same result for a cube. + """ + if radius <= 0: + return mask + k = 2 * radius + 1 + x = mask.float() + for kernel in ((k, 1, 1), (1, k, 1), (1, 1, k)): + x = F.max_pool3d(x, kernel, stride=1, padding=tuple(s // 2 for s in kernel)) + return x > 0 + + +def _shift(volume: Tensor, offsets: list[list[int]], fill: float | bool) -> Tensor: + """Translate each row of `[N, C, D, H, W]` by whole voxels; vacated voxels get `fill`. + + `offsets[b]` is the `(d, h, w)` displacement of row `b`. Integer and exact, so the + image and the mask move identically and no label is interpolated. + """ + out = torch.full_like(volume, fill) + spatial = volume.shape[2:] + for b, offset in enumerate(offsets): + src, dst = [], [] + for size, o in zip(spatial, offset): + if abs(o) >= size: + break + src.append(slice(max(0, -o), size - max(0, o))) + dst.append(slice(max(0, o), size - max(0, -o))) + else: + out[b][(slice(None), *dst)] = volume[b][(slice(None), *src)] + return out + + +def _em_regions(x: Tensor, weight: Tensor, n_regions: Tensor, k_max: int, eps: float = 1e-6) -> Tensor: + """Fit a 1-D Gaussian mixture per row by EM and return each voxel's most likely component. + + The batched, weighted form of `synthseg.functional._em_gmm_1d`: every row is fitted at + once, on `_EM_SAMPLES` voxels drawn uniformly and weighted by `weight`, so selecting the + background costs no boolean indexing and no host synchronisation. Row `b` uses its first + `n_regions[b]` of `k_max` components; the rest are masked out of the softmax. + + Args: + x: `[N, V]` intensities, all voxels of the patch. + weight: `[N, V]` 1.0 where a voxel may be fitted, else 0.0. + n_regions: `[N]` long, components per row, at least 1. + k_max: number of component slots. + + Returns: + `[N, V]` long, the component index of every voxel. + """ + n, v = x.shape + idx = torch.randint(0, v, (n, _EM_SAMPLES), device=x.device) + xs, ws = x.gather(1, idx), weight.gather(1, idx) + + comp = torch.arange(k_max, device=x.device) + active = comp.view(1, -1) < n_regions.view(-1, 1) # [N, K] + total = ws.sum(1, keepdim=True).clamp_min(eps) + shift = (ws * xs).sum(1, keepdim=True) / total + centred = xs - shift + + big = torch.finfo(x.dtype).max + lo = torch.where(ws > 0, centred, big).amin(1, keepdim=True) + hi = torch.where(ws > 0, centred, -big).amax(1, keepdim=True) + lo, hi = torch.where(total > eps, lo, 0.0), torch.where(total > eps, hi, 0.0) + means = lo + (hi - lo) / (n_regions.view(-1, 1) - 1).clamp_min(1) * comp.view(1, -1) + var = ((ws * centred * centred).sum(1, keepdim=True) / total).clamp_min(eps).expand(n, k_max) + mix = active.float() / n_regions.view(-1, 1) + + design = torch.stack([centred * centred, centred, torch.ones_like(centred)], dim=2) # [N, S, 3] + for _ in range(_EM_ITERS): + terms = _quadratic_log_likelihood_terms(means, var, mix, eps).permute(1, 0, 2) # [N, 3, K] + resp = torch.softmax((design @ terms).masked_fill(~active.unsqueeze(1), -math.inf), dim=2) * ws.unsqueeze(2) + # Sums of r*x^2, r*x, r per component. A broadcast product rather than + # `resp.mT @ design`: cuBLAS takes ~7.5 ms for that reduction shape, this ~0.04 ms. + moments = (resp.unsqueeze(3) * design.unsqueeze(2)).sum(1) # [N, K, 3] + nk = moments[..., 2].clamp_min(eps) + mix = torch.where(active, nk / total, 0.0) + means = moments[..., 1] / nk + var = (moments[..., 0] / nk - means * means).clamp_min(eps) + + terms = _quadratic_log_likelihood_terms(means, var, mix, eps).permute(1, 0, 2) + c = (x - shift).unsqueeze(2) + logits = torch.cat([c * c, c, torch.ones_like(c)], dim=2) @ terms # [N, V, K] + return logits.masked_fill(~active.unsqueeze(1), -math.inf).argmax(dim=2) + + +def _random_similarity(n: int, device: torch.device) -> Tensor: + """`[n, 3, 3]` random rotation (uniform axis, angle in [-pi, pi]) x log-uniform scale x mirror.""" + axis = F.normalize(torch.randn(n, 3, device=device), dim=1) + angle = (torch.rand(n, device=device) * 2.0 - 1.0) * math.pi + kx = torch.zeros(n, 3, 3, device=device) + kx[:, 0, 1], kx[:, 0, 2], kx[:, 1, 2] = -axis[:, 2], axis[:, 1], -axis[:, 0] + kx = kx - kx.transpose(1, 2) + eye = torch.eye(3, device=device).expand(n, 3, 3) + sin, cos = angle.sin().view(n, 1, 1), angle.cos().view(n, 1, 1) + rotation = eye + sin * kx + (1.0 - cos) * (kx @ kx) # Rodrigues + lo, hi = math.log(_SCALE_RANGE[0]), math.log(_SCALE_RANGE[1]) + scale = torch.exp(lo + torch.rand(n, 1, device=device) * (hi - lo)) + mirror = torch.where(torch.rand(n, 3, device=device) < 0.5, -1.0, 1.0) + return rotation @ torch.diag_embed(scale * mirror) + + +def _extend_regions(regions: Tensor, todo: Tensor, layers: int) -> Tensor: + """Write region labels into `todo` from warped copies of the labels already present. + + `regions` is `[N, 1, D, H, W]` long, `-1` where unlabelled. Each layer warps the + current map -- including what earlier layers filled -- by a random rotation, scale, + mirror and translation with reflection padding, and keeps the voxels of `todo` it + lands a label on. Whatever is still unlabelled after the layers takes a neighbour's + label by max-pool propagation, then region 0. + """ + n = regions.shape[0] + shape = regions.shape[2:] + half = torch.tensor(shape, device=regions.device, dtype=torch.float32).view(1, 3, 1) / 2.0 + cur = regions.clone() + for _ in range(layers): + affine = torch.zeros(n, 4, 4, device=regions.device) + affine[:, :3, :3] = _random_similarity(n, regions.device) + affine[:, :3, 3:] = (torch.rand(n, 3, 1, device=regions.device) * 2.0 - 1.0) * half + affine[:, 3, 3] = 1.0 + warped = warp_volume((cur + 1).float(), affine=affine, interp="nearest", padding_mode="reflection").round().long() - 1 + take = todo & (warped >= 0) + cur = torch.where(take, warped, cur) + todo = todo & ~take + for _ in range(_CLOSE_ITERS): + grown = F.max_pool3d((cur + 1).float(), 3, stride=1, padding=1).long() - 1 + take = todo & (grown >= 0) + cur = torch.where(take, grown, cur) + todo = todo & ~take + return torch.where(todo, 0, cur) + + +def _render(image: Tensor, atlas: Tensor, layout: Tensor, k_max: int, eps: float = 1e-6) -> Tensor: + """Draw every voxel of `layout` from its region's Gaussian, per channel, then blur as SynthSeg does. + + Means and variances are measured where `atlas` labels the patch itself (`-1` elsewhere), + so the fill has the patch's own contrast; every intensity augmentation after this + transform then acts on fill and anatomy alike. The within-region variation is spatially + correlated noise with the region's standard deviation: SynthSeg's iid draws read as grain + next to the smooth scan, a cue that alone would separate fill from anatomy. + """ + n, c = image.shape[:2] + # A comparison rather than `F.one_hot`, which reads the label range back to the host on CUDA. + member = (atlas.flatten(1).unsqueeze(2) == torch.arange(k_max, device=atlas.device)).float() # [N, V, K], -1 matches none + flat = image.flatten(2) # [N, C, V] + count = member.sum(1).clamp_min(eps).unsqueeze(1) # [N, 1, K] + mean = (flat @ member) / count + std = (((flat * flat) @ member) / count - mean * mean).clamp_min(0.0).sqrt() + index = layout.flatten(1).clamp_min(0).unsqueeze(1).expand(n, c, -1) # -1 lies outside the fill + texture = gaussian_blur_3d(torch.randn_like(image), torch.full((3,), _TEXTURE_SIGMA, device=image.device)).flatten(2) + texture = texture / texture.std(dim=2, keepdim=True).clamp_min(eps) + drawn = mean.gather(2, index) + std.gather(2, index) * texture + sigma = torch.full((3,), _RENDER_SIGMA, device=image.device) + return gaussian_blur_3d(drawn.view_as(image), sigma, blur_range=_RENDER_BLUR_RANGE) + + +class FOVExtensionGenerator3D(RandomGeneratorBase): + """Per row: the shift as a fraction of the patch along `(d, h, w)`, and the EM region count.""" + + def __init__(self, shift: tuple[float, float, float], regions: tuple[int, int]) -> None: + super().__init__() + self.shift = shift + self.regions = regions + + def make_samplers(self, device: torch.device, dtype: torch.dtype) -> None: + self.shift_samplers = [ + UniformDistribution( + torch.tensor(-s, device=device, dtype=dtype), torch.tensor(s, device=device, dtype=dtype), validate_args=False + ) + for s in self.shift + ] + lo, hi = self.regions + self.regions_sampler = UniformDistribution( + torch.tensor(float(lo), device=device, dtype=dtype), + torch.tensor(float(hi + 1), device=device, dtype=dtype), + validate_args=False, + ) + + def forward(self, batch_shape: tuple[int, ...], same_on_batch: bool = False) -> dict[str, Tensor]: + n = batch_shape[0] + shift = torch.stack([_adapted_rsampling((n,), s, same_on_batch) for s in self.shift_samplers], dim=1) + regions = _adapted_rsampling((n,), self.regions_sampler, same_on_batch).floor().clamp(max=float(self.regions[1])) + return {"shift": shift, "regions": regions} + + +@register( + aug_id=AugId.FOV_EXTENSION, + backend=Backend.GPU, + group=AugType.GEO, + smoke_kwargs={"shift": 0.5, "regions": (2, 4), "layers": 2}, +) +class RandomFOVExtensionGPU(RigidAffineAugmentationBase3D): + """Shift the scan inside the patch and fill everything outside it with EM-synthesised background. + + Per selected sample: + + 1. Image and segmentation move by the same whole-voxel shift, up to `shift` of the patch + per axis. At `shift=1.0` the spine can leave the window entirely. + 2. The *fill region* is every voxel that is not scan: the padding (every channel exactly + `pad_value`, which is what nnU-Net pads with) plus whatever the shift vacated. + 3. The patch's own background -- label 0, more than `margin` voxels from any label, not + on the padding's interpolated edge -- is split into a random 3 to 10 (`regions`) + intensity regions by EM, as SynthSeg's label completion splits its background. + 4. The region map is post-processed: a 3x3x3 majority filter removes the EM speckle, and + a random diffeomorphic field (SynthSeg's) deforms it. + 5. `layers` randomly rotated, scaled, mirrored and translated copies of it are composited + into the fill region; any gap left takes a neighbour's region. + 6. Each fill voxel is drawn from its region's Gaussian, measured on the patch itself, + with spatially correlated texture, and blurred at SynthSeg's atlas-resolution sigma. + Its label is 0. + + Everything outside the fill region is the shifted input, bit for bit. A sample with no + usable background still shifts and keeps the fill at `pad_value`, so image and mask agree. + Without a segmentation (an image-only pipeline) the background is every non-padding voxel. + + **Row selection is always per sample**, whatever the pipeline's `same_on_batch` says. + Under the paper config the trainer forces `same_on_batch=True` onto every child, which + would make `p` one coin per batch and could empty both samples of a step of spine. + + Place it before every intensity transform, which `PIPELINE_ORDER` does, so the fill is + augmented like real tissue. A label-to-image generator (SynthSeg, PALETTE synthesis) + later in the pipeline regenerates the fill from label 0 and overwrites it. + + Args: + shift: maximum shift as a fraction of the patch, one value or one per `(d, h, w)` + axis, drawn uniformly in `[-shift, shift]`. 0 fills the existing padding only. + regions: `(min, max)` number of EM regions, drawn per sample. + margin: voxels between any label and the nearest background voxel used as a source. + layers: warped copies of the region map composited into the fill region. + pad_value: the intensity that marks padding. + """ + + def __init__( + self, + shift: float | Sequence[float] = 0.5, + regions: Sequence[int] = (3, 10), + margin: int = 4, + layers: int = 6, + pad_value: float = 0.0, + p: float = 0.5, + p_batch: float = 1.0, + keepdim: bool = True, + ) -> None: + super().__init__(p=p, p_batch=p_batch, same_on_batch=False, keepdim=keepdim) + shifts = (float(shift),) * 3 if isinstance(shift, (int, float)) else tuple(float(s) for s in shift) + if len(shifts) != 3 or not all(0.0 <= s <= 1.0 for s in shifts): + raise ValueError(f"shift must be one value or three, each in [0, 1]. Got {shift!r}.") + if len(regions) != 2 or not 1 <= int(regions[0]) <= int(regions[1]): + raise ValueError(f"regions must be (min, max) with 1 <= min <= max. Got {regions!r}.") + if int(margin) < 0 or int(layers) < 1: + raise ValueError(f"margin must be >= 0 and layers >= 1. Got margin={margin!r}, layers={layers!r}.") + self.regions = (int(regions[0]), int(regions[1])) + self.margin = int(margin) + self.layers = int(layers) + self.pad_value = float(pad_value) + self._param_generator = FOVExtensionGenerator3D(shifts, self.regions) + + @property # type: ignore[override] + def same_on_batch(self) -> bool: + """Always False: see the class docstring. The setter ignores what a container writes.""" + return False + + @same_on_batch.setter + def same_on_batch(self, value: bool) -> None: + pass + + def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor: + return self.identity_matrix(input) + + @staticmethod + def _offsets(params: dict[str, Tensor], spatial: Sequence[int]) -> list[list[int]]: + size = torch.tensor(list(spatial), dtype=torch.float32) + # One transfer for the whole batch, as in RandomCropTransformGPU. + return (params["shift"].detach().float().cpu() * size).round().long().tolist() + + @torch.no_grad() + def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: + offsets = self._offsets(params, input.shape[2:]) + # fp32 throughout: under the trainer's autocast `warp_volume`'s bmm and the EM + # matmuls would run in half precision. + with torch.autocast(device_type=input.device.type, enabled=False): + image = input.float() + n = image.shape[0] + d, h, w = image.shape[2:] + padding = (image == self.pad_value).all(dim=1, keepdim=True) + seg = segmentation_from(params) + labelled = torch.zeros_like(padding) if seg is None else _foreground(seg, dim=1).unsqueeze(1) + + image = _shift(image, offsets, self.pad_value) + fill = _shift(padding, offsets, True) + labelled = _shift(labelled, offsets, False) + + sources = ~(_dilate(labelled, self.margin) | _dilate(fill, _SEAM)) + k_max = self.regions[1] + n_regions = params["regions"].to(image.device).long() + em = _em_regions(image[:, 0].flatten(1), sources.flatten(1).float(), n_regions, k_max).view_as(padding) + + # Majority filter over the 3x3x3 neighbourhood, among source voxels only. + votes = (em == torch.arange(k_max, device=em.device).view(1, -1, 1, 1, 1)).float() * sources + atlas = F.avg_pool3d(votes, 3, stride=1, padding=1).argmax(dim=1, keepdim=True) + atlas = torch.where(sources, atlas, -1) + + displacement = random_svf_field(n, (d, h, w), image.device) + deformed = warp_volume((atlas + 1).float(), displacement=displacement, interp="nearest").round().long() - 1 + layout = _extend_regions(torch.where(fill, -1, deformed), fill, self.layers) + + rendered = _render(image, atlas, layout, k_max) + has_sources = sources.flatten(1).any(dim=1).view(n, 1, 1, 1, 1) + rendered = torch.where(has_sources, rendered, self.pad_value) + out = torch.where(fill, rendered, image) + return out.to(input.dtype) + + def apply_non_transform_mask( + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None + ) -> Tensor: + return input + + def apply_transform_mask( + self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None + ) -> Tensor: + return _shift(input, self._offsets(params, input.shape[2:]), 0) diff --git a/unit_tests/helpers.py b/unit_tests/helpers.py index 8c12622..84cc94a 100644 --- a/unit_tests/helpers.py +++ b/unit_tests/helpers.py @@ -102,9 +102,10 @@ def all_config_paths() -> list[Path]: def gpu_config_paths() -> list[Path]: """Configs that drive the GPU augmentation pipeline.""" paths = [p for p in all_config_paths() if p.name.startswith("transform_params_gpu")] - extra = configs_dir() / "transform_params_one-sequence-to-segment-them-all.json" - if extra.is_file(): - paths.append(extra) + for name in ("transform_params_one-sequence-to-segment-them-all.json", "transform_params_one-sequence-to-segment-them-all-fov.json"): + extra = configs_dir() / name + if extra.is_file(): + paths.append(extra) return sorted(paths) diff --git a/unit_tests/test_container_fast_path.py b/unit_tests/test_container_fast_path.py index 2523487..61c13b7 100644 --- a/unit_tests/test_container_fast_path.py +++ b/unit_tests/test_container_fast_path.py @@ -194,7 +194,7 @@ def test_it_matches_kornia_for_every_image_only_transform(self): def test_a_geometric_transform_is_not_short_circuited(self): """The ones that do move a mask have to keep going through kornia.""" - for name in ("RandomFlipTransformGPU", "RandomAffineGPU", "RandomCropTransformGPU"): + for name in ("RandomFlipTransformGPU", "RandomAffineGPU", "RandomCropTransformGPU", "RandomFOVExtensionGPU"): with self.subTest(transform=name): module = registry.get(name, Backend.GPU).cls() self.assertFalse(_leaves_masks_untouched(module)) diff --git a/unit_tests/test_fov_extension.py b/unit_tests/test_fov_extension.py new file mode 100644 index 0000000..b529727 --- /dev/null +++ b/unit_tests/test_fov_extension.py @@ -0,0 +1,183 @@ +"""`RandomFOVExtensionGPU`: shift the scan inside the patch, fill what lies outside it. + +The fixture is a small sagittal-slab stand-in: tissue in [1, 2], a slab of exact zeros at +both ends of the first axis (nnU-Net's padding), and a labelled block at intensity 1000. +The sentinel makes two properties checkable voxel by voxel: the image and the mask moved +together (1000 sits exactly under the label), and no foreground intensity reached the +synthesised background (nothing else comes near 1000). +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from smauglab.transforms.gpu.base import AugmentationSequentialCustom +from smauglab.transforms.gpu.context import RandomFOVExtensionGPU, _dilate, _em_regions, _shift +from unit_tests.helpers import SmaugLabTestCase + +SHAPE = (3, 1, 28, 24, 20) # non-cubic on purpose +SENTINEL = 1000.0 + + +def slab(shape=SHAPE): + image = torch.rand(shape) + 1.0 + image[:, :, :6] = 0.0 + image[:, :, -6:] = 0.0 + seg = torch.zeros(shape) + seg[:, :, 10:18, 6:16, 7:13] = 1.0 + image[seg.bool()] = SENTINEL + return image, seg + + +def run(transform, image, seg, same_on_batch=True): + pipeline = AugmentationSequentialCustom(transform, data_keys=["input", "mask"], same_on_batch=same_on_batch) + return pipeline(image.clone(), seg.clone()) + + +class TestShiftAndFill(SmaugLabTestCase): + def test_image_and_mask_move_together(self): + image, seg = slab() + for seed in range(6): + with self.subTest(seed=seed): + torch.manual_seed(seed) + out, out_seg = run(RandomFOVExtensionGPU(shift=1.0, p=1.0), image, seg) + self.assertTrue(bool(torch.equal(out > SENTINEL / 2, out_seg > 0))) + + def test_the_shift_moves_something(self): + """The guard against passing the alignment test by never shifting.""" + image, seg = slab() + torch.manual_seed(0) + _, out_seg = run(RandomFOVExtensionGPU(shift=1.0, p=1.0), image, seg) + self.assertFalse(bool(torch.equal(out_seg, seg))) + + def test_outside_the_fill_is_the_shifted_input(self): + image, seg = slab() + transform = RandomFOVExtensionGPU(shift=0.5, p=1.0) + torch.manual_seed(3) + params = transform.generate_parameters(image.shape) + params["seg"] = seg + out = transform.apply_transform(image, params, transform.flags) + offsets = transform._offsets(params, image.shape[2:]) + shifted = _shift(image, offsets, 0.0) + scan = _shift(torch.ones_like(image, dtype=torch.bool), offsets, False) & (shifted != 0.0) + self.assertTrue(bool(torch.equal(out[scan], shifted[scan]))) + + def test_no_shift_touches_only_the_padding(self): + image, seg = slab() + out, out_seg = run(RandomFOVExtensionGPU(shift=0.0, p=1.0), image, seg) + scan = image != 0.0 + self.assertTrue(bool(torch.equal(out_seg, seg))) + self.assertTrue(bool(torch.equal(out[scan], image[scan]))) + self.assertEqual(int((out[~scan] == 0.0).sum()), 0, "padding left unfilled") + + def test_no_foreground_intensity_reaches_the_fill(self): + image, seg = slab() + for margin in (0, 2): + with self.subTest(margin=margin): + torch.manual_seed(1) + out, out_seg = run(RandomFOVExtensionGPU(shift=1.0, margin=margin, p=1.0), image, seg) + background = out[out_seg == 0] + self.assertLess(float(background.max()), 10.0) + self.assertTrue(bool(torch.isfinite(out).all())) + + def test_unselected_rows_are_untouched(self): + image, seg = slab() + torch.manual_seed(0) + out, out_seg = run(RandomFOVExtensionGPU(p=0.0), image, seg) + self.assertTrue(bool(torch.equal(out, image))) + self.assertTrue(bool(torch.equal(out_seg, seg))) + + def test_no_background_still_shifts_consistently(self): + """Nothing to fit: the fill stays at pad_value, and the mask still follows the image.""" + image, seg = slab() + seg[:, :, 6:-6] = 1.0 + image[seg.bool()] = SENTINEL + torch.manual_seed(2) + out, out_seg = run(RandomFOVExtensionGPU(shift=0.5, p=1.0), image, seg) + self.assertTrue(bool(torch.equal(out > SENTINEL / 2, out_seg > 0))) + self.assertTrue(bool(((out == 0.0) | (out == SENTINEL)).all())) + + def test_without_a_segmentation_every_scan_voxel_is_a_source(self): + image, _ = slab() + transform = RandomFOVExtensionGPU(shift=0.0, p=1.0) + params = transform.generate_parameters(image.shape) + out = transform.apply_transform(image, params, transform.flags) + self.assertEqual(int((out == 0.0).sum()), 0) + self.assertTrue(bool(torch.isfinite(out).all())) + + def test_padding_means_every_channel(self): + image, seg = slab((2, 2, 28, 24, 20)) + image[:, 1, :6] = 5.0 # channel 1 is not padding there, so neither is the voxel + out, _ = run(RandomFOVExtensionGPU(shift=0.0, p=1.0), image, seg) + self.assertTrue(bool(torch.equal(out[:, :, :6], image[:, :, :6]))) + self.assertEqual(int((out[:, :, -6:] == 0.0).all(dim=1).sum()), 0) + + +class TestSelectionAndDeterminism(SmaugLabTestCase): + def test_rows_are_selected_per_sample_under_same_on_batch(self): + """The paper config hands every child same_on_batch=True; this transform ignores it.""" + image, seg = slab((8, 1, 28, 24, 20)) + counts = [] + for seed in range(8): + torch.manual_seed(seed) + transform = RandomFOVExtensionGPU(shift=0.0, p=0.5) + out, _ = run(transform, image, seg, same_on_batch=True) + self.assertFalse(transform.same_on_batch) + counts.append(int((out != image).flatten(1).any(dim=1).sum())) + self.assertTrue(any(0 < c < 8 for c in counts), counts) + + def test_seed_determinism(self): + image, seg = slab() + results = [] + for _ in range(2): + torch.manual_seed(11) + results.append(run(RandomFOVExtensionGPU(shift=0.5, p=1.0), image, seg)) + self.assertTrue(bool(torch.equal(results[0][0], results[1][0]))) + self.assertTrue(bool(torch.equal(results[0][1], results[1][1]))) + + def test_autocast_matches_fp32(self): + image, seg = slab() + torch.manual_seed(5) + ref, ref_seg = run(RandomFOVExtensionGPU(shift=0.5, p=1.0), image, seg) + torch.manual_seed(5) + with torch.autocast("cpu", dtype=torch.bfloat16): + out, out_seg = run(RandomFOVExtensionGPU(shift=0.5, p=1.0), image, seg) + self.assertTrue(bool(torch.equal(out_seg, ref_seg))) + self.assertTrue(bool(torch.allclose(out.float(), ref, atol=1e-5))) + + def test_rejects_bad_arguments(self): + for kwargs in ({"shift": 1.5}, {"shift": (0.5, 0.5)}, {"regions": (0, 3)}, {"regions": (5, 3)}, {"layers": 0}, {"margin": -1}): + with self.subTest(**kwargs), self.assertRaises(ValueError): + RandomFOVExtensionGPU(**kwargs) + + +class TestHelpers(SmaugLabTestCase): + def test_em_separates_two_populations(self): + x = torch.cat([torch.randn(1, 4000) * 0.1, torch.randn(1, 4000) * 0.1 + 5.0], dim=1) + assign = _em_regions(x, torch.ones_like(x), torch.tensor([2]), k_max=4) + first, second = assign[0, :4000], assign[0, 4000:] + self.assertGreater(float((first == first.mode().values).float().mean()), 0.99) + self.assertGreater(float((second == second.mode().values).float().mean()), 0.99) + self.assertNotEqual(int(first.mode().values), int(second.mode().values)) + self.assertLess(int(assign.max()), 2, "a masked-out component was assigned") + + def test_em_fits_only_weighted_voxels(self): + """With weight on one population only, both components land inside it.""" + x = torch.cat([torch.randn(1, 4000) * 0.1, torch.randn(1, 4000) * 0.1 + 5.0], dim=1) + weight = torch.cat([torch.ones(1, 4000), torch.zeros(1, 4000)], dim=1) + assign = _em_regions(x, weight, torch.tensor([2]), k_max=2) + self.assertEqual(int(assign[0, 4000:].unique().numel()), 1) + + def test_dilate_matches_a_dense_pool(self): + mask = torch.rand(2, 1, 12, 10, 9) > 0.95 + for radius in (0, 1, 3): + with self.subTest(radius=radius): + k = 2 * radius + 1 + dense = F.max_pool3d(mask.float(), k, stride=1, padding=radius) > 0 + self.assertTrue(bool(torch.equal(_dilate(mask, radius), dense))) + + def test_shift_beyond_the_patch_empties_the_row(self): + volume = torch.ones(1, 1, 4, 4, 4) + self.assertTrue(bool((_shift(volume, [[0, 4, 0]], 0.0) == 0).all())) + self.assertTrue(bool(torch.equal(_shift(volume, [[0, 0, 0]], 0.0), volume))) diff --git a/unit_tests/test_transforms_gpu.py b/unit_tests/test_transforms_gpu.py index d370f68..427b3c9 100644 --- a/unit_tests/test_transforms_gpu.py +++ b/unit_tests/test_transforms_gpu.py @@ -22,6 +22,7 @@ from unit_tests.helpers import SmaugLabTestCase, first_output TRANSFORM_MODULES = [ + "smauglab.transforms.gpu.context", "smauglab.transforms.gpu.contrast", "smauglab.transforms.gpu.spatial", "smauglab.transforms.gpu.fromSeg",