From 1a3aaa4a2688724793a3dbde0d9fcddc9fed6192 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 15 Oct 2026 12:02:26 +0000 Subject: [PATCH 1/3] feat: add the region-selection module behind gpu/regions.py Nothing calls it yet. The parameters it serves -- `in_seg`, `out_seg` and `mix_in_out` -- expressed about a third of what they looked like they expressed, and this is the one implementation the whole group will share, in the shape `gpu/mixing.py` established for `mix_prob` / `mix_alpha`. What it adds over the three helpers in gpu/contrast.py it will replace: * `in_seg` and `out_seg` both firing now means a *mixture* of foreground and background rather than falling through to "the whole patch", so `in_seg=1, out_seg=1` stops being indistinguishable from `in_seg=0, out_seg=0`. * `seg_unit` / `seg_unit_prob` replace `mix_in_out`, which keyed its keep/drop draw on the mask's *channel* axis. On the single-channel label map nnU-Net's trainer passes that is one draw per sample, deleting the whole mask half the time -- so the thirteen shipped configs setting it have been running a coin flip rather than a label mix. The replacement switches whole label values or whole connected components. * `bg_split` / `bg_em_clusters` let the background be split into intensity-coherent subregions by EM (SynthSeg sec. 5.4) so `out_seg` can apply to part of it instead of all or none. `em_subdivide_labels` does the work; this wraps it. * `seg_feather` turns the composite weight from a binary step into a Gaussian ramp. The old composite put a one-voxel discontinuity exactly on the label boundary -- an edge perfectly correlated with the segmentation target. `connected_components_3d` is the only new algorithm: label-equivalence propagation with a pointer jump, 6-connected to match `scipy.ndimage.label`'s default. Verified against scipy on random blobs (exact partition agreement over six trials), on corner-touching voxels (two components, which is what pins 6- over 26-connectivity), and on multi-label volumes where the component count must equal the sum of per-label scipy counts. Label boundaries are respected in a single pass: the propagation is gated on same-label neighbours, so it does not need a pass per label value. int64 throughout, because float32 is exact only to 2^24 and a 256^3 patch has exactly that many voxels. Three properties the code it replaces broke, each now tested: * a knob that is off spends no random numbers -- `_choose_region_mode` drew its two uniforms whenever a mask was attached, even at zero; * the number of uniforms a draw spends is a function of the configuration and never of the data, so a four-label patch and a five-label one leave the stream in the same place; * the EM draw is forked, so a cache hit and a cache miss leave the global stream in the same place. The decomposition is cached as an attribute on the mask tensor rather than in a dict keyed on it: `Tensor.__hash__` is identity-based but `__eq__` is elementwise, so a hash-bucket collision makes `dict` evaluate `bool(a == b)` and raise. Measured: a WeakKeyDictionary with five hundred same-shaped tensor keys raises reliably. `normalise_blend_range` gains a `bounds=` argument so `seg_feather`, a sigma in voxels with no upper limit, can share it with the three blend weights that need the unit interval. The default keeps the `[0, 1]` wording its existing callers' tests assert on. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/gpu/mixing.py | 29 +- smauglab/transforms/gpu/regions.py | 825 +++++++++++++++++++++++++++++ 2 files changed, 850 insertions(+), 4 deletions(-) create mode 100644 smauglab/transforms/gpu/regions.py diff --git a/smauglab/transforms/gpu/mixing.py b/smauglab/transforms/gpu/mixing.py index c5fa185..e79d6df 100644 --- a/smauglab/transforms/gpu/mixing.py +++ b/smauglab/transforms/gpu/mixing.py @@ -83,8 +83,23 @@ class Mixing(Protocol): _MIX_ALPHA_HINT = "It is the weight on the original input: 0 is the transform's output, 1 is the input untouched." -def normalise_blend_range(spec: MixAlphaSpec, key: str = "mix_alpha", owner: str = "", hint: str = _MIX_ALPHA_HINT) -> tuple[float, float]: - """A scalar or ``[lo, hi]`` as a validated pair inside ``[0, 1]``. +def _bound_text(value: float | None) -> str: + """A bound as it appears in a message: ``0``, ``1``, ``inf``. + + `%g` rather than `str`, so the default bounds still render as the ``[0, 1]`` the + error message has always said rather than as ``[0.0, 1.0]``. + """ + return "inf" if value is None else f"{value:g}" + + +def normalise_blend_range( + spec: MixAlphaSpec, + key: str = "mix_alpha", + owner: str = "", + hint: str = _MIX_ALPHA_HINT, + bounds: tuple[float, float | None] = (0.0, 1.0), +) -> tuple[float, float]: + """A scalar or ``[lo, hi]`` as a validated pair inside `bounds`. Checked here rather than in `smauglab/config.py`, which only ever asks whether a key is accepted, never what its value is -- every value check in this package happens in @@ -93,6 +108,11 @@ def normalise_blend_range(spec: MixAlphaSpec, key: str = "mix_alpha", owner: str `key` and `owner` are for the message alone. `AnatomicalLabelOverlay.blend_strength` is a different blend with the same shape of argument, and shares this. + + `bounds` is what makes it shareable by a parameter that is *not* a weight: `gpu/regions.py` + passes ``(0.0, None)`` for `seg_feather`, a Gaussian sigma in voxels with no upper + limit. The default is the unit interval every blend weight here needs, so the three + existing callers are unaffected -- including the message, which still reads ``[0, 1]``. """ where = f"{owner}: " if owner else "" if isinstance(spec, bool): # an int subtype, and never a blend weight anyone meant @@ -111,9 +131,10 @@ def normalise_blend_range(spec: MixAlphaSpec, key: str = "mix_alpha", owner: str lo, hi = values if lo > hi: raise ValueError(f"{where}{key} range must be non-decreasing, got [{lo}, {hi}]") - if lo < 0.0 or hi > 1.0: + low, high = bounds + if lo < low or (high is not None and hi > high): suffix = f" {hint}" if hint else "" - raise ValueError(f"{where}{key} must lie inside [0, 1], got [{lo}, {hi}].{suffix}") + raise ValueError(f"{where}{key} must lie inside [{_bound_text(low)}, {_bound_text(high)}], got [{lo}, {hi}].{suffix}") return (lo, hi) diff --git a/smauglab/transforms/gpu/regions.py b/smauglab/transforms/gpu/regions.py new file mode 100644 index 0000000..a34c543 --- /dev/null +++ b/smauglab/transforms/gpu/regions.py @@ -0,0 +1,825 @@ +"""Where a transform applies, and how sharply it stops applying. + +Three parameters predate this module -- `in_seg`, `out_seg` and `mix_in_out` -- and +between them they expressed about a third of what they looked like they expressed: + +* `in_seg` and `out_seg` were two independent uniforms, and *both fired* fell through + to the same ``"all"`` as *neither fired*. "A mixture of foreground and background" + was not expressible, and `in_seg=1, out_seg=1` was indistinguishable from + `in_seg=0, out_seg=0`. +* `mix_in_out` keyed its keep/drop draw on `seg_mask.shape[:2]`, the **channel** axis. + For the single-channel integer label map nnU-Net's trainer passes, that is `(B, 1)` + -- one draw per sample, deleting the whole mask half the time. So ``"in"`` applied + the transform nowhere and ``"out"`` applied it everywhere, and the thirteen shipped + configs that set `mix_in_out: true` were running a coin flip rather than a label mix + in every nnU-Net run. It only ever meant "per label" on a one-hot mask. +* The background was one blob. `out_seg` took all of it or none of it. +* The composite was `m * transformed + (1 - m) * orig` with a **binary** `m`, so every + region-selected augmentation left a one-voxel discontinuity sitting exactly on the + label boundary -- an edge perfectly correlated with the segmentation target, which is + the one thing an augmentation must not hand the network. + +The replacement separates four questions that were tangled together: + +================ ============================================================ +`in_seg`, **where**: foreground labels, background, a mixture of both, +`out_seg` or the whole patch. Both fired now means ``"both"``. +`seg_unit`, **at what granularity**: the whole side as one piece, one unit +`seg_unit_prob` per label value, or one unit per connected component. +`bg_split`, **how finely the background is divided**: one blob, or EM +`bg_em_clusters` intensity clusters, SynthSeg-style (Billot et al. 2023, 5.4). +`seg_feather` **how sharp the boundary is**: a hard step, or a Gaussian ramp. +================ ============================================================ + +Everything resolves to a single integer **unit id map** (:class:`RegionUnits`) and one +gather off it. Nothing materialises per-unit volumes: a noisy mask under +``seg_unit="component"`` can hold hundreds of components, and `[2, 300, 128, 128, 128]` +of bool is 1.2 GB. + +Three properties this module goes out of its way to keep, each of which the code it +replaces broke: + +**A knob that is off spends no random numbers.** `_choose_region_mode` drew its two +uniforms whenever a mask was attached, including at `in_seg = out_seg = 0`, so an +unselected run and a selected one at the same seed were never comparable. That is the +defect `gpu/mixing.py` removed for `mix_prob`, and :func:`choose_region_modes` returns +before drawing anything for the same reason. + +**The number of uniforms a draw spends is a function of the configuration, never of the +data.** A per-unit keep draw would spend one uniform per label *present*, so a patch +with four labels and one with five would leave the stream in different places and every +transform after this one would diverge on the data. :data:`UNIT_DRAW_WIDTH` fixes the +width per decomposition kind instead. + +**The decomposition is computed once per pipeline pass, not once per transform.** See +:func:`region_units`. Connected components and EM each cost tens of milliseconds on a +128^3 patch; a pipeline runs twenty-odd region-selecting transforms over one batch. + +One consequence worth stating plainly rather than leaving to be discovered: +``bg_split="em"`` is the only decomposition that reads the image, and the image moves as +the pass runs -- every transform's first statement is `input = input.clone()`. The +clusters are therefore fitted **once per pass, from the image as it stood when the first +region-selecting transform asked**, and not re-fitted afterwards. That is deliberate +twice over. Re-fitting is twenty-odd EM runs a batch, each with a 100 000-element +subsample draw and a host synchronisation per sample; and a single fit means two +region-selected transforms in the same pass cut the background along the same lines +instead of each along its own. What it costs is that a transform late in the pass +clusters an image a few transforms old. The units are a partition to switch on and off, +not a measurement, so that is the cheap half of the trade. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any, NamedTuple, Protocol, Union + +import torch +from torch import Tensor + +from smauglab.transforms.gpu.mixing import normalise_blend_range +from smauglab.transforms.kernels import gaussian_blur3d +from smauglab.transforms.rng import next_shared_seed + +#: A fixed sigma in voxels, or a ``[lo, hi]`` range drawn from per call. Mirrors +#: `MixAlphaSpec` in gpu/mixing.py, which is the same shape of argument for a weight +#: rather than a width. +SegFeatherSpec = Union[float, Sequence[float]] + +SEG_UNITS = ("whole", "label", "component") +BG_SPLITS = ("none", "em") + +DEFAULT_SEG_UNIT = "whole" +DEFAULT_SEG_UNIT_PROB = 0.5 +DEFAULT_BG_SPLIT = "none" +DEFAULT_BG_EM_CLUSTERS = (3, 10) +DEFAULT_SEG_FEATHER = 0.0 + +#: How many keep/drop uniforms a draw spends, per decomposition kind. Fixed per *kind* +#: so the count is a function of the configuration and never of the segmentation -- see +#: the module docstring. A unit id at or above the width folds into the last unit, which +#: only bites a patch with more than 32 label values or more than 256 components per +#: side, and then only by tying their fates together. +UNIT_DRAW_WIDTH = {"whole": 2, "label": 32, "component": 256} + +#: Mode codes, as they travel in the `[N]` tensor `choose_region_modes` returns. +MODE_ALL = 0 +MODE_IN = 1 +MODE_OUT = 2 +MODE_BOTH = 3 + +#: A label value no real label can take, for the gated shifts in +#: :func:`connected_components_3d`: the out-of-volume neighbour must never compare equal +#: to the voxel it is shifted past. +_NO_LABEL = -1 + +#: Where the attribute cache lives on a mask tensor. See :func:`region_units`. +_CACHE_ATTR = "_smauglab_region_units" + +#: Appended to a `seg_feather` bound error, because "sigma in voxels" is not what the +#: name suggests to someone who has just read `mix_alpha`. +_FEATHER_HINT = "It is a Gaussian sigma in voxels, not a weight: 0 is a hard boundary, and there is no upper limit." + + +class RegionUnits(NamedTuple): + """A segmentation decomposed into the units a mode can switch on and off. + + `ids` is dense in ``[0, n_units)`` and background units come **first**, so "is this + unit foreground" is the scalar comparison ``id >= n_background`` rather than a + per-unit tensor. Label units are numbered batch-wide, so a given label value is the + same unit in every sample of a batch -- the choice `seg_region_masks` already + documents for its channel ordering. Component and EM-cluster units cannot be: + they are per-sample by nature, and sample `b`'s unit `k` is unrelated to sample + `b+1`'s. The keep draw is per sample either way, so that costs nothing. + """ + + ids: Tensor + n_units: int + n_background: int + + +class UnitSpec(NamedTuple): + """What a decomposition depends on, and therefore the cache key. + + `seg_unit_prob`, `in_seg`, `out_seg` and `seg_feather` are deliberately absent: they + are redrawn per call and must not fragment the cache. + """ + + seg_unit: str + bg_split: str + bg_em_clusters: tuple[int, int] + + +class RegionSelecting(Protocol): + """What the helpers below need off the transform they are handed. + + A Protocol rather than `ImageOnlyTransform`, for the same reason `mixing.Mixing` is + one: these are plain attributes on an `nn.Module`, so reading them through the base + class resolves via `Module.__getattr__`, which is typed as returning + `Tensor | Module` -- and neither of those is a float. Naming them here is both what + makes the arithmetic below type-check and a statement of the requirement. + """ + + in_seg: float + out_seg: float + seg_unit: str + seg_unit_prob: float + bg_split: str + bg_em_clusters: tuple[int, int] + seg_feather: tuple[float, float] + + +# ── validation ─────────────────────────────────────────────────────────────── + + +def _one_of(value: Any, choices: tuple[str, ...], key: str, owner: str) -> str: + where = f"{owner}: " if owner else "" + if not isinstance(value, str) or value not in choices: + raise ValueError(f"{where}{key} must be one of {', '.join(choices)}, got {value!r}") + return value + + +def _cluster_range(spec: Any, owner: str) -> tuple[int, int]: + """A ``[lo, hi]`` pair of positive cluster counts. + + Not `normalise_blend_range`: these are counts, and silently accepting a float would + put a non-integer into `torch.randint`'s bounds. + """ + where = f"{owner}: " if owner else "" + if isinstance(spec, (bool, int, float)): + raise ValueError(f"{where}bg_em_clusters must be a [lo, hi] pair of counts, got {spec!r}") # noqa: TRY004 + try: + values = [int(v) for v in spec] + except (TypeError, ValueError) as exc: + raise ValueError(f"{where}bg_em_clusters must be a [lo, hi] pair of counts, got {spec!r}") from exc + if len(values) != 2: + raise ValueError(f"{where}bg_em_clusters must be a [lo, hi] pair of counts, got {len(values)} values: {spec!r}") + lo, hi = values + if lo < 1 or hi < lo: + raise ValueError(f"{where}bg_em_clusters must be a non-decreasing pair of counts >= 1, got [{lo}, {hi}]") + if hi > UNIT_DRAW_WIDTH["whole"] + max(UNIT_DRAW_WIDTH.values()): + # Only a sanity bound: a background split into thousands of clusters is a + # configuration mistake, and the keep draw would fold most of them together. + raise ValueError(f"{where}bg_em_clusters upper bound {hi} is implausibly large") + return (lo, hi) + + +def configure_regions( + transform: RegionSelecting, + in_seg: float, + out_seg: float, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, +) -> None: + """Validate and assign the region parameters. One line per constructor. + + Every value check in this package happens in a constructor -- `smauglab/config.py` + only ever asks whether a key is *accepted*, never what its value is. A pipeline is + built before the first batch, so a bad value still fails long before training. + + `in_seg` and `out_seg` used to be clamped into range silently + (``max(0.0, min(1.0, p))``), which quietly accepted `in_seg: 5.0` and + `in_seg: true`. They are rejected now. + """ + owner = type(transform).__name__ + transform.in_seg = normalise_blend_range(in_seg, key="in_seg", owner=owner, hint="")[0] + transform.out_seg = normalise_blend_range(out_seg, key="out_seg", owner=owner, hint="")[0] + transform.seg_unit = _one_of(seg_unit, SEG_UNITS, "seg_unit", owner) + transform.seg_unit_prob = normalise_blend_range(seg_unit_prob, key="seg_unit_prob", owner=owner, hint="")[0] + transform.bg_split = _one_of(bg_split, BG_SPLITS, "bg_split", owner) + transform.bg_em_clusters = _cluster_range(bg_em_clusters, owner) + transform.seg_feather = normalise_blend_range(seg_feather, key="seg_feather", owner=owner, hint=_FEATHER_HINT, bounds=(0.0, None)) + + +def unit_spec(transform: RegionSelecting) -> UnitSpec: + """The decomposition `transform` asks for, which is also its cache key.""" + return UnitSpec(transform.seg_unit, transform.bg_split, transform.bg_em_clusters) + + +def region_selection_active(transform: RegionSelecting) -> bool: + """Whether either probability can fire. False means every helper here is a no-op.""" + return transform.in_seg > 0.0 or transform.out_seg > 0.0 + + +# ── segmentation layout ────────────────────────────────────────────────────── + + +def label_map_from(mask: Tensor) -> Tensor: + """A `[B, 1, *spatial]` integer label map, background 0, from any of three layouts. + + The layouts reach a transform interchangeably and the values alone cannot always + tell them apart: + + * ``[B, 1, *spatial]`` -- an ordinary label map. The ids are the values and + background is 0. Rounded first: a mask that has been through a geometric transform + comes back as float, and unrounded floats would invent labels. + * ``[B, C, *spatial]`` that leaves some voxel unset -- a background-*implicit* + one-hot, this repository's convention (`collapse_onehot_to_index` in + gpu/fromSeg.py): channel `c` encodes label `c + 1`. + * ``[B, C, *spatial]`` that leaves no voxel unset -- a background-*inclusive* + one-hot, which is what `seg_region_masks` emits for a single-channel label map and + what `unit_tests/test_seg_layout.py` builds. Channel 0 **is** background. + + The third is the one `_foreground` already had to detect by coverage rather than by + value, and the one `collapse_onehot_to_index` still gets wrong: it does `argmax + 1` + unconditionally, so a background-inclusive one-hot comes back with every label + shifted by one and background relabelled 1. That is why `em_subdivide_labels` carries + a fallback for "the real background is no longer 0". Deriving the map here, once, is + what lets everything else in this module be a read off it. + + One case no function can read off the values: a background-implicit one-hot patch + labelled edge to edge is indistinguishable from a background-inclusive one. A real + segmentation patch is not labelled edge to edge, which is the same admission + `_foreground`'s docstring has always carried. + """ + if mask.dim() < 3: + raise ValueError(f"expected [B, C, *spatial], got {tuple(mask.shape)}") + if mask.shape[1] == 1: + return mask.round().long().clamp_min(0) + + covered = mask.amax(dim=1, keepdim=True) > 0 + index = mask.argmax(dim=1, keepdim=True).long() + # `torch.where`, not `if bool(covered.all())`: reading that test back to Python + # stalls the host on everything queued behind it, and both branches are a cheap + # elementwise pass. Same reason `_foreground` is written this way. + return torch.where(covered.all(), index, (index + 1) * covered) + + +def _foreground(mask: Tensor, dim: int) -> Tensor: + """Which voxels the segmentation covers, reducing over the class axis. + + This used to be `torch.argmax(mask, dim) > 0`, which is only "is anything labelled + here" if class 0 is background -- and it is not: + + * For a single-channel `[B, 1, D, H, W]` mask (an ordinary nnU-Net target, and what + the tests build) `argmax` over a length-1 axis is always 0, so the result was + **all False**. `in_seg` then applied the transform nowhere and `out_seg` applied + it everywhere: the two knobs did nothing and the opposite of nothing. + * For a one-hot mask, this repository's convention (`collapse_onehot_to_index` in + gpu/fromSeg.py) is that channel `c` encodes label `c + 1` with background + implicit, so `argmax == 0` is a real foreground class and was being dropped. + + `amax > 0` asks the question that was meant, and matches what + `collapse_onehot_to_index` already does with `seg_raw.any(dim=1)`. + + That still leaves one layout it cannot read off the values alone: a one-hot + that *includes* a background channel, which is what `seg_region_masks` emits + for a single-channel label map and what `unit_tests/test_seg_layout.py` + builds. Every voxel then has some channel set, `amax > 0` is True everywhere, + and `in_seg` applies the transform to the whole patch while `out_seg` applies + it nowhere -- the exact failure this function was written to fix, in the one + layout it was not checked against. + + A background-inclusive one-hot is recognisable: it is multi-channel and leaves + no voxel unset, because the background channel covers whatever the foreground + channels do not. A background-implicit one-hot always has all-zero background + voxels unless the patch is labelled edge to edge, which a real segmentation + patch is not. So full coverage plus more than one channel means channel 0 is + background, and the reduction skips it. + + Kept as its own function, rather than folded into :func:`label_map_from`, because + `dim` is not always 1: `_apply_region_mode`'s unbatched path reduces a `[C, *spatial]` + mask over axis 0. The two agree on every layout, which + `unit_tests/test_foreground_layouts.py` is the check on. + """ + covered = mask.amax(dim=dim) > 0 + if mask.shape[dim] == 1: + return covered + # `torch.where`, not `if bool(...)`: reading that test back to Python stalls the host + # on everything queued behind it. Both branches are a reduction over the class axis, + # so computing the unused one costs a fraction of what the stall does. + without_background = mask.narrow(dim, 1, mask.shape[dim] - 1).amax(dim=dim) > 0 + return torch.where(covered.all(), without_background, covered) + + +def densify_region_ids(region_ids: Tensor) -> tuple[Tensor, int]: + """Remap sparse non-negative integer region ids onto a contiguous [0, R) range. + + `em_subdivide_labels` encodes a fine id as `parent_idx * mult + assign`, which + is sparse. Anything that indexes a per-region tensor of length R with these -- the + keep/drop gather here, `signed_alpha_affine_remap` in the PALETTE blocks -- would + index out of bounds on a sparse id. + + Lives here rather than in `gpu/palette/partitioners.py`, where it was written, + because this module cannot import that one: `palette/base.py` reaches `gpu/base.py`, + which imports this module for its "off" defaults. Re-exported from there. + """ + # A lookup table off a bincount, not `torch.unique(return_inverse=True)`: the ids + # are small non-negative integers, so counting the bins and numbering the + # non-empty ones gives the same ascending remap without radix-sorting two + # million voxels. + counts = torch.bincount(region_ids.reshape(-1)) + present = counts.nonzero().flatten() + lookup = torch.zeros(counts.numel(), dtype=torch.long, device=region_ids.device) + lookup[present] = torch.arange(present.numel(), device=region_ids.device) + return lookup[region_ids], int(present.numel()) + + +# ── connected components ───────────────────────────────────────────────────── + +#: Hard ceiling on the propagation below. Pointer jumping squares the reach each round, +#: so even a serpentine component spanning a 128^3 patch settles in about twenty; the cap +#: only exists so a bug cannot hang a training run. +_CC_MAX_ITERS = 64 + +#: How often the fixpoint test runs. `bool(...)` on a device tensor stalls the host until +#: the queue drains, and the loop body is six shifts over the volume, so testing every +#: round would spend a synchronisation to save work that is already cheap. Compact blobs +#: converge in a handful of rounds, which is why the test exists at all rather than +#: running a fixed count. Same trade `_kmeans_1d` made in the other direction. +_CC_CHECK_EVERY = 4 + + +def _shift(volume: Tensor, dim: int, delta: int, fill: Any) -> Tensor: + """`volume` shifted along `dim` so that index `i` holds what index `i - delta` held. + + `fill` enters at the face that is uncovered. Not `torch.roll`, which wraps the far + face round to the near one and would join two components that only touch through + opposite sides of the patch. + """ + size = volume.shape[dim] + step = abs(delta) + kept = volume.narrow(dim, 0, size - step) if delta > 0 else volume.narrow(dim, step, size - step) + pad_shape = list(volume.shape) + pad_shape[dim] = step + block = torch.full(pad_shape, fill, dtype=volume.dtype, device=volume.device) + return torch.cat([block, kept] if delta > 0 else [kept, block], dim=dim) + + +def connected_components_3d(labels: Tensor, foreground: Tensor) -> tuple[Tensor, int]: + """Connected components of each label value, 6-connected, per sample. + + `labels` is `[B, 1, D, H, W]` integer and `foreground` the bool mask of the voxels to + partition. Returns `(ids, n_max)` where `ids` is `[B, 1, D, H, W]`, numbered from 0 + within each sample and `-1` outside `foreground`, and `n_max` is the largest + per-sample count -- which is what sizes the keep/drop draw. + + **6-connected**, matching `scipy.ndimage.label`'s default structuring element + (`generate_binary_structure(3, 1)`), which is also the conventional reading of + "connected components of a label": two voxels touching only at a corner are two + components. The obvious implementation -- one `max_pool3d(kernel=3)` -- is a 3^3 + neighbourhood and therefore 26-connected, which merges them. Six explicit shifts + cost twice the kernels of three axis pools and buy the right answer. + + Components respect label boundaries without a pass per label value. The gate is the + `same` mask below: a neighbour contributes only where it carries the *same* label, so + one propagation partitions every label at once. Doing it with pooled values instead + would need a pass per label (a spine patch has five to twenty-five), because a pool + reports the largest value in the neighbourhood without saying which voxel it came + from, and a higher-labelled neighbour would mask the same-labelled one. + + The algorithm is label-equivalence propagation: every voxel starts as its own + component, named by its flat index; each round takes the largest name reachable from + a same-label neighbour, then **jumps to that name's own current name**, which doubles + the distance information travels per round and turns a linear walk along a long thin + structure into a logarithmic one. + + Indices are `int64` throughout. The float alternative is tempting because + `max_pool3d` requires it, but `float32` represents integers exactly only to 2^24, and + a 256^3 patch has exactly that many voxels -- so a larger patch would silently merge + components whose flat indices round together. + """ + batch = labels.shape[0] + spatial = labels.shape[2:] + voxels = int(labels[0, 0].numel()) + device = labels.device + + # Gated once, outside the loop: the label map does not move, so the six + # "neighbour carries the same label, and both of us are foreground" masks are + # constants of the call. Recomputing them per round would double its cost. + same = [] + for dim in (2, 3, 4): + for delta in (1, -1): + neighbour_label = _shift(labels, dim, delta, _NO_LABEL) + neighbour_fg = _shift(foreground, dim, delta, False) + same.append(((neighbour_label == labels) & neighbour_fg & foreground, dim, delta)) + + # Globally unique seeds, not per-sample ones: the densify below counts over the whole + # batch, so two samples whose components happened to root at the same within-sample + # index would be counted as one. + ids = torch.arange(batch * voxels, device=device, dtype=torch.long).view(batch, 1, *spatial) + ids = torch.where(foreground, ids, torch.full_like(ids, -1)) + + for iteration in range(1, _CC_MAX_ITERS + 1): + best = ids + for gate, dim, delta in same: + best = torch.maximum(best, torch.where(gate, _shift(ids, dim, delta, -1), torch.full_like(ids, -1))) + flat = best.flatten(1) + # The name that our current name currently answers to. `clamp_min(0)` only keeps + # the gather in range for the background entries, which the `where` then drops. + jumped = torch.gather(flat, 1, (flat - batch_offset(batch, voxels, device)).clamp_min(0)) + advanced = torch.maximum(flat, torch.where(flat >= 0, jumped, flat)).view_as(ids) + advanced = torch.where(foreground, advanced, torch.full_like(advanced, -1)) + settled = iteration % _CC_CHECK_EVERY == 0 and bool(torch.equal(advanced, ids)) + ids = advanced + if settled: + break + + # Densified over the batch, then rebased per sample. Global ids are ascending in the + # root's flat index and every index of sample b is below every index of sample b+1, so + # each sample's ids are a contiguous ascending block and subtracting its lowest gives + # a 0-based numbering without a second densify per sample. + dense, _total = densify_region_ids(ids.clamp_min(0)) + dense = torch.where(foreground, dense, torch.full_like(dense, -1)) + flat_dense = dense.flatten(1) + sentinel = torch.full_like(flat_dense, torch.iinfo(torch.int64).max) + lowest = torch.where(flat_dense >= 0, flat_dense, sentinel).min(dim=1).values + highest = flat_dense.max(dim=1).values + counts = torch.where(highest >= lowest, highest - lowest + 1, torch.zeros_like(highest)) + rebased = torch.where(foreground, dense - lowest.view(-1, *([1] * (dense.dim() - 1))), torch.full_like(dense, -1)) + return rebased, max(int(counts.max()), 1) + + +def batch_offset(batch: int, voxels: int, device: torch.device) -> Tensor: + """`[B, 1]` of each sample's first global seed index. + + The names in :func:`connected_components_3d` are global so the densify can count them + together, but `torch.gather` indexes within a row, so the row's offset comes off again + before the lookup. + """ + return (torch.arange(batch, device=device, dtype=torch.long) * voxels).view(batch, 1) + + +# ── background decomposition ───────────────────────────────────────────────── + + +def background_em_units(image: Tensor, labels: Tensor, clusters_range: tuple[int, int]) -> tuple[Tensor, int]: + """The background split into intensity-coherent subregions by EM, SynthSeg-style. + + Returns `(ids, n)` with `ids` numbered from 0 on background voxels and 0 elsewhere + (foreground ids are never read: the caller overwrites them with the foreground + numbering). + + `em_subdivide_labels` does the work -- this is a two-label call that asks it to split + label 0 into `clusters_range` clusters and leave the foreground as a single one, which + is exactly the paper's background handling without its foreground handling. + + Two things are wrapped around it rather than left to it: + + * **The draw is forked.** It spends one `randint` per sample *and* a + `max_fit_voxels`-element subsample index -- a hundred thousand uniforms, which the + background of a 128^3 patch always triggers. Since the result is cached, a miss + would spend them and a hit would not, which would make the global RNG stream a + function of cache state: of which transform asked first, of which rows `p` selected, + of whether a geometric transform replaced the mask. Forking keeps the stream exactly + where it was found. `next_shared_seed` rather than an arbitrary seed so DDP ranks + still agree. + * **A patch with no background is left alone.** When the configured background label + is absent, `em_subdivide_labels` falls back to the largest-area label, which here + would silently split a *foreground* structure into three to ten pieces and call them + background. + + `same_on_batch` is fixed at False: the decomposition is shared by every transform in + the pass, so it cannot follow any one transform's flag. + """ + # Imported here, not at module scope. `smauglab/transforms/synthseg/__init__.py` + # imports synthseg/transforms.py, which imports gpu/base.py -- and gpu/base.py imports + # this module for its "off" defaults. A module-scope import would close that loop and + # break every `import smauglab.transforms.gpu.base`; `unit_tests/test_imports.py` is + # the canary. By the time any transform calls this, the package is fully imported. + from smauglab.transforms.synthseg.functional import em_subdivide_labels + + background = labels <= 0 + if not bool(background.any()): + return torch.zeros_like(labels), 1 + + binary = (~background).long() + device = image.device + with torch.random.fork_rng(devices=[device] if device.type == "cuda" else [], enabled=True): + torch.manual_seed(next_shared_seed()) + fine, _generation, _merge = em_subdivide_labels( + image, + binary, + n_foreground_clusters=1, + background_clusters_range=clusters_range, + background_label=0, + same_on_batch=False, + ) + # Background is parent index 0, so its fine ids are `0 * mult + assign` -- already the + # 0-based cluster numbering this returns. Foreground ids are zeroed because the caller + # replaces them. + ids = torch.where(background, fine, torch.zeros_like(fine)) + return ids, max(int(ids.max()) + 1, 1) + + +# ── the decomposition, and its cache ───────────────────────────────────────── + + +def _label_units(labels: Tensor) -> tuple[Tensor, int]: + """One unit per distinct positive label value, numbered batch-wide from 0.""" + counts = torch.bincount(labels.reshape(-1).clamp_min(0)) + present = counts.nonzero().flatten() + positive = present[present > 0] + lookup = torch.zeros(counts.numel(), dtype=torch.long, device=labels.device) + lookup[positive] = torch.arange(positive.numel(), device=labels.device) + return lookup[labels.clamp_min(0)], max(int(positive.numel()), 1) + + +def _as_batched(volume: Tensor) -> Tensor: + """`[C, *spatial]` or `[B, C, *spatial]` as `[B, C, *spatial]`. + + The region helpers are called from two shapes of `apply_transform`: the vectorised + ones hand over the whole batch, and the per-sample ones (`RandomInverseGPU`, + `RandomDomainTransferGPU`) hand over one sample with no batch axis. + """ + return volume if volume.dim() == 5 else volume.unsqueeze(0) + + +def _build_units(mask: Tensor, image: Tensor | None, spec: UnitSpec) -> RegionUnits: + labels = label_map_from(_as_batched(mask)) + foreground = labels > 0 + + if spec.bg_split == "em" and image is not None: + background_ids, n_background = background_em_units(_as_batched(image), labels, spec.bg_em_clusters) + else: + background_ids, n_background = torch.zeros_like(labels), 1 + + if spec.seg_unit == "label": + foreground_ids, n_foreground = _label_units(labels) + elif spec.seg_unit == "component": + foreground_ids, n_foreground = connected_components_3d(labels, foreground) + else: + foreground_ids, n_foreground = torch.zeros_like(labels), 1 + + ids = torch.where(foreground, n_background + foreground_ids.clamp_min(0), background_ids) + return RegionUnits(ids, n_background + n_foreground, n_background) + + +def region_units(mask: Tensor, image: Tensor | None, spec: UnitSpec) -> RegionUnits: + """The decomposition of `mask`, computed once per mask tensor per spec. + + Connected components and EM each cost tens of milliseconds on a 128^3 patch, and a + pipeline runs twenty-odd region-selecting transforms over one batch. Without a cache + the feature is not trainable; with one, every transform of a pass shares the same + units, which is also what makes two region-selected transforms cut the background + along the same lines rather than each along its own. + + **The cache is an attribute on the mask tensor**, keyed by spec. That choice is load + bearing in three ways: + + * **Lifetime.** It dies with the tensor, so there is nothing to bound, evict or clear. + * **Invalidation.** `MaskSequentialOpsCustom.transform` hands back the *same* tensor + object for an image-only module and a *new* one for a geometric module, so the + identity changes exactly when the mask changes. A geometric transform mid-pipeline + invalidates this for free. + * **Not a dict keyed on tensors.** `Tensor.__hash__` is identity-based but + `Tensor.__eq__` is elementwise, so a hash-bucket collision between two different + masks makes `dict` evaluate `bool(a == b)` and raise *"Boolean value of Tensor with + more than one value is ambiguous"*. Measured: a `WeakKeyDictionary` with five + hundred same-shaped tensor keys raises reliably. An attribute has no buckets. + + Keyed on the mask **as the transform was handed it**, which is what keeps the image + and the units row-aligned. kornia slices the image down to the rows that drew below + `p` and `_seg_for_applied_rows` slices the mask to match, so a decomposition built + from the full batch would pair B-row units with a B'-row image -- and `bg_split="em"` + reads that image. When every row applies, which is the common case and always true at + `p = 1`, the slice *is* the full mask object and every module of the pass hits the + same entry. A module whose row selection is a strict subset pays for its own + decomposition, which is the honest price of having the rows line up. + """ + cached = getattr(mask, _CACHE_ATTR, None) + if cached is not None and spec in cached: + return cached[spec] + units = _build_units(mask, image, spec) + if cached is not None: + cached[spec] = units + else: + try: + setattr(mask, _CACHE_ATTR, {spec: units}) + except AttributeError: # pragma: no cover -- a tensor subclass that forbids it + pass + return units + + +# ── the per-call draws ─────────────────────────────────────────────────────── + + +def choose_region_modes(transform: RegionSelecting, batch: int, device: torch.device) -> Tensor | None: + """One mode per sample, or None when neither probability can fire. + + Returns `MODE_ALL` / `MODE_IN` / `MODE_OUT` / `MODE_BOTH` codes. **Both fired means + `MODE_BOTH`**, where it used to fall through to `MODE_ALL` -- so `in_seg=1, out_seg=1` + stopped being indistinguishable from `in_seg=0, out_seg=0`. + + Returns None, **having drawn nothing**, when both probabilities are zero. The version + this replaces drew its two uniforms whenever a mask was attached, which is the defect + `984f037` removed for `mix_prob`: with the draws unconditional, a run with region + selection off and a run with it on were never on the same stream, so no sweep could be + compared against its own baseline at a fixed seed. + + Per sample, not per batch. Ten of the eleven call sites used to take one decision for + the whole batch, so `in_seg=0.5` meant "half the batches" rather than "half the + samples" -- while `p` and `mix_prob` were per-sample all along. On the image's device + and with no `dtype=`, mirroring `mix_with_original`, so a half-precision batch stays on + the same stream as a single-precision one. + """ + if not region_selection_active(transform): + return None + draws = torch.rand(batch, 2, device=device) + fired_in = draws[:, 0] < transform.in_seg + fired_out = draws[:, 1] < transform.out_seg + # The code arithmetic is the definition: 0 neither, 1 in, 2 out, 3 both. + return fired_in.long() + 2 * fired_out.long() + + +def draw_feather(transform: RegionSelecting) -> float: + """The feather sigma for this call, in voxels. Zero means a hard boundary. + + One draw per call rather than per sample: `gaussian_blur3d` reads its sigmas back to + the host (`.tolist()`), so a per-sample sigma would be one blur call and one + synchronisation per sample. The modes and the per-unit keeps stay per sample; a + boundary *width* is not a decision about where. From the CPU default generator as a + Python float, the way `blend_one` draws its gate, so nothing has to be read off the + device. + + Spends nothing when `seg_feather` is zero, and otherwise spends exactly one uniform + whatever the modes came out as. Gating the draw on "did any sample pick a mode that + uses it" would make the stream a function of the data. Drawn even when `lo == hi`, for + the reason `mix_alpha` is: otherwise the stream becomes a function of the range, and + two runs differing only in `seg_feather` would not be comparable. + """ + low, high = transform.seg_feather + if high <= 0.0: + return 0.0 + return float(torch.rand(1).item()) * (high - low) + low + + +def region_alpha(transform: RegionSelecting, units: RegionUnits, modes: Tensor) -> Tensor: + """The `[B, 1, *spatial]` composite weight: 1 where the transform applies, 0 where not. + + Built as one gather off the unit id map. Nothing per-unit is materialised: a noisy + mask under `seg_unit="component"` can hold hundreds of components, and + `[2, 300, 128, 128, 128]` of bool is 1.2 GB. + + The keep/drop draw is a fixed-width block (:data:`UNIT_DRAW_WIDTH`) so the number of + uniforms it spends depends on the configuration and not on how many labels this + particular patch happens to contain -- otherwise every transform after this one would + diverge between a four-label patch and a five-label one. Unit ids at or above the + width fold into the last column, which ties those units' fates together rather than + reading past the draw. + + The fall-back rule -- **if a side kept nothing, keep all of it** -- is applied per + side, and costs no extra draw. It is what makes the degenerate cases come out right: + an empty mask has no occupied foreground unit, so `MODE_IN` is a no-op and `MODE_OUT` + covers everything, which is exactly what the binary version did. + """ + device = units.ids.device + width = UNIT_DRAW_WIDTH[transform.seg_unit] + transform.bg_em_clusters[1] + columns = min(units.n_units, width) + column = units.ids.clamp(0, columns - 1) + + keep = torch.rand(modes.numel(), width, device=device)[:, :columns] < transform.seg_unit_prob + # Background units come first, and only the high end is folded, so a column is + # foreground exactly when its index is at or past `n_background`. + is_foreground = torch.arange(columns, device=device) >= units.n_background + wants_foreground = (modes == MODE_IN) | (modes == MODE_BOTH) + wants_background = (modes == MODE_OUT) | (modes == MODE_BOTH) + wanted = torch.where(is_foreground.unsqueeze(0), wants_foreground.unsqueeze(1), wants_background.unsqueeze(1)) + + selected = keep & wanted + for side in (is_foreground, ~is_foreground): + on_this_side = selected & side + nothing_kept = ~on_this_side.any(dim=1, keepdim=True) + selected = selected | (nothing_kept & wanted & side) + + alpha = torch.gather(selected, 1, column.flatten(1)).view(column.shape) + everywhere = (modes == MODE_ALL).view(-1, *([1] * (alpha.dim() - 1))) + return (alpha | everywhere).to(torch.float32) + + +def feather_alpha(alpha: Tensor, sigma: float) -> Tensor: + """`alpha` with its boundaries turned from a step into a Gaussian ramp. + + Reflect padding, which is `gaussian_blur3d`'s default and the reason to use it rather + than a local copy: zero padding pulls the weight toward 0 at the patch border, so a + unit touching the border would get a soft edge that is an artefact of the patch rather + than of the segmentation. `gpu/fromSeg.py` carries a note about exactly that bug in + the local blur it replaced. + + Every boundary in `alpha` is feathered, not only foreground against background: under + `seg_unit="label"` or `"component"` that includes the boundary between a kept unit and + a dropped neighbour, and under `bg_split="em"` the boundaries between clusters -- which + is what stops an EM split being visible as a hard intensity edge. + + The ramp is centred on the boundary and passes through 0.5 there, so feathering + shrinks the full-strength region slightly. Dilating by `sigma` first would grow the + ramp outward instead and keep each unit at full strength; that is one extra + `max_pool3d` and a second parameter, and worth adding only if the shrink turns out to + matter. + + Two-dimensional spatial inputs pass through unfeathered: the blur is 3-D. The sigma is + still drawn, because the draw must not depend on the data. + """ + if sigma <= 0.0 or alpha.dim() != 5: + return alpha + return gaussian_blur3d(alpha, sigma).clamp(0.0, 1.0) + + +# ── the public entry point ─────────────────────────────────────────────────── + + +def _rescaled_into(source: Tensor, target: Tensor, reduce_dims: tuple[int, ...]) -> Tensor: + """`source` linearly remapped onto `target`'s per-sample min-max range.""" + source_min = torch.amin(source, dim=reduce_dims, keepdim=True) + source_max = torch.amax(source, dim=reduce_dims, keepdim=True) + target_min = torch.amin(target, dim=reduce_dims, keepdim=True) + target_max = torch.amax(target, dim=reduce_dims, keepdim=True) + return (source - source_min) / (source_max - source_min + 1e-8) * (target_max - target_min) + target_min + + +def apply_region_selection( + transform: RegionSelecting, + orig: Tensor, + transformed: Tensor, + seg: Tensor | None, + normalize: bool = False, +) -> Tensor: + """Composite `transformed` over `orig` wherever the drawn region says it applies. + + `orig` and `transformed` are either `[N, *spatial]` -- one channel of a batch, which is + what most `apply_transform` methods work in -- or `[*spatial]` for the transforms that + assemble their result one sample at a time. `seg` is the segmentation in any of the + three layouts :func:`label_map_from` reads. + + `normalize` rescales `transformed` into `orig`'s range first. That matters for the + transforms whose output lives in a different space from their input: a SynthSeg or + PALETTE synthesis is in `[0, 1]` while the patch it is composited onto is z-scored, so + a raw composite would put an intensity step on the mask boundary that has nothing to do + with the augmentation. The parameter existed on the function this replaces and had no + caller; these are the callers it was written for. + + Returns `transformed` untouched, having drawn nothing, when there is no segmentation or + when neither probability can fire. + """ + if seg is None or not region_selection_active(transform): + return transformed + + batched = orig.dim() >= 4 + reference = orig if batched else orig.unsqueeze(0) + modes = choose_region_modes(transform, reference.shape[0], orig.device) + assert modes is not None # region_selection_active was just checked + sigma = draw_feather(transform) + + units = region_units(seg, reference.unsqueeze(1), unit_spec(transform)) + if units.ids.shape[0] != reference.shape[0]: + # kornia slices the image down to the rows that drew below `p`, and + # `_seg_for_applied_rows` slices the mask to match, so these agree by + # construction. Said out loud because the alternative is a broadcast that + # silently pairs one sample's units with another sample's intensities. + raise ValueError( + f"{type(transform).__name__}: the segmentation has {units.ids.shape[0]} rows but the image has " + f"{reference.shape[0]}; params['seg'] must be sliced to the rows apply_transform was handed." + ) + alpha = feather_alpha(region_alpha(transform, units, modes), sigma) + # `[B, 1, *spatial]` down to the caller's shape. The channel axis exists only so the + # blur above sees a 5-D volume. + weight = alpha.squeeze(1).to(transformed.dtype) + if not batched: + weight = weight.squeeze(0) + + if normalize: + reduce_dims = tuple(range(1, orig.dim())) if batched else tuple(range(orig.dim())) + transformed = _rescaled_into(transformed, orig, reduce_dims) + return weight * transformed + (1.0 - weight) * orig From 8df003f38c2e66d4fb550bb74932828d5e4fccbf Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 15 Oct 2026 12:15:16 +0000 Subject: [PATCH 2/3] feat: give every GPU augmentation a configurable region, granularity and boundary Wires `gpu/regions.py` into all twenty-eight transforms that read the segmentation, replacing the three helpers in gpu/contrast.py. The surface, on every transform that had `in_seg` / `out_seg` and on the four TAs that had neither: in_seg / out_seg where: foreground labels, background, both (a mixture), or the whole patch seg_unit what a unit is: "whole", "label", "component" seg_unit_prob the chance each unit takes part bg_split "em" splits the background into intensity bg_em_clusters clusters so part of it can be used seg_feather Gaussian sigma for the boundary, 0 = a hard step `mix_in_out` is gone. It keyed its keep/drop draw on the mask's *channel* axis, which for nnU-Net's single-channel label map is one draw per sample that deletes the whole mask half the time -- so `in_seg` applied the transform nowhere and `out_seg` applied it everywhere. All thirty-two shipped configs that set it are migrated to `seg_unit="label"`, which is what it was trying to be; this is a behaviour fix, not a rename, and the keep draw also moves from `randint(0, 2, ...)` to `rand(...) < prob`. `RandomRedistributeSegGPU.in_seg` becomes `confine_prob`, on both backends. It never meant what `in_seg` means elsewhere -- it confines each region's Gaussian to that region's own voxels -- and twenty-eight shipped configs set it. So this is the one transform that does not accept `in_seg` / `out_seg`: had it accepted them, every stale config would have kept validating while silently changing meaning, and no diagnostic was possible, because difflib scores `in_seg` against `out_seg` at 0.615 and would answer "did you mean out_seg?" before reaching the hint. Unaccepted, the hint fires. What the class does take is the rest of the surface, which its per-region machinery consumes directly: `seg_unit` chooses what a region is, `bg_split` gives it more than one region of air, `seg_unit_prob` which regions contribute. Its defaults are `seg_unit="label"`, `seg_unit_prob=1.0` -- the region set and participation it has always had -- and the keep draw spends nothing at 1.0, so the default is bit-compatible. Verified: the unit decomposition is tensor-identical to the `seg_region_masks` set it replaces, including when a patch has no background at all. The four TAs that held the segmentation but could not be confined to it now can be. Each composites after its own stats restore, which is where `_select_and_check` already sat in the other twenty-four: * RandomSynthSegGPU, RandomPaletteGPU, PaletteSynthesisGPU with `normalize=True`, reviving a parameter `_apply_region_mode` had carried with no caller. Their output is a [0, 1] synthesis and the patch it lands on is z-scored, so a raw composite would leave an intensity step on the mask boundary -- a feature perfectly correlated with the target. * RandomDomainTransferGPU after the z-restore rather than after its LUT-space blend, which is the point at which both sides share a scale. RandomPaletteGPU and PaletteSynthesisGPU also gain the non-finite guard the other transforms have had; they used to write the synthesis unconditionally. Two deliberate stream changes, both once and both documented in gpu/regions.py: the mode draw is per sample rather than batch-wide (it is the last decision in the package that was not), and it spends nothing when both probabilities are zero. `RandomInverseGPU`'s per-sample loop existed only to get a per-sample mode and is gone with it; its non-finite guard is now per channel, which only matters for an input that was already broken, since the transform is `max - x`. Also removed: `_apply_region_mode`'s dead `normalize` branch is now live, its dead `mix_in_out` branch is gone, `_choose_region_mode`'s unused `seg_mask` parameter is gone, out-of-range probabilities are rejected rather than silently clamped, and `ZscoreNormalizationGPU` no longer has to assign `mix_in_out` by hand outside its signature. Suite is at the documented 11-failure baseline in seg11 (9 from nnunetv2 2.4.2, 2 from the stale random-normalization config); `template --check` and `matrix --check` both pass. Every migrated config gets a new content hash, so experiment directory names change -- the old files are in git history. Co-Authored-By: Claude Opus 5 --- smauglab/configs/all_augmentations.json | 267 +++++++-- smauglab/configs/transform_params.json | 2 +- smauglab/configs/transform_params_gpu.json | 46 +- smauglab/configs/transform_params_hybrid.json | 50 +- .../configs/transform_params_hybrid_TAGE.json | 50 +- ...rams_one-sequence-to-segment-them-all.json | 46 +- .../transform_params_paper-synthseg.json | 46 +- smauglab/registry.py | 10 + smauglab/transforms/cpu/fromSeg.py | 22 +- smauglab/transforms/gpu/base.py | 25 + smauglab/transforms/gpu/contrast.py | 565 ++++++++---------- smauglab/transforms/gpu/domain_transfer.py | 30 +- smauglab/transforms/gpu/fromSeg.py | 132 +++- smauglab/transforms/gpu/palette/transform.py | 25 +- smauglab/transforms/gpu/regions.py | 77 ++- smauglab/transforms/synthseg/transforms.py | 41 +- unit_tests/test_redistribute_autocast.py | 4 +- ...in_seg.py => test_redistribute_confine.py} | 28 +- unit_tests/test_region_and_stats.py | 90 ++- unit_tests/test_seg_layout.py | 4 +- 20 files changed, 1040 insertions(+), 520 deletions(-) rename unit_tests/{test_redistribute_in_seg.py => test_redistribute_confine.py} (71%) diff --git a/smauglab/configs/all_augmentations.json b/smauglab/configs/all_augmentations.json index 7301826..9da857e 100644 --- a/smauglab/configs/all_augmentations.json +++ b/smauglab/configs/all_augmentations.json @@ -41,6 +41,11 @@ "apply_resolution": true, "apply_to_channel": null, "atlas_res": 1.0, + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "bias_field_std": 0.7, "bias_scale": 0.025, "blur_range": 1.03, @@ -61,6 +66,7 @@ "gamma_std": 0.5, "generation_classes": null, "generation_labels": null, + "in_seg": 0.0, "keepdim": true, "max_res_aniso": 8.0, "max_res_iso": 4.0, @@ -74,6 +80,7 @@ "nonlin_scale": 0.04, "nonlin_std": 4.0, "normalise": true, + "out_seg": 0.0, "output_labels": null, "p": 0.5, "p_batch": 1.0, @@ -84,6 +91,9 @@ "rotation_bounds": 15.0, "same_on_batch": false, "scaling_bounds": 0.2, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "shearing_bounds": 0.012, "svf_integration_steps": 7, "thickness": null, @@ -94,6 +104,11 @@ 0.5, 2.0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "blur_sigmas": [ 0.0, 0.0, @@ -110,6 +125,7 @@ 6 ], "dark_threshold": 0.01, + "in_seg": 0.0, "keepdim": false, "label_classes": null, "label_remap_prob": 0.5, @@ -120,6 +136,7 @@ ], "mix_prob": 0.0, "n_kmeans_subsample": 10000, + "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "s_choices": [ @@ -134,6 +151,9 @@ 10 ], "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "skip_parcellation_prob": 0.1, "skip_sub_parc_prob": 0.4 }, @@ -142,6 +162,11 @@ 0.5, 2.0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "blur_sigmas_post": [ 0.0, 0.0, @@ -159,6 +184,7 @@ 0.8 ], "dark_threshold": 0.01, + "in_seg": 0.0, "initial_partitioner": null, "keepdim": false, "mix_alpha": [ @@ -166,20 +192,30 @@ 1.0 ], "mix_prob": 0.0, + "out_seg": 0.0, "overlay": null, "p": 1.0, "p_batch": 1.0, "refinement_partitioners": null, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomDomainTransferGPU": { "any_source": false, "apply_to_channel": null, "bank_path": null, + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "bias_field_std": 0.0, "bias_scale": 0.03, "blend_concentration": 1.0, "blend_targets": 1, + "in_seg": 0.0, "include_self": true, "keepdim": true, "mix_alpha": [ @@ -187,12 +223,16 @@ 1.0 ], "mix_prob": 0.0, + "out_seg": 0.0, "p": 0.2, "p_batch": 1.0, "p_class_mix": 0.0, "p_spatial_mix": 0.0, "pct": 1.0, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "sigma": 2.0, "source_label": null, "spatial_mix_gain": 3.0, @@ -204,47 +244,66 @@ "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomHistogramEqualizationGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomRedistributeSegGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", + "confine_prob": 0.2, "dilation_iterations_range": [ 1, 3 ], - "in_seg": 0.2, "keepdim": true, "mix_alpha": [ 0.0, @@ -255,6 +314,9 @@ "p_batch": 1.0, "retain_stats": true, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "label", + "seg_unit_prob": 1.0, "std_noise_range": [ 0.1, 0.3 @@ -265,56 +327,77 @@ "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": true, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomLaplaceGPU": { "absolute": false, "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomUnsharpMaskGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "sigma": 1.0, "unsharp_amount": 1.5 }, @@ -322,6 +405,11 @@ "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "kernel_sizes": [ @@ -334,217 +422,303 @@ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomClampGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "max_clamp_amount": 0.0, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomGaussianNoiseGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mean": 0.0, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "std": 1.0 }, "RandomGaussianBlurGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5, "sigma": 1.0 }, "RandomBrightnessGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "brightness_range": [ 0.5, 1.5 ], "in_seg": 0.0, "keepdim": true, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomGammaGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "gamma_range": [ 0.7, 1.5 ], "in_seg": 0.0, "keepdim": false, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomInvGammaGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "gamma_range": [ 0.7, 1.5 ], "in_seg": 0.0, "keepdim": false, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomContrastGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "contrast_range": [ 0.75, 1.25 ], "in_seg": 0.0, "keepdim": true, - "mix_in_out": false, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomLog1pGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomSqrtGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomSinGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomExpGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomSigmoidGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "mix_alpha": [ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "RandomLowResTransformGPU": { "keepdim": true, @@ -588,6 +762,11 @@ "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "coefficients": 0.5, "in_seg": 0.0, "invert": false, @@ -596,24 +775,34 @@ 0.0, 1.0 ], - "mix_in_out": false, "mix_prob": 0.0, "order": 3, "out_seg": 0.0, "p": 1.0, "p_batch": 1.0, "retain_stats": false, - "same_on_batch": false + "same_on_batch": false, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 }, "ZscoreNormalizationGPU": { "apply_to_channel": [ 0 ], + "bg_em_clusters": [ + 3, + 10 + ], + "bg_split": "none", "in_seg": 0.0, "keepdim": true, "out_seg": 0.0, "p": 1.0, - "p_batch": 1.0 + "p_batch": 1.0, + "seg_feather": 0.0, + "seg_unit": "whole", + "seg_unit_prob": 0.5 } }, "CPU": { @@ -653,7 +842,7 @@ }, "RedistributeTransform": { "classes": null, - "in_seg": 0.2, + "confine_prob": 0.2, "p": 1.0, "retain_stats": false }, diff --git a/smauglab/configs/transform_params.json b/smauglab/configs/transform_params.json index bbaaeaf..d1b5394 100644 --- a/smauglab/configs/transform_params.json +++ b/smauglab/configs/transform_params.json @@ -30,7 +30,7 @@ "retain_stats": true }, "RedistributeTransform": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": true, "p": 0.5 }, diff --git a/smauglab/configs/transform_params_gpu.json b/smauglab/configs/transform_params_gpu.json index 05e170b..493ea1b 100644 --- a/smauglab/configs/transform_params_gpu.json +++ b/smauglab/configs/transform_params_gpu.json @@ -74,7 +74,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.0, "p": 0.25 }, @@ -82,12 +83,13 @@ "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.6, "p": 0.25 }, "RandomRedistributeSegGPU": { - "in_seg": 0.25, + "confine_prob": 0.25, "retain_stats": true, "p": 0.4 }, @@ -97,7 +99,8 @@ "mix_prob": 0.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.25 }, "RandomUnsharpMaskGPU": { @@ -105,7 +108,8 @@ "unsharp_amount": 1.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.25 }, @@ -118,7 +122,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.2 }, @@ -127,7 +132,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.0 }, "RandomGaussianNoiseGPU": { @@ -135,14 +141,16 @@ "std": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomBrightnessGPU": { @@ -152,7 +160,8 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomGammaGPU": { @@ -163,7 +172,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -174,7 +184,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -185,42 +196,38 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "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": { @@ -256,7 +263,8 @@ "coefficients": 0.2, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "ZscoreNormalizationGPU": { diff --git a/smauglab/configs/transform_params_hybrid.json b/smauglab/configs/transform_params_hybrid.json index 14175bf..db87125 100644 --- a/smauglab/configs/transform_params_hybrid.json +++ b/smauglab/configs/transform_params_hybrid.json @@ -35,18 +35,20 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomHistogramEqualizationGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomRedistributeSegGPU": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": false, "p": 0.3 }, @@ -55,7 +57,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomUnsharpMaskGPU": { @@ -63,7 +66,8 @@ "unsharp_amount": 1, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomRandConvGPU": { @@ -76,7 +80,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.2, "p": 0.1 }, @@ -85,7 +90,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.4 }, "RandomGaussianNoiseGPU": { @@ -93,14 +99,14 @@ "std": 0.1, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": false, "p": 0.1 }, "RandomBrightnessGPU": { @@ -110,7 +116,6 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": false, "p": 0.15 }, "RandomGammaGPU": { @@ -121,7 +126,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -132,7 +138,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -143,42 +150,46 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": false, "p": 0.15 }, "RandomLog1pGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSqrtGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSinGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomExpGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSigmoidGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomLowResTransformGPU": { @@ -202,7 +213,6 @@ "coefficients": 0.5, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": false, "p": 0.2 }, "ZscoreNormalizationGPU": { @@ -242,7 +252,7 @@ "p": 0.1 }, "RedistributeTransform": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": false, "p": 0.5 }, diff --git a/smauglab/configs/transform_params_hybrid_TAGE.json b/smauglab/configs/transform_params_hybrid_TAGE.json index 2e2e052..5d62a9b 100644 --- a/smauglab/configs/transform_params_hybrid_TAGE.json +++ b/smauglab/configs/transform_params_hybrid_TAGE.json @@ -36,18 +36,20 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomHistogramEqualizationGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomRedistributeSegGPU": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": false, "p": 0.3 }, @@ -56,7 +58,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomUnsharpMaskGPU": { @@ -64,7 +67,8 @@ "unsharp_amount": 1, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomRandConvGPU": { @@ -77,7 +81,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.2, "p": 0.1 }, @@ -86,7 +91,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.4 }, "RandomGaussianNoiseGPU": { @@ -94,14 +100,14 @@ "std": 0.1, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": false, "p": 0.1 }, "RandomBrightnessGPU": { @@ -111,7 +117,6 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": false, "p": 0.15 }, "RandomGammaGPU": { @@ -122,7 +127,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -133,7 +139,8 @@ "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -144,42 +151,46 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": false, "p": 0.15 }, "RandomLog1pGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSqrtGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSinGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomExpGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomSigmoidGPU": { "retain_stats": false, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.05 }, "RandomLowResTransformGPU": { @@ -203,7 +214,6 @@ "coefficients": 0.5, "in_seg": 0.5, "out_seg": 0.5, - "mix_in_out": false, "p": 0.2 }, "ZscoreNormalizationGPU": { @@ -243,7 +253,7 @@ "p": 0.1 }, "RedistributeTransform": { - "in_seg": 0.2, + "confine_prob": 0.2, "retain_stats": false, "p": 0.5 }, diff --git a/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json b/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json index 6d124c9..6ed93c6 100755 --- a/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json +++ b/smauglab/configs/transform_params_one-sequence-to-segment-them-all.json @@ -34,7 +34,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.0, "p": 0.25 }, @@ -42,12 +43,13 @@ "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.6, "p": 0.25 }, "RandomRedistributeSegGPU": { - "in_seg": 0.25, + "confine_prob": 0.25, "retain_stats": true, "p": 0.4 }, @@ -57,7 +59,8 @@ "mix_prob": 0.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.25 }, "RandomUnsharpMaskGPU": { @@ -65,7 +68,8 @@ "unsharp_amount": 1.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.25 }, @@ -78,7 +82,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.2 }, @@ -87,7 +92,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.0 }, "RandomGaussianNoiseGPU": { @@ -95,14 +101,16 @@ "std": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomBrightnessGPU": { @@ -112,7 +120,8 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomGammaGPU": { @@ -123,7 +132,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -134,7 +144,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -145,42 +156,38 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "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": { @@ -204,7 +211,8 @@ "coefficients": 0.2, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "ZscoreNormalizationGPU": { diff --git a/smauglab/configs/transform_params_paper-synthseg.json b/smauglab/configs/transform_params_paper-synthseg.json index dc1b461..f571427 100644 --- a/smauglab/configs/transform_params_paper-synthseg.json +++ b/smauglab/configs/transform_params_paper-synthseg.json @@ -76,7 +76,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.0, "p": 0.0 }, @@ -84,12 +85,13 @@ "retain_stats": true, "in_seg": 0, "out_seg": 0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.6, "p": 0.0 }, "RandomRedistributeSegGPU": { - "in_seg": 0.25, + "confine_prob": 0.25, "retain_stats": true, "p": 0.0 }, @@ -99,7 +101,8 @@ "mix_prob": 0.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.0 }, "RandomUnsharpMaskGPU": { @@ -107,7 +110,8 @@ "unsharp_amount": 1.5, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.0 }, @@ -120,7 +124,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "mix_prob": 0.5, "p": 0.0 }, @@ -129,7 +134,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.0 }, "RandomGaussianNoiseGPU": { @@ -137,14 +143,16 @@ "std": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomGaussianBlurGPU": { "sigma": 1.0, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.2 }, "RandomBrightnessGPU": { @@ -154,7 +162,8 @@ ], "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "RandomGammaGPU": { @@ -165,7 +174,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.3 }, "RandomInvGammaGPU": { @@ -176,7 +186,8 @@ "retain_stats": true, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.1 }, "RandomContrastGPU": { @@ -187,42 +198,38 @@ "retain_stats": false, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "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": { @@ -246,7 +253,8 @@ "coefficients": 0.2, "in_seg": 0.0, "out_seg": 0.0, - "mix_in_out": true, + "seg_unit": "label", + "seg_unit_prob": 0.5, "p": 0.15 }, "ZscoreNormalizationGPU": { diff --git a/smauglab/registry.py b/smauglab/registry.py index 36805bc..4dd7351 100644 --- a/smauglab/registry.py +++ b/smauglab/registry.py @@ -177,6 +177,16 @@ def __init__(self, source: str, problems: list[str]) -> None: "kernel_type": "the class name now carries the kernel (e.g. RandomScharrGPU)", "func": "the class name now carries the function (e.g. RandomLog1pGPU)", "one_dim": "use RandomAcqTransformGPU (one_dim) or RandomLowResTransformGPU", + "mix_in_out": "use seg_unit='label' with seg_unit_prob", + # Only reachable on RandomRedistributeSegGPU, the one transform that does not + # accept `in_seg`. That is deliberate: it used to accept it with a *different* + # meaning (confine each region's Gaussian to its own voxels), twenty-eight shipped + # configs set it, and if the class also took the standard `in_seg` they would all + # have kept validating while silently changing what they do. Keeping the name + # unaccepted is what turns that into this message. Note difflib would otherwise + # answer "did you mean out_seg?" -- it scores 0.615, above the 0.6 cutoff -- which + # is exactly why `out_seg` is not accepted there either. + "in_seg": "on RandomRedistributeSegGPU it is now confine_prob; elsewhere in_seg is unchanged", } ) diff --git a/smauglab/transforms/cpu/fromSeg.py b/smauglab/transforms/cpu/fromSeg.py index 6f467a2..11e64db 100644 --- a/smauglab/transforms/cpu/fromSeg.py +++ b/smauglab/transforms/cpu/fromSeg.py @@ -19,14 +19,18 @@ class RedistributeTransform(BasicTransform): Based on https://github.com/neuropoly/totalspineseg/blob/main/totalspineseg/utils/augment.py """ - def __init__(self, classes=None, in_seg=0.2, retain_stats=False): + # `confine_prob` was `in_seg`, renamed with its GPU twin: it is the probability that + # each label's Gaussian is confined to that label's own voxels, not the probability of + # applying the transform inside the segmentation, which is what `in_seg` means + # everywhere else. Both halves of `AugId.REDISTRIBUTE_SEG` use the one name. + def __init__(self, classes=None, confine_prob=0.2, retain_stats=False): super().__init__() self.classes = classes - self.in_seg = in_seg + self.confine_prob = confine_prob self.retain_stats = retain_stats def get_parameters(self, **data_dict) -> dict: - return {"classes": self.classes, "in_seg": self.in_seg, "retain_stats": self.retain_stats} + return {"classes": self.classes, "confine_prob": self.confine_prob, "retain_stats": self.retain_stats} def apply(self, data_dict: dict, **params) -> dict: if data_dict.get("image") is not None and data_dict.get("segmentation") is not None: @@ -36,12 +40,12 @@ def apply(self, data_dict: dict, **params) -> dict: def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> tuple[torch.Tensor, torch.Tensor]: for c in range(1): # Works on the first channel only img[c], seg[c] = aug_redistribute_seg( - img[c], seg[c], classes=params["classes"], in_seg=params["in_seg"], retain_stats=params["retain_stats"] + img[c], seg[c], classes=params["classes"], confine_prob=params["confine_prob"], retain_stats=params["retain_stats"] ) return img, seg -def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False): +def aug_redistribute_seg(img, seg, classes=None, confine_prob=0.2, retain_stats=False): """ Augment the image by redistributing the values of the image within the regions defined by the segmentation. @@ -50,7 +54,7 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) """ device = img.device _seg = seg - in_seg_bool = 1 - torch.rand(1, device=device) <= in_seg + confine = 1 - torch.rand(1, device=device) <= confine_prob if classes: _seg = combine_classes(_seg, classes) @@ -112,7 +116,7 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) ) sign = 2 * torch.rand(1, device=device) - 1 - if in_seg_bool: + if confine: to_add[l_mask] += _normal_pdf(img[l_mask], l_mean, redist_std) * sign else: to_add += _normal_pdf(img, l_mean, redist_std) * sign @@ -120,8 +124,8 @@ def aug_redistribute_seg(img, seg, classes=None, in_seg=0.2, retain_stats=False) # The min-max form shifted zero: with `2 * (to_add - min) / range`, a voxel where # nothing was redistributed picks up `-2 * min / range`, non-zero whenever min < 0 -- # and it always is, since the per-label scale is `2 * rand - 1` in [-1, 1]. On the - # `in_seg` branch `to_add` is exactly zero outside the labels, so that constant became - # a DC offset over the whole patch: `in_seg` restricted where the redistribution was + # confined branch `to_add` is exactly zero outside the labels, so that constant became + # a DC offset over the whole patch: confining restricted where the redistribution was # *computed* but not where it landed. # # Scaling by the largest magnitude keeps the peak amplitude of 2 and maps zero to zero, diff --git a/smauglab/transforms/gpu/base.py b/smauglab/transforms/gpu/base.py index 62f78d3..d43e55a 100644 --- a/smauglab/transforms/gpu/base.py +++ b/smauglab/transforms/gpu/base.py @@ -33,6 +33,13 @@ from torch.nn import Module from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, +) try: # private to kornia, and the only part of the faster parameter draw that needs it from kornia.augmentation.container.image import ImageSequentialBase, _get_new_batch_shape @@ -95,6 +102,24 @@ class ImageOnlyTransform(RigidAffineAugmentationBase3D): mix_prob: float = 0.0 mix_alpha: tuple[float, float] = DEFAULT_MIX_ALPHA + #: Which part of the segmentation the transform applies to, and how sharply it stops; + #: see `gpu/regions.py`. Declared here for the same reason as the two above, and with + #: the values that mean "off": `in_seg = out_seg = 0` is "the whole patch", and then + #: none of the other five is read. `seg_feather` is the normalised pair + #: `configure_regions` assigns, not the scalar-or-range a config may write. + #: + #: These used to be assigned only in the constructors, with no class-level default, so + #: `ZscoreNormalizationGPU` had to set `mix_in_out` by hand outside its signature to + #: keep the helper's contract -- the registry derives a config's accepted keys from + #: `__init__`, so declaring it there would have invented a setting. + in_seg: float = 0.0 + out_seg: float = 0.0 + seg_unit: str = DEFAULT_SEG_UNIT + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB + bg_split: str = DEFAULT_BG_SPLIT + bg_em_clusters: tuple[int, int] = DEFAULT_BG_EM_CLUSTERS + seg_feather: tuple[float, float] = (DEFAULT_SEG_FEATHER, DEFAULT_SEG_FEATHER) + def compute_transformation(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any]) -> Tensor: # These transforms never move anything, so the matrix is always the # identity -- shared and cached rather than rebuilt per module. See diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index e2f9c56..23321eb 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -1,7 +1,7 @@ import functools import math from collections.abc import Callable, Sequence -from typing import Any, Protocol, Union +from typing import Any, Union import torch import torchvision.transforms._functional_tensor as F_t @@ -11,33 +11,25 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA, MixAlphaSpec, blend_one, configure_mix, mix_with_original + +# Re-exported, not used here: `_foreground` moved to gpu/regions.py with the rest of the +# region machinery, and the tests that pin its three segmentation layouts import it from +# this module, where it used to live. +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _foreground, # noqa: F401 + _select_and_check, + configure_regions, +) from smauglab.transforms.kernels import depthwise_conv3d, gaussian_kernel3d, laplace_kernel, scharr_kernels, stacked_scharr_kernels from smauglab.transforms.rng import shared_choice -def _choose_region_mode( - p_in: float, - p_out: float, - seg_mask: torch.Tensor | None, # noqa: ARG001 -) -> str: - """Sample where to apply the transform: 'in', 'out', or 'all'. - - - p_in, p_out are probabilities in [0,1]. - - If both probs are 0, or both fire at once, return 'all'. - - seg_mask is accepted but unused here; _apply_region_mode treats a None - mask as 'all' regardless of the mode chosen. - """ - p_in = float(max(0.0, min(1.0, p_in))) - p_out = float(max(0.0, min(1.0, p_out))) - in_bool = torch.rand(()) < p_in - out_bool = torch.rand(()) < p_out - if in_bool and not out_bool: - return "in" - if out_bool and not in_bool: - return "out" - return "all" - - def _channel_stats(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Per-sample mean and std of a `[N, ...spatial]` channel, each shaped `[N]`.""" reduce_dims = tuple(range(1, x.dim())) @@ -76,159 +68,81 @@ def _check_channel(transform: torch.nn.Module, channel: int, channels: int) -> N ) -class _RegionSelecting(Protocol): - """What `_select_and_check` needs off the transform it is handed. +def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torch.Tensor: + """ + Based on https://github.com/pytorch/vision/blob/e3b5d3a8bf5e8636462fd8bce9897bccc690b2a0/torchvision/transforms/_functional_tensor.py#L746 + """ + if not (isinstance(img, torch.Tensor)): + raise TypeError(f"img should be Tensor. Got {type(img)}") + + if dim == 2: + kernel = kernel.expand(img.shape[-(1 + dim)], 1, kernel.shape[0], kernel.shape[1]) + padding = [kernel.shape[2] // 2, kernel.shape[2] // 2, kernel.shape[3] // 2, kernel.shape[3] // 2] + elif dim == 3: + kernel = kernel.expand(img.shape[-(1 + dim)], 1, kernel.shape[0], kernel.shape[1], kernel.shape[2]) + padding = [ + kernel.shape[2] // 2, + kernel.shape[2] // 2, + kernel.shape[3] // 2, + kernel.shape[3] // 2, + kernel.shape[4] // 2, + kernel.shape[4] // 2, + ] + else: + raise ValueError(f"Only 2D and 3D convolution are supported. Got {dim}D.") - A Protocol rather than `ImageOnlyTransform`, because these are plain attributes on - an nn.Module: reading them through the base class resolves via `Module.__getattr__`, - which is typed as returning `Tensor | Module`. Naming them here is both what makes - that type-check and a statement of the helper's actual requirement. + img, need_cast, need_squeeze, out_dtype = F_t._cast_squeeze_in(img, [kernel.dtype]) - All three are attributes rather than constructor parameters as far as this - Protocol is concerned; the registry derives a transform's config surface from its - __init__ signature, so declaring one here does not make it settable in a config. - """ + # padding = (left, right, top, bottom) + img = F.pad(img, padding, mode="reflect") + if dim == 2: # noqa: SIM108 -- the 2d/3d split reads better spelled out than as a ternary + img = F.conv2d(img, kernel, groups=img.shape[-(1 + dim)]) + else: # dim == 3 + # Via `depthwise_conv3d`, not `F.conv3d`: every row of `kernel` is the same + # filter (it came from an `expand` above), so one row is handed over and the + # helper broadcasts it back across the planes, which is also what lets it make + # `groups > 1` from a single-plane call. The bucketed pipelines hand over one + # sample at a time, precisely the case cuDNN serves 35x slower than depthwise. + leading = img.shape[:-3] + planes = int(torch.tensor(leading).prod()) if leading else 1 + out = depthwise_conv3d(img.reshape(planes, 1, *img.shape[-3:]), kernel[:1, :1]) + img = out.reshape(*leading, *out.shape[-3:]) - in_seg: float - out_seg: float - mix_in_out: bool + img = F_t._cast_squeeze_out(img, need_cast, need_squeeze, out_dtype) + return img -def _select_and_check( - transform: _RegionSelecting, - orig: torch.Tensor, - x: torch.Tensor, - seg_mask: torch.Tensor | None, - note: str = "", -) -> torch.Tensor | None: - """Apply region selection, then reject the result if it went non-finite. +def apply_convolution_per_sample(channel_data: torch.Tensor, kernels: Sequence[Tensor]) -> torch.Tensor: + """Convolve `[N, D, H, W]` with one 3-D kernel per sample, in a single conv. - Returns None when the channel should be left as it was -- the "Final safety" check - that closed all nine of these loops identically. + The alternative -- and what every caller here used to do -- is a Python loop + calling `apply_convolution` on `channel_data[b : b + 1]`. That is a + single-channel `F.conv3d`, which cuDNN serves from its implicit-GEMM path: 4.3 ms + per sample for a 3x3x3 kernel over a 128^3 patch on an A40, against 0.12 ms for + the whole batch here. See `kernels.depthwise_conv3d` for why the shape matters. - Takes the transform rather than its three region attributes: every call site read - exactly `self.in_seg`, `self.out_seg` and `self.mix_in_out`, and spelling them out - made the call longer than the code it replaces. - """ - if seg_mask is not None: - region_mode = _choose_region_mode(transform.in_seg, transform.out_seg, seg_mask) - x = _apply_region_mode(orig, x, seg_mask, region_mode, mix_in_out=transform.mix_in_out) - # One `isfinite(...).all()`, not `isnan().any() or isinf().any()`: each of those - # reads a reduction back to the host, and the `or` forces the second on the common - # path. Runs for every channel of every region-selecting transform. - if not bool(torch.isfinite(x).all()): - print(f"Warning nan: {type(transform).__name__}{note}", flush=True) - return None - return x - - -def _foreground(mask: torch.Tensor, dim: int) -> torch.Tensor: - """Which voxels the segmentation covers, reducing over the class axis. - - This used to be `torch.argmax(mask, dim) > 0`, which is only "is anything labelled - here" if class 0 is background -- and it is not: - - * For a single-channel `[B, 1, D, H, W]` mask (an ordinary nnU-Net target, and what - the tests build) `argmax` over a length-1 axis is always 0, so the result was - **all False**. `in_seg` then applied the transform nowhere and `out_seg` applied - it everywhere: the two knobs did nothing and the opposite of nothing. - * For a one-hot mask, this repository's convention (`collapse_onehot_to_index` in - gpu/fromSeg.py) is that channel `c` encodes label `c + 1` with background - implicit, so `argmax == 0` is a real foreground class and was being dropped. - - `amax > 0` asks the question that was meant, and matches what - `collapse_onehot_to_index` already does with `seg_raw.any(dim=1)`. - - That still leaves one layout it cannot read off the values alone: a one-hot - that *includes* a background channel, which is what `seg_region_masks` emits - for a single-channel label map and what `unit_tests/test_seg_layout.py` - builds. Every voxel then has some channel set, `amax > 0` is True everywhere, - and `in_seg` applies the transform to the whole patch while `out_seg` applies - it nowhere -- the exact failure this function was written to fix, in the one - layout it was not checked against. - - A background-inclusive one-hot is recognisable: it is multi-channel and leaves - no voxel unset, because the background channel covers whatever the foreground - channels do not. A background-implicit one-hot always has all-zero background - voxels unless the patch is labelled edge to edge, which a real segmentation - patch is not. So full coverage plus more than one channel means channel 0 is - background, and the reduction skips it. - """ - covered = mask.amax(dim=dim) > 0 - if mask.shape[dim] == 1: - return covered - # `torch.where`, not `if bool(...)`: reading that test back to Python stalls the host - # on everything queued behind it. Both branches are a reduction over the class axis, - # so computing the unused one costs a fraction of what the stall does. - without_background = mask.narrow(dim, 1, mask.shape[dim] - 1).amax(dim=dim) > 0 - return torch.where(covered.all(), without_background, covered) - - -def _apply_region_mode( - orig: torch.Tensor, - transformed: torch.Tensor, - seg_mask: torch.Tensor | None, - mode: str, - normalize: bool = False, - mix_in_out: bool = False, -) -> torch.Tensor: - """Blend transformed with orig based on region selection mode. - - - mode 'all': return transformed - - mode 'in': apply transform inside seg, keep orig outside - - mode 'out': apply transform outside seg, keep orig inside - - mix_in_out: if True, randomly apply transform to some of the segmentation, not all. + Kernels of different sizes are zero-padded up to the widest one. That is exact, + not an approximation: a zero tap contributes nothing, and `reflect` padding gives + the same value at a given offset however wide the pad is, so the extra taps read + real (if irrelevant) voxels and multiply them by zero. `RandConv` draws its kernel + size per sample, so without this the batch could not be convolved in one call. """ - if seg_mask is None or mode == "all": - return transformed - - # Rescale transformed based on min max orig - # Needed due to the important change in the image - if orig.dim() == 4: - if normalize: - orig_min = torch.amin(orig, dim=tuple(range(1, orig.dim())), keepdim=True) - orig_max = torch.amax(orig, dim=tuple(range(1, orig.dim())), keepdim=True) - transformed_min = torch.amin(transformed, dim=tuple(range(1, transformed.dim())), keepdim=True) - transformed_max = torch.amax(transformed, dim=tuple(range(1, transformed.dim())), keepdim=True) - transformed = (transformed - transformed_min) / (transformed_max - transformed_min + 1e-8) * (orig_max - orig_min) + orig_min - - # No `.clone()`: nothing below writes into `m`, and on a [B, C, D, H, W] mask - # that copy is the size of the batch. - m = seg_mask.to(transformed.dtype) - if mix_in_out: - # One keep/drop draw per (sample, class), in a single call rather than one - # `randint` per sample: the same draws in the same order, one kernel. - keep = torch.randint(0, 2, seg_mask.shape[:2], device=seg_mask.device, dtype=m.dtype) - m = m * keep.view(*seg_mask.shape[:2], *([1] * (m.dim() - 2))) - - m = _foreground(m, dim=1) - m = m.to(transformed.dtype) - if mode == "out": - m = 1.0 - m - - elif orig.dim() == 3: - if normalize: - orig_min = torch.amin(orig) - orig_max = torch.amax(orig) - transformed_min = torch.amin(transformed) - transformed_max = torch.amax(transformed) - transformed = (transformed - transformed_min) / (transformed_max - transformed_min + 1e-8) * (orig_max - orig_min) + orig_min - - m = seg_mask.to(transformed.dtype) - if mix_in_out: - # Create a tensor with random one and zero - o = torch.randint(0, 2, (seg_mask.shape[0],), device=seg_mask.device, dtype=m.dtype) - m = m * o.view(-1, 1, 1, 1) # Broadcasting o to match the dimensions of m - m = _foreground(m, dim=0) - m = m.to(transformed.dtype) - if mode == "out": - m = 1.0 - m - - else: - raise ValueError(f"Only 4D and 3D images are supported. Got {orig.dim()}D.") + sizes = [int(k.shape[-1]) for k in kernels] + widest = max(sizes) + stacked = [] + for kernel, size in zip(kernels, sizes): + if size == widest: + stacked.append(kernel) + else: + lo = (widest - size) // 2 + padded = torch.zeros((widest, widest, widest), device=kernel.device, dtype=kernel.dtype) + padded[lo : lo + size, lo : lo + size, lo : lo + size] = kernel + stacked.append(padded) + weight = torch.stack(stacked).unsqueeze(1) # [N, 1, k, k, k] - return m * transformed + (1.0 - m) * orig + pad = widest // 2 + padded_img = F.pad(channel_data.unsqueeze(1), [pad] * 6, mode="reflect") + return depthwise_conv3d(padded_img, weight).squeeze(1) ## Convolution transform @@ -262,7 +176,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, # Kernel-specific, and explicit rather than read from **kwargs, where # `inspect.signature` could not see them and a config typo silently took the # default. Defaults are the historical kwargs.get() ones, so behaviour is unchanged. @@ -285,9 +203,7 @@ def __init__( self.absolute = absolute self.sigma = sigma self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) self.unsharp_amount = unsharp_amount self.kernel_sizes = kernel_sizes configure_mix(self, mix_prob, mix_alpha) @@ -443,7 +359,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -458,7 +378,11 @@ def __init__( retain_stats=retain_stats, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -483,7 +407,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -498,7 +426,11 @@ def __init__( same_on_batch=same_on_batch, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -521,7 +453,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -536,7 +472,11 @@ def __init__( retain_stats=retain_stats, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, mix_prob=mix_prob, mix_alpha=mix_alpha, p=p, @@ -564,7 +504,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -580,7 +524,11 @@ def __init__( retain_stats=retain_stats, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -605,7 +553,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -620,90 +572,17 @@ def __init__( retain_stats=retain_stats, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, ) -def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torch.Tensor: - """ - Based on https://github.com/pytorch/vision/blob/e3b5d3a8bf5e8636462fd8bce9897bccc690b2a0/torchvision/transforms/_functional_tensor.py#L746 - """ - if not (isinstance(img, torch.Tensor)): - raise TypeError(f"img should be Tensor. Got {type(img)}") - - if dim == 2: - kernel = kernel.expand(img.shape[-(1 + dim)], 1, kernel.shape[0], kernel.shape[1]) - padding = [kernel.shape[2] // 2, kernel.shape[2] // 2, kernel.shape[3] // 2, kernel.shape[3] // 2] - elif dim == 3: - kernel = kernel.expand(img.shape[-(1 + dim)], 1, kernel.shape[0], kernel.shape[1], kernel.shape[2]) - padding = [ - kernel.shape[2] // 2, - kernel.shape[2] // 2, - kernel.shape[3] // 2, - kernel.shape[3] // 2, - kernel.shape[4] // 2, - kernel.shape[4] // 2, - ] - else: - raise ValueError(f"Only 2D and 3D convolution are supported. Got {dim}D.") - - img, need_cast, need_squeeze, out_dtype = F_t._cast_squeeze_in(img, [kernel.dtype]) - - # padding = (left, right, top, bottom) - img = F.pad(img, padding, mode="reflect") - if dim == 2: # noqa: SIM108 -- the 2d/3d split reads better spelled out than as a ternary - img = F.conv2d(img, kernel, groups=img.shape[-(1 + dim)]) - else: # dim == 3 - # Via `depthwise_conv3d`, not `F.conv3d`: every row of `kernel` is the same - # filter (it came from an `expand` above), so one row is handed over and the - # helper broadcasts it back across the planes, which is also what lets it make - # `groups > 1` from a single-plane call. The bucketed pipelines hand over one - # sample at a time, precisely the case cuDNN serves 35x slower than depthwise. - leading = img.shape[:-3] - planes = int(torch.tensor(leading).prod()) if leading else 1 - out = depthwise_conv3d(img.reshape(planes, 1, *img.shape[-3:]), kernel[:1, :1]) - img = out.reshape(*leading, *out.shape[-3:]) - - img = F_t._cast_squeeze_out(img, need_cast, need_squeeze, out_dtype) - return img - - -def apply_convolution_per_sample(channel_data: torch.Tensor, kernels: Sequence[Tensor]) -> torch.Tensor: - """Convolve `[N, D, H, W]` with one 3-D kernel per sample, in a single conv. - - The alternative -- and what every caller here used to do -- is a Python loop - calling `apply_convolution` on `channel_data[b : b + 1]`. That is a - single-channel `F.conv3d`, which cuDNN serves from its implicit-GEMM path: 4.3 ms - per sample for a 3x3x3 kernel over a 128^3 patch on an A40, against 0.12 ms for - the whole batch here. See `kernels.depthwise_conv3d` for why the shape matters. - - Kernels of different sizes are zero-padded up to the widest one. That is exact, - not an approximation: a zero tap contributes nothing, and `reflect` padding gives - the same value at a given offset however wide the pad is, so the extra taps read - real (if irrelevant) voxels and multiply them by zero. `RandConv` draws its kernel - size per sample, so without this the batch could not be convolved in one call. - """ - sizes = [int(k.shape[-1]) for k in kernels] - widest = max(sizes) - stacked = [] - for kernel, size in zip(kernels, sizes): - if size == widest: - stacked.append(kernel) - else: - lo = (widest - size) // 2 - padded = torch.zeros((widest, widest, widest), device=kernel.device, dtype=kernel.dtype) - padded[lo : lo + size, lo : lo + size, lo : lo + size] = kernel - stacked.append(padded) - weight = torch.stack(stacked).unsqueeze(1) # [N, 1, k, k, k] - - pad = widest // 2 - padded_img = F.pad(channel_data.unsqueeze(1), [pad] * 6, mode="reflect") - return depthwise_conv3d(padded_img, weight).squeeze(1) - - ## Noise transform @register( aug_id=AugId.GAUSSIAN_NOISE, @@ -733,7 +612,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -742,9 +625,7 @@ def __init__( self.apply_to_channel = apply_to_channel self.mean = mean self.std = std - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @torch.no_grad() # disable gradients for efficiency def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -801,7 +682,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -809,9 +694,7 @@ def __init__( super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) self.brightness_range = brightness_range self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @torch.no_grad() # disable gradients for efficiency def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -870,7 +753,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = False, @@ -880,9 +767,7 @@ def __init__( self.invert_image = invert_image self.retain_stats = retain_stats self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @torch.no_grad() # disable gradients for efficiency def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -963,7 +848,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = False, @@ -976,7 +865,11 @@ def __init__( same_on_batch=same_on_batch, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -999,7 +892,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = False, @@ -1012,7 +909,11 @@ def __init__( same_on_batch=same_on_batch, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, p=p, p_batch=p_batch, keepdim=keepdim, @@ -1049,7 +950,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -1058,9 +963,7 @@ def __init__( self.contrast_range = contrast_range self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @torch.no_grad() # disable gradients for efficiency def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -1129,7 +1032,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -1140,9 +1047,7 @@ def __init__( self.func = func self.retain_stats = retain_stats self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() # disable gradients for efficiency @@ -1211,7 +1116,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -1225,7 +1134,11 @@ def __init__( same_on_batch=same_on_batch, in_seg=in_seg, out_seg=out_seg, - mix_in_out=mix_in_out, + seg_unit=seg_unit, + seg_unit_prob=seg_unit_prob, + bg_split=bg_split, + bg_em_clusters=bg_em_clusters, + seg_feather=seg_feather, mix_prob=mix_prob, mix_alpha=mix_alpha, p=p, @@ -1316,7 +1229,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -1326,9 +1243,7 @@ def __init__( super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() # disable gradients for efficiency @@ -1340,9 +1255,15 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ for c in self.apply_to_channel: # The per-sample reductions -- max, and the two retain_stats passes -- are the # expensive part and are independent across samples, so they run once for the - # whole channel. Only the region selection stays in the loop: - # `_choose_region_mode` picks `in` / `out` / `all` per sample, and hoisting it - # would make that one decision for the batch. + # whole channel. The region selection used to be the one thing left in a + # per-sample loop, because the mode draw was batch-wide everywhere else and + # hoisting it here would have made one decision for the batch. It is per + # sample everywhere now (see `choose_region_modes`), so the loop is gone. + # + # What that costs: the non-finite guard was per sample and is now per channel, + # so one bad sample withholds the whole channel rather than just itself. This + # transform is `max - x`, which is finite whenever its input is, so the guard + # only fires on an input that was already broken. channel = input[:, c] reduce_dims = tuple(range(1, channel.dim())) keep_shape = (-1, *([1] * (channel.dim() - 1))) @@ -1368,14 +1289,10 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ new_std = inverted.std(dim=reduce_dims).view(keep_shape) inverted = (inverted - new_mean) / (new_std + eps) * orig_stds + orig_means - for i in range(input.shape[0]): - orig = channel[i] - x = inverted[i] - - checked = _select_and_check(self, orig, x, None if seg_mask is None else seg_mask[i]) - if checked is None: - continue - input[i, c] = checked + checked = _select_and_check(self, channel, inverted, seg_mask) + if checked is None: + continue + input[:, c] = checked return input @@ -1413,7 +1330,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -1423,9 +1344,7 @@ def __init__( super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) self.retain_stats = retain_stats self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() # disable gradients for efficiency @@ -1574,7 +1493,11 @@ def __init__( retain_stats: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, same_on_batch: bool = False, @@ -1595,9 +1518,7 @@ def __init__( self.apply_to_channel = apply_to_channel self.invert = invert self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) def _num_coeffs(self, dim: int) -> int: @@ -1795,7 +1716,11 @@ def __init__( same_on_batch: bool = False, in_seg: float = 0.0, out_seg: float = 0.0, - mix_in_out: bool = False, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, keepdim: bool = True, @@ -1804,9 +1729,7 @@ def __init__( self.max_clamp_amount = max_clamp_amount self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @torch.no_grad() # disable gradients for efficiency def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[str, Any], transform: Tensor | None = None) -> Tensor: @@ -1865,17 +1788,17 @@ def __init__( keepdim: bool = True, in_seg: float = 0.0, out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, p: float = 1.0, p_batch: float = 1.0, ) -> None: super().__init__(p=p, p_batch=p_batch, same_on_batch=False, keepdim=keepdim) self.apply_to_channel = apply_to_channel - self.in_seg = in_seg - self.out_seg = out_seg - # Not a constructor parameter: this transform has no mix_in_out knob and the - # registry derives a config's accepted keys from __init__, so adding one there - # would invent a setting. Set here so the region-selection helper's contract holds. - self.mix_in_out = False + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) @torch.no_grad() def apply_transform( @@ -1909,8 +1832,10 @@ def apply_transform( flat = channel.reshape(channel.shape[0], -1) spread = (flat.amax(dim=1) - flat.amin(dim=1)).view(-1, *([1] * (channel.dim() - 1))) channel = torch.where(spread == 0, channel, (channel - mean) / std) - # No mix_in_out here: z-scoring is applied whole, never to a random subset - # of the mask channels. + # `seg_unit` and `bg_split` are accepted here like everywhere else, but a + # z-score applied to a random subset of the units is a strange thing to ask + # for: the point of this transform is to put the whole channel on one scale. + # Left configurable rather than refused, since the surface is uniform. checked = _select_and_check(self, orig, channel, seg_mask) if checked is None: continue diff --git a/smauglab/transforms/gpu/domain_transfer.py b/smauglab/transforms/gpu/domain_transfer.py index 5f972aa..b00a0c0 100644 --- a/smauglab/transforms/gpu/domain_transfer.py +++ b/smauglab/transforms/gpu/domain_transfer.py @@ -36,6 +36,7 @@ import math import os +from collections.abc import Sequence from typing import Any import numpy as np @@ -48,6 +49,16 @@ from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from from smauglab.transforms.gpu.fromSeg import seg_region_masks from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA, MixAlphaSpec, blend_at, configure_mix, sample_mix +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _select_and_check, + configure_regions, +) from smauglab.transforms.kernels import gaussian_blur3d, random_bias_field3d # The transfer LUT bank is a multi-hundred-MB artefact built offline by @@ -146,6 +157,13 @@ def __init__( p_spatial_mix: float = 0.0, spatial_mix_scale: float = 0.03, spatial_mix_gain: float = 3.0, + in_seg: float = 0.0, + out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, same_on_batch: bool = False, @@ -223,6 +241,7 @@ def __init__( # Whole-transform blend back towards the input, distinct from `blend_targets` and # `p_spatial_mix` above: those interpolate between LUTs, which stays a monotone # remap of the source intensities, while this interpolates in image space. + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) def _sample_blended_luts(self, lut_bank: Tensor, n_seg_c: int, device) -> tuple[Tensor, bool]: @@ -375,5 +394,14 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # transferred distribution SHAPE while placing it back in the z-score scale the # network trains on, so the augmented patch matches the rest of the batch. acc = (acc - acc.mean()) / acc.std().clamp_min(1e-6) * sd + mu - out[b, ch] = acc + + # Region selection after the z-restore above, not after the LUT-space blend: + # at that point `acc` is still in the LUT's [0, 1] domain while `x` is in the + # input's, and a composite has to have both sides on one scale. By here they + # do, which is also why this is the one of the four synthesising transforms + # that needs no `normalize`. + checked = _select_and_check(self, x, acc, seg[b]) + if checked is None: + continue + out[b, ch] = checked return out diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index aac4140..7ffc3c1 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -10,6 +10,21 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA, MixAlphaSpec, blend_at, configure_mix, mix_with_original, sample_mix +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _select_and_check, + configure_regions, + draw_feather, + feather_alpha, + keep_units, + region_units, + unit_spec, +) from smauglab.transforms.kernels import gaussian_blur3d from smauglab.transforms.rng import shared_choice @@ -249,6 +264,24 @@ class RandomRedistributeSegGPU(ImageOnlyTransform): Mirrors the CPU `RedistributeTransform` behavior using GPU-friendly ops. Works with inputs shaped [N, C, H, W] or [N, C, D, H, W]. + `confine_prob` is the per-sample probability that each region's Gaussian perturbation + is restricted to that region's own voxels rather than added across the whole patch. + It was called `in_seg`, which it should never have been: everywhere else in the + package `in_seg` is the probability of applying the *whole transform* inside the + segmentation, and the two are different questions. The rename is not cosmetic -- + twenty-eight shipped configs set the old name, and had this class also accepted the + standard `in_seg` they would all have kept validating while silently changing meaning, + with no diagnostic possible because the key would be accepted. So this is the one + transform that does **not** take `in_seg` / `out_seg`, and a config still naming + `in_seg` here fails as an unknown parameter and is told what to write instead. + + What it does take is the rest of the region surface, because its work is already + per region: `seg_unit` chooses what a region *is* (one per label value, which is what + it has always used and therefore the default here, or one per connected component), + `bg_split="em"` splits the background into intensity clusters so there is more than + one region of air to redistribute between, `seg_unit_prob` is the chance each region + contributes, and `seg_feather` softens the edge of a confined contribution. + `retain_stats` defaults to True, and wants to stay that way. The whole method operates in a per-sample [0, 1] min-max space and adds a perturbation of up to 2.0 *in that space* -- twice the input's full dynamic range. With @@ -270,7 +303,7 @@ class RandomRedistributeSegGPU(ImageOnlyTransform): def __init__( self, - in_seg: float = 0.2, + confine_prob: float = 0.2, apply_to_channel: Sequence[int] = (0,), retain_stats: bool = True, same_on_batch: bool = False, @@ -279,15 +312,28 @@ def __init__( keepdim: bool = True, std_noise_range: Sequence[float] = (0.1, 0.3), dilation_iterations_range: Sequence[int] = (1, 3), + # `seg_unit` defaults to "label", not the base class's "whole": one region per + # distinct label value is the region set this transform has always used, so this + # is what keeps the default behaviour unchanged. `seg_unit_prob` defaults to 1.0 + # for the same reason -- every region contributing is what it has always done, and + # at 1.0 the keep draw spends nothing, so the default is bit-compatible. + seg_unit: str = "label", + seg_unit_prob: float = 1.0, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, ) -> None: super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) - self.in_seg = in_seg + self.confine_prob = confine_prob self.apply_to_channel = apply_to_channel self.retain_stats = retain_stats self.std_noise_range = std_noise_range self.dilation_iterations_range = dilation_iterations_range + # `in_seg` and `out_seg` are passed as zero and are deliberately not constructor + # parameters; see the class docstring. The rest of the region surface is read. + configure_regions(self, 0.0, 0.0, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() @@ -313,10 +359,34 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ N = input.shape[0] - # One region per label, whichever layout the caller used. Reading regions off - # the channel axis alone collapses a single-channel label map -- which is what - # nnU-Net's trainer passes -- to one foreground blob. - regions = seg_region_masks(seg) + if not len(self.apply_to_channel): + return input + + # The regions to redistribute between, from `gpu/regions.py` rather than from + # `seg_region_masks`. At this class's defaults -- `seg_unit="label"` -- the two + # produce the same set: one region per distinct label value, background included, + # numbered batch-wide. What the decomposition adds is the other two settings: + # `seg_unit="component"` splits a label into its connected pieces, and + # `bg_split="em"` splits the background into intensity clusters, so the + # redistribution has more than one region of air to work with. + # + # The reference image for `bg_split="em"` is the first selected channel. The + # decomposition is shared by every channel of the call -- and by every other + # transform in the pass -- so it cannot be per channel. + units = region_units(seg, input[:, self.apply_to_channel[0] : self.apply_to_channel[0] + 1], unit_spec(self)) + # Only the units some voxel actually lands in. `RegionUnits` always carries a + # background unit, occupied or not, and an unoccupied region is not harmless here: + # `_region_moments` clamps its count to 1 and hands back mean 0, so it would + # contribute a full-amplitude Gaussian at zero. Dropping it is also what makes the + # default region set identical to the `seg_region_masks` one this replaced, whose + # count was the number of label values *present*. + present = torch.bincount(units.ids.reshape(-1), minlength=units.n_units).nonzero().flatten() + region_index = present.view(1, present.numel(), *([1] * spatial_dims)) + regions = units.ids == region_index + # Which regions contribute at all. Spends nothing at the default `seg_unit_prob` + # of 1.0, where every region does. + keep_rows = keep_units(self, units, N, input.device)[:, present].tolist() + confine_sigma = draw_feather(self) # Apply per selected image channel and per batch sample for c in self.apply_to_channel: @@ -344,7 +414,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # sample's draw independent but not in the loop's interleaved order, so a # seeded run differs from the previous release's. has_foreground = (seg.reshape(N, -1) > 0).any(dim=1).tolist() - in_seg_draws = (torch.rand(N, device=input.device) <= self.in_seg).tolist() + confine_draws = (torch.rand(N, device=input.device) <= self.confine_prob).tolist() dilation_draws = torch.randint( self.dilation_iterations_range[0], self.dilation_iterations_range[1] + 1, (N,), device=input.device ).tolist() @@ -368,7 +438,7 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ continue # Decide redistribution mode once per sample - in_seg_bool = in_seg_draws[b] + confine = confine_draws[b] # Binary masks for regions masks = regions[b] # (R, ...) @@ -415,10 +485,22 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ # `counts` is `clamp_min(1)`. to_add = torch.zeros_like(x) rand_sign = 2 * torch.rand(R, device=input.device) - 1 # random sign factor per region + # Feathered only where it is *multiplied in*, never where the statistics + # come from: `_region_moments` relies on the mask being binary (it squares + # it and calls the result a voxel count), so a soft mask there would + # quietly corrupt every per-region mean and standard deviation. Here it + # softens the edge of the confined contribution, which is the whole point. + confine_flat = mask_flat.to(x.dtype) + if confine and confine_sigma > 0.0: + confine_flat = feather_alpha(confine_flat.view(1, R, *x.shape), confine_sigma).view(R, -1) for r in range(R): + # The sign is drawn for every region whether or not it contributes, so + # `seg_unit_prob` does not shift the stream for the regions that do. + if not keep_rows[b][r]: + continue contribution = _normal_pdf(x, means[r], redist_std[r]) * rand_sign[r] - if in_seg_bool: - contribution = contribution * mask_flat[r].view(x.shape) + if confine: + contribution = contribution * confine_flat[r].view(x.shape) to_add += contribution # Normalize to_add if non-zero. `torch.where`, not `if`: the test is one @@ -484,6 +566,16 @@ class RandomPaletteGPU(ImageOnlyTransform): min_label_voxels: minimum voxels in a label to attempt the remap. label_classes: if set, restrict label remap to these class indices (None = all foreground classes present in the batch). + in_seg: per-sample probability of confining the synthesis to the foreground labels. + out_seg: per-sample probability of confining it to the background. Both firing means a + mixture of the two; neither means the whole patch, which is the default and what + this transform has always done. + seg_unit: what a unit is when only part of a side is used -- "whole", "label" or + "component". `seg_unit_prob` is the chance each unit takes part. + bg_split: "em" splits the background into intensity clusters so `out_seg` can land on + part of it. Not to be confused with the generator's own `em_*` parameters. + seg_feather: Gaussian sigma in voxels for the boundary between the composited region + and the untouched one. 0 is a hard edge. mix_prob: per-sample probability of blending the synthesis back with the original. mix_alpha: weight on the original when blending, fixed or sampled from a [lo, hi] range. 0 is the pure synthesis, 1 is the input's own contrast. Both sides of @@ -505,6 +597,13 @@ def __init__( label_remap_prob: float = 0.5, min_label_voxels: int = 4, label_classes: list[int] | None = None, + in_seg: float = 0.0, + out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -527,6 +626,7 @@ def __init__( self.label_remap_prob = label_remap_prob self.min_label_voxels = min_label_voxels self.label_classes = label_classes + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() @@ -680,7 +780,17 @@ def apply_transform( synth_z = ((synth - b_mean) / b_std * flat_m_all).reshape(B, 1, D, H, W) out = input.clone() - out[:, 0:1] = synth_z.to(input.dtype) + # Region selection last, after the foreground z-score above -- which is this + # transform's stats restore -- as it is in every other transform. `normalize=True` + # because the synthesis is masked to zero on the background while the input's + # background is z-scored air: compositing them raw would put an intensity step on + # the mask boundary that has nothing to do with the augmentation. + # + # The non-finite guard is new here. This used to write `synth_z` unconditionally, + # so a NaN out of the synthesis reached the batch. + checked = _select_and_check(self, input[:, 0], synth_z[:, 0], seg_raw, normalize=True) + if checked is not None: + out[:, 0:1] = checked.unsqueeze(1).to(input.dtype) return out diff --git a/smauglab/transforms/gpu/palette/transform.py b/smauglab/transforms/gpu/palette/transform.py index da16cdb..5af481a 100644 --- a/smauglab/transforms/gpu/palette/transform.py +++ b/smauglab/transforms/gpu/palette/transform.py @@ -18,6 +18,16 @@ ) from smauglab.transforms.gpu.palette.factory import make_initial, make_overlay, make_refinement from smauglab.transforms.gpu.palette.overlay import AnatomicalLabelOverlay +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _select_and_check, + configure_regions, +) from smauglab.transforms.rng import shared_choice #: The composition `RandomPaletteGPU` hard-codes, which is what this defaults to. @@ -87,6 +97,13 @@ def __init__( dark_threshold: float = 0.01, blur_sigmas_pre: Sequence[float] = (0.0, 0.0, 0.0, 0.3, 0.5, 0.8), blur_sigmas_post: Sequence[float] = (0.0, 0.0, 0.0, 0.3, 0.5, 0.8), + in_seg: float = 0.0, + out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 1.0, @@ -103,6 +120,7 @@ def __init__( self.dark_threshold = float(dark_threshold) self.blur_sigmas_pre = list(blur_sigmas_pre) self.blur_sigmas_post = list(blur_sigmas_post) + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) @torch.no_grad() @@ -180,5 +198,10 @@ def apply_transform( synth_z = ((synth - b_mean) / b_std * flat_m_all).reshape(B, 1, D, H, W) out = input.clone() - out[:, 0:1] = synth_z.to(input.dtype) + # Region selection last, after the foreground z-score -- this transform's stats + # restore -- as everywhere else. `normalize=True` because the synthesis is masked + # to zero on the background while the input's background is z-scored air. + checked = _select_and_check(self, input[:, 0], synth_z[:, 0], seg_raw, normalize=True) + if checked is not None: + out[:, 0:1] = checked.unsqueeze(1).to(input.dtype) return out diff --git a/smauglab/transforms/gpu/regions.py b/smauglab/transforms/gpu/regions.py index a34c543..f3a291b 100644 --- a/smauglab/transforms/gpu/regions.py +++ b/smauglab/transforms/gpu/regions.py @@ -686,6 +686,42 @@ def draw_feather(transform: RegionSelecting) -> float: return float(torch.rand(1).item()) * (high - low) + low +def _keep_draw(transform: RegionSelecting, batch: int, columns: int, width: int, device: torch.device) -> Tensor: + """`[batch, columns]` of which units are kept, at a fixed draw width. + + Spends nothing at `seg_unit_prob >= 1`, which means "every unit" and is therefore the + setting that turns granularity off. Same rule as `in_seg = out_seg = 0` spending + nothing in :func:`choose_region_modes`, and it is what lets `seg_unit="label"` at + probability 1 reproduce `"whole"` bit for bit rather than merely in distribution. + """ + if transform.seg_unit_prob >= 1.0: + return torch.ones(batch, columns, dtype=torch.bool, device=device) + return torch.rand(batch, width, device=device)[:, :columns] < transform.seg_unit_prob + + +def keep_units(transform: RegionSelecting, units: RegionUnits, batch: int, device: torch.device) -> Tensor: + """`[batch, n_units]` of which units take part, with no reference to the in/out mode. + + For the transforms whose own work is already per region -- `RandomRedistributeSegGPU` + adds a Gaussian per region and then sums them -- rather than a composite over a + selected area. There is no "inside or outside" to choose there, only which of the + regions it already iterates should contribute, so this is the half of + :func:`region_alpha` that applies. + + The fall-back rule applies as it does there: if the draw kept nothing, keep + everything, so a sample never silently loses the whole augmentation. + """ + width = UNIT_DRAW_WIDTH[transform.seg_unit] + transform.bg_em_clusters[1] + columns = min(units.n_units, width) + keep = _keep_draw(transform, batch, columns, width, device) + keep = keep | ~keep.any(dim=1, keepdim=True) + if units.n_units > columns: + # The fold: units past the draw width share the last column's fate, as they do in + # `region_alpha`. + keep = torch.cat([keep, keep[:, -1:].expand(batch, units.n_units - columns)], dim=1) + return keep + + def region_alpha(transform: RegionSelecting, units: RegionUnits, modes: Tensor) -> Tensor: """The `[B, 1, *spatial]` composite weight: 1 where the transform applies, 0 where not. @@ -710,7 +746,7 @@ def region_alpha(transform: RegionSelecting, units: RegionUnits, modes: Tensor) columns = min(units.n_units, width) column = units.ids.clamp(0, columns - 1) - keep = torch.rand(modes.numel(), width, device=device)[:, :columns] < transform.seg_unit_prob + keep = _keep_draw(transform, modes.numel(), columns, width, device) # Background units come first, and only the high end is folded, so a column is # foreground exactly when its index is at or past `n_background`. is_foreground = torch.arange(columns, device=device) >= units.n_background @@ -749,8 +785,8 @@ def feather_alpha(alpha: Tensor, sigma: float) -> Tensor: `max_pool3d` and a second parameter, and worth adding only if the shrink turns out to matter. - Two-dimensional spatial inputs pass through unfeathered: the blur is 3-D. The sigma is - still drawn, because the draw must not depend on the data. + The `dim` guard is defensive only: `apply_region_selection` rejects anything but three + spatial dimensions before it gets here, because the blur is 3-D. """ if sigma <= 0.0 or alpha.dim() != 5: return alpha @@ -795,8 +831,12 @@ def apply_region_selection( """ if seg is None or not region_selection_active(transform): return transformed + if orig.dim() not in (3, 4): + # The contract `_apply_region_mode` had: a channel of a batch, or one sample's + # volume. Three spatial dimensions either way, because the feather blur is 3-D. + raise ValueError(f"Only 4D [N, D, H, W] and 3D [D, H, W] channels are supported. Got {orig.dim()}D.") - batched = orig.dim() >= 4 + batched = orig.dim() == 4 reference = orig if batched else orig.unsqueeze(0) modes = choose_region_modes(transform, reference.shape[0], orig.device) assert modes is not None # region_selection_active was just checked @@ -823,3 +863,32 @@ def apply_region_selection( reduce_dims = tuple(range(1, orig.dim())) if batched else tuple(range(orig.dim())) transformed = _rescaled_into(transformed, orig, reduce_dims) return weight * transformed + (1.0 - weight) * orig + + +def _select_and_check( + transform: RegionSelecting, + orig: Tensor, + x: Tensor, + seg_mask: Tensor | None, + note: str = "", + normalize: bool = False, +) -> Tensor | None: + """Apply region selection, then reject the result if it went non-finite. + + Returns None when the channel should be left as it was -- the "Final safety" check + that closed all nine of these loops identically. + + Takes the transform rather than its region attributes: every call site read exactly + the same ones off `self`, and spelling them out made the call longer than the code it + replaces. The selection itself lives in `gpu/regions.py`, which is also what decides + whether it runs at all -- a transform with both probabilities at zero gets its input + back having drawn nothing. + """ + x = apply_region_selection(transform, orig, x, seg_mask, normalize=normalize) + # One `isfinite(...).all()`, not `isnan().any() or isinf().any()`: each of those + # reads a reduction back to the host, and the `or` forces the second on the common + # path. Runs for every channel of every region-selecting transform. + if not bool(torch.isfinite(x).all()): + print(f"Warning nan: {type(transform).__name__}{note}", flush=True) + return None + return x diff --git a/smauglab/transforms/synthseg/transforms.py b/smauglab/transforms/synthseg/transforms.py index 735cc2b..ec6f201 100644 --- a/smauglab/transforms/synthseg/transforms.py +++ b/smauglab/transforms/synthseg/transforms.py @@ -21,6 +21,7 @@ import json import os +from collections.abc import Sequence from typing import Any import torch @@ -29,6 +30,16 @@ from smauglab.registry import AugId, AugType, Backend, register from smauglab.transforms.gpu.base import ImageOnlyTransform, segmentation_from from smauglab.transforms.gpu.mixing import DEFAULT_MIX_ALPHA, MixAlphaSpec, configure_mix, mix_with_original +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + SegFeatherSpec, + _select_and_check, + configure_regions, +) from smauglab.transforms.synthseg.generator import SynthSegGenerator # Keys understood from the JSON config / kwargs, forwarded to SynthSegGenerator. @@ -124,6 +135,13 @@ def __init__( self, apply_to_channel: list[int] | None = None, same_on_batch: bool = False, + in_seg: float = 0.0, + out_seg: float = 0.0, + seg_unit: str = DEFAULT_SEG_UNIT, + seg_unit_prob: float = DEFAULT_SEG_UNIT_PROB, + bg_split: str = DEFAULT_BG_SPLIT, + bg_em_clusters: Sequence[int] = DEFAULT_BG_EM_CLUSTERS, + seg_feather: SegFeatherSpec = DEFAULT_SEG_FEATHER, mix_prob: float = 0.0, mix_alpha: MixAlphaSpec = DEFAULT_MIX_ALPHA, p: float = 0.5, @@ -134,8 +152,15 @@ def __init__( super().__init__(p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim) self.apply_to_channel = apply_to_channel if apply_to_channel is not None else [0] # Named explicitly rather than left to **kwargs: those are filtered against - # `_GENERATOR_KEYS` on the way to SynthSegGenerator, so a mix parameter arriving - # that way would be dropped without a word. + # `_GENERATOR_KEYS` on the way to SynthSegGenerator, so a mix or region parameter + # arriving that way would be dropped without a word. + # + # `bg_em_clusters` and the generator's own `em_background_clusters_range` are + # different features one word apart. The generator's splits the *label map it + # synthesises from*, to cope with an incomplete segmentation; this one splits the + # background into units that `out_seg` can land on part of. Setting one does not + # set the other. + configure_regions(self, in_seg, out_seg, seg_unit, seg_unit_prob, bg_split, bg_em_clusters, seg_feather) configure_mix(self, mix_prob, mix_alpha) gen_kwargs = _filter_generator_kwargs(kwargs) # Intensity-only: never deform/flip internally (geometry comes from the @@ -165,10 +190,16 @@ def apply_transform(self, input: Tensor, params: dict[str, Tensor], flags: dict[ src = n if n < synth.shape[1] else synth.shape[1] - 1 x = synth[:, src] x = mix_with_original(self, input[:, c], x) - if torch.isnan(x).any() or torch.isinf(x).any(): - print(f"Warning nan: {self.__class__.__name__}", flush=True) + # `_select_and_check` subsumes the hand-written guard this replaces, and costs + # one reduction where `isnan().any() or isinf().any()` read two back to the + # host. `normalize=True` because the synthesis is a [0, 1] GMM render while the + # patch it is composited onto is z-scored; without it a confined synthesis would + # leave an intensity step on the mask boundary, which is a feature the network + # would happily learn. + checked = _select_and_check(self, input[:, c], x, seg, normalize=True) + if checked is None: continue - input[:, c] = x + input[:, c] = checked return input diff --git a/unit_tests/test_redistribute_autocast.py b/unit_tests/test_redistribute_autocast.py index cef98bc..e9d7934 100644 --- a/unit_tests/test_redistribute_autocast.py +++ b/unit_tests/test_redistribute_autocast.py @@ -77,7 +77,7 @@ def test_it_moves_the_patch(self): with self.subTest(seed=seed): torch.manual_seed(seed) img, seg = pair() - transform = RandomRedistributeSegGPU(in_seg=1.0, p=1.0, p_batch=1.0) + transform = RandomRedistributeSegGPU(confine_prob=1.0, p=1.0, p_batch=1.0) with torch.autocast("cpu", torch.float16): out = transform.apply_transform(img.clone(), {"seg": seg}, {}) @@ -91,7 +91,7 @@ def test_it_moves_the_patch(self): def test_autocast_does_not_change_the_result(self): """The moments are the only thing autocast could touch here, and they are pinned.""" img, seg = pair() - transform = RandomRedistributeSegGPU(in_seg=1.0, p=1.0, p_batch=1.0) + transform = RandomRedistributeSegGPU(confine_prob=1.0, p=1.0, p_batch=1.0) torch.manual_seed(0) plain = transform.apply_transform(img.clone(), {"seg": seg}, {}) diff --git a/unit_tests/test_redistribute_in_seg.py b/unit_tests/test_redistribute_confine.py similarity index 71% rename from unit_tests/test_redistribute_in_seg.py rename to unit_tests/test_redistribute_confine.py index d3e19c0..0437711 100644 --- a/unit_tests/test_redistribute_in_seg.py +++ b/unit_tests/test_redistribute_confine.py @@ -1,13 +1,19 @@ -"""`in_seg` must confine the redistribution to the segmentation. +"""`confine_prob` must confine the redistribution to the segmentation. + +The parameter was called `in_seg` until this file was renamed with it: everywhere else +in the package `in_seg` is the probability of applying the whole transform inside the +segmentation, while here it decides whether each label's Gaussian is restricted to that +label's own voxels. Two different questions under one name, and twenty-eight shipped +configs set the old one. `aug_redistribute_seg` builds a `to_add` field -- zero outside the labels on the -`in_seg` branch -- and then normalised it with +confined branch -- and then normalised it with img += 2 * (to_add - to_add.min()) / (to_add.max() - to_add.min() + 1e-6) A min-max rescale moves zero. A voxel where nothing was redistributed picked up `-2 * min / range`, which is non-zero whenever `min < 0` -- and it always is, -because the per-label scale is `2 * rand - 1` on [-1, 1]. So `in_seg` restricted +because the per-label scale is `2 * rand - 1` on [-1, 1]. So confining restricted where the redistribution was *computed* but not where it landed: the whole patch got a uniform DC offset. @@ -36,21 +42,21 @@ def normalised(img: torch.Tensor) -> torch.Tensor: return (img - img.min()) / (img.max() - img.min()).clamp_min(1e-6) -class TestInSegStaysInsideTheSegmentation(SmaugLabTestCase): +class TestConfiningStaysInsideTheSegmentation(SmaugLabTestCase): def test_nothing_changes_outside_the_labels(self): for seed in range(5): with self.subTest(seed=seed): torch.manual_seed(seed) img, seg = pair() - out, _ = aug_redistribute_seg(img.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(img.clone(), seg, confine_prob=1.0) delta = out - normalised(img) self.assertAlmostEqual( float(delta[seg == 0].abs().max()), 0.0, places=6, - msg="in_seg applied a DC offset to the background", + msg="confining applied a DC offset to the background", ) def test_something_still_changes_inside_the_labels(self): @@ -58,7 +64,7 @@ def test_something_still_changes_inside_the_labels(self): torch.manual_seed(0) img, seg = pair() - out, _ = aug_redistribute_seg(img.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(img.clone(), seg, confine_prob=1.0) delta = out - normalised(img) self.assertGreater(float(delta[seg == 1].abs().max()), 0.1) @@ -68,17 +74,17 @@ def test_the_peak_amplitude_is_unchanged(self): torch.manual_seed(0) img, seg = pair() - out, _ = aug_redistribute_seg(img.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(img.clone(), seg, confine_prob=1.0) delta = (out - normalised(img)).abs().max() self.assertLessEqual(float(delta), 2.0 + 1e-4) def test_the_whole_image_branch_still_redistributes(self): - """`in_seg=0` takes the other branch, which must keep working.""" + """`confine_prob=0` takes the other branch, which must keep working.""" torch.manual_seed(0) img, seg = pair() - out, _ = aug_redistribute_seg(img.clone(), seg, in_seg=0.0) + out, _ = aug_redistribute_seg(img.clone(), seg, confine_prob=0.0) delta = out - normalised(img) self.assertGreater(float(delta.abs().max()), 0.1) @@ -89,6 +95,6 @@ def test_a_constant_patch_is_still_finite(self): torch.manual_seed(0) _, seg = pair() - out, _ = aug_redistribute_seg(torch.full(SHAPE, -2.709), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(torch.full(SHAPE, -2.709), seg, confine_prob=1.0) self.assertTrue(bool(torch.isfinite(out).all())) diff --git a/unit_tests/test_region_and_stats.py b/unit_tests/test_region_and_stats.py index 896aff1..5b49589 100644 --- a/unit_tests/test_region_and_stats.py +++ b/unit_tests/test_region_and_stats.py @@ -21,9 +21,15 @@ RandomHistogramEqualizationGPU, RandomScharrGPU, RandomSqrtGPU, - _apply_region_mode, _foreground, ) +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + apply_region_selection, +) from unit_tests.helpers import SmaugLabTestCase, first_output @@ -62,7 +68,14 @@ def test_background_stays_background(self): class TestApplyRegionMode(SmaugLabTestCase): - """The behaviour a config actually asks for when it sets in_seg or out_seg.""" + """The behaviour a config actually asks for when it sets in_seg or out_seg. + + The mode is no longer passed in as a string: `in_seg` and `out_seg` are drawn per + sample, so 1.0/0.0 is how a test pins mode "in". `seg_unit_prob=1.0` keeps every unit, + which is what makes these assertions about *where* rather than about which units the + keep draw happened to pick, and `seg_feather=0` keeps the boundary a step so the + changed-voxel set is exactly the selected region. + """ def setUp(self): super().setUp() @@ -74,29 +87,43 @@ def setUp(self): self.mask[0, 0, :2, :2, :2] = 1.0 self.inside = (slice(None), slice(0, 2), slice(0, 2), slice(0, 2)) + def _select(self, **kwargs): + return apply_region_selection(_Stub(**kwargs), self.orig, self.transformed, self.mask) + def test_mode_in_changes_only_the_masked_voxels(self): - out = _apply_region_mode(self.orig, self.transformed, self.mask, "in") + out = self._select(in_seg=1.0) self.assertTrue(bool((out[self.inside] == 1.0).all()), "'in' did not apply the transform inside the mask") self.assertEqual(int(out.sum()), 8, "'in' leaked outside the mask") def test_mode_out_changes_only_the_unmasked_voxels(self): - out = _apply_region_mode(self.orig, self.transformed, self.mask, "out") + out = self._select(out_seg=1.0) self.assertTrue(bool((out[self.inside] == 0.0).all()), "'out' applied the transform inside the mask") self.assertEqual(int(out.sum()), 64 - 8) def test_in_and_out_partition_the_volume(self): - inside = _apply_region_mode(self.orig, self.transformed, self.mask, "in") - outside = _apply_region_mode(self.orig, self.transformed, self.mask, "out") + inside = self._select(in_seg=1.0) + outside = self._select(out_seg=1.0) self.assertTrue(torch.equal(inside + outside, self.transformed)) + def test_both_fired_is_a_mixture_and_not_the_whole_patch(self): + """The one semantic change: `in_seg` and `out_seg` both firing used to mean "all". + + At `seg_unit_prob=1.0` every unit of both sides is kept, so "both" does cover the + whole patch -- the difference from "all" only shows up below that, where the two + sides are sampled independently. + """ + covered = apply_region_selection(_Stub(in_seg=1.0, out_seg=1.0, seg_unit_prob=1.0), self.orig, self.transformed, self.mask) + self.assertTrue(torch.equal(covered, self.transformed)) + def test_mode_all_ignores_the_mask(self): - out = _apply_region_mode(self.orig, self.transformed, self.mask, "all") + """Neither probability firing. The mask is read but nothing is confined.""" + out = self._select(in_seg=0.0, out_seg=0.0) self.assertTrue(torch.equal(out, self.transformed)) def test_a_missing_mask_means_apply_everywhere(self): - out = _apply_region_mode(self.orig, self.transformed, None, "in") + out = apply_region_selection(_Stub(in_seg=1.0), self.orig, self.transformed, None) self.assertTrue(torch.equal(out, self.transformed)) def test_the_unbatched_3d_path_behaves_the_same(self): @@ -105,12 +132,28 @@ def test_the_unbatched_3d_path_behaves_the_same(self): mask = torch.zeros(2, 4, 4, 4) mask[0, :2, :2, :2] = 1.0 - out = _apply_region_mode(orig, transformed, mask, "in") + out = apply_region_selection(_Stub(in_seg=1.0), orig, transformed, mask) self.assertEqual(int(out.sum()), 8) def test_an_unsupported_rank_is_rejected(self): with self.assertRaises(ValueError): - _apply_region_mode(torch.rand(2, 2), torch.rand(2, 2), torch.rand(1, 2, 2), "in") + apply_region_selection(_Stub(in_seg=1.0), torch.rand(2, 2), torch.rand(2, 2), torch.rand(1, 2, 2)) + + def test_a_feathered_boundary_is_a_ramp_not_a_step(self): + """`seg_feather` is the whole point of the alpha being a float and not a bool.""" + step = self._select(in_seg=1.0) + ramp = apply_region_selection(_Stub(in_seg=1.0, seg_feather=(1.0, 1.0)), self.orig, self.transformed, self.mask) + + self.assertEqual(sorted(set(step.flatten().tolist())), [0.0, 1.0], "the unfeathered composite is a step") + between = (ramp > 1e-6) & (ramp < 1.0 - 1e-6) + self.assertGreater(int(between.sum()), 0, "a feathered boundary must put voxels strictly between the two") + self.assertTrue(bool((ramp >= 0.0).all() and (ramp <= 1.0).all()), "the alpha left [0, 1]") + + def test_feathered_in_and_out_still_sum_to_the_transform(self): + """blur(m) + blur(1 - m) == blur(1) == 1, for a normalised kernel and reflect padding.""" + inside = apply_region_selection(_Stub(in_seg=1.0, seg_feather=(1.0, 1.0)), self.orig, self.transformed, self.mask) + outside = apply_region_selection(_Stub(out_seg=1.0, seg_feather=(1.0, 1.0)), self.orig, self.transformed, self.mask) + self.assertTrue(torch.allclose(inside + outside, self.transformed, atol=1e-5)) class TestRegionModeReachesTheRealTransforms(SmaugLabTestCase): @@ -237,7 +280,7 @@ def test_a_2d_image_does_not_raise(self): seg = torch.zeros(12, 12) seg[3:8, 3:8] = 1.0 - out, _ = aug_redistribute_seg(image.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(image.clone(), seg, confine_prob=1.0) self.assertEqual(tuple(out.shape), tuple(image.shape)) self.assertTrue(bool(np.isfinite(out.numpy()).all())) @@ -249,18 +292,31 @@ def test_the_3d_path_is_unchanged(self): seg = torch.zeros(10, 10, 10) seg[2:6, 2:6, 2:6] = 1.0 - out, _ = aug_redistribute_seg(image.clone(), seg, in_seg=1.0) + out, _ = aug_redistribute_seg(image.clone(), seg, confine_prob=1.0) self.assertEqual(tuple(out.shape), tuple(image.shape)) class _Stub: - """Stands in for a transform: _select_and_check only reads these three attributes.""" + """Stands in for a transform: the region helpers only read these attributes. + + Built from `gpu/regions.py`'s own "off" defaults rather than from literals, so this + also fails if one of those constants disappears. + """ - def __init__(self, in_seg: float = 0.0, out_seg: float = 0.0, mix_in_out: bool = False): - self.in_seg = in_seg - self.out_seg = out_seg - self.mix_in_out = mix_in_out + def __init__(self, **overrides): + defaults = { + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": DEFAULT_SEG_UNIT, + "seg_unit_prob": 1.0, + "bg_split": DEFAULT_BG_SPLIT, + "bg_em_clusters": DEFAULT_BG_EM_CLUSTERS, + "seg_feather": (DEFAULT_SEG_FEATHER, DEFAULT_SEG_FEATHER), + } + defaults.update(overrides) + for name, value in defaults.items(): + setattr(self, name, value) class TestExtractedTailHelpers(SmaugLabTestCase): diff --git a/unit_tests/test_seg_layout.py b/unit_tests/test_seg_layout.py index c3f5e42..2a111a1 100644 --- a/unit_tests/test_seg_layout.py +++ b/unit_tests/test_seg_layout.py @@ -87,7 +87,7 @@ def test_label_map_and_one_hot_give_the_same_image(self): image = torch.rand(*SHAPE) def build(): - return RandomRedistributeSegGPU(p=1.0, in_seg=1.0) + return RandomRedistributeSegGPU(p=1.0, confine_prob=1.0) torch.testing.assert_close(_apply(build, labels, image), _apply(build, onehot, image)) @@ -99,7 +99,7 @@ def test_the_labels_actually_drive_the_result(self): image = torch.rand(*SHAPE) def build(): - return RandomRedistributeSegGPU(p=1.0, in_seg=1.0) + return RandomRedistributeSegGPU(p=1.0, confine_prob=1.0) per_label = _apply(build, labels, image) one_blob = _apply(build, merged, image) From 3693b54e0c46fe9b86c5b000efb72c4e02bcc4d3 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 15 Oct 2026 12:20:37 +0000 Subject: [PATCH 3/3] test: cover the region surface, and check the components against scipy `test_connected_components.py` checks the one new algorithm against `scipy.ndimage.label` rather than against itself: partition equality on random blobby volumes, on a serpentine ribbon whose geodesic length is far beyond its Euclidean extent (which is what the pointer jump is for), and on a multi-label volume where the count must equal the sum of scipy's per-label runs -- i.e. that one propagation pass respects label boundaries. Two tests pin the connectivity: corner-touching voxels are two components and face-touching ones are one, so a future switch to the cheaper single `max_pool3d` (a 3x3x3 neighbourhood, and therefore 26-connected) fails here rather than silently merging structures. `test_region_selection.py` is registry-driven over every GPU transform that accepts the surface, so one added later cannot arrive without it or with it wired into the constructor and never read. The "is it read" case runs several seeds, because a draw that happens to keep every label legitimately makes `seg_unit="label"` identical to `"whole"`, and skips a transform that leaves the test input alone -- `RandomClampGPU` does, which is a statement about the input rather than about the knob. Three of its cases are about the RNG stream rather than about pixels, and each is a property the replaced code broke: * region selection off, and `seg_unit_prob=1`, spend no random numbers, so a sweep is comparable with its own baseline at a fixed seed; * the draw width is a function of the configuration and never of the data -- a two-label patch and a five-label one leave the stream in the same place, which a per-unit draw would not; * `bg_split="em"` leaves the global stream exactly where `"none"` does, so whether the decomposition was cached cannot change the augmentation. Plus: the cache is shared across a pass and rebuilt per spec and per mask object; the whole-unit invariant (a label or component is applied or it is not, never partially); the feathered boundary is a monotone ramp bounded in [0, 1] whose `in` and `out` halves still sum to the full transform; an empty mask makes `in` a no-op and `out` everywhere, which is what the binary version did; and no shipped config still names `mix_in_out` or `RandomRedistributeSegGPU.in_seg`. `helpers.multi_label_seg` is new because neither existing fixture can express the granularities: `tiny_seg` is one binary cube, so it has one label value and one connected component and all three `seg_unit` settings agree on it. The new one has three values with one of them in two separated blocks, which is the only shape that tells "component" from "label" apart. Suite in seg11: 556 passed, 2701 subtests, and the same documented 11-failure baseline as before this branch. Co-Authored-By: Claude Opus 5 --- unit_tests/helpers.py | 20 + unit_tests/test_connected_components.py | 225 ++++++++ unit_tests/test_region_selection.py | 682 ++++++++++++++++++++++++ 3 files changed, 927 insertions(+) create mode 100644 unit_tests/test_connected_components.py create mode 100644 unit_tests/test_region_selection.py diff --git a/unit_tests/helpers.py b/unit_tests/helpers.py index 8c12622..2e7a035 100644 --- a/unit_tests/helpers.py +++ b/unit_tests/helpers.py @@ -49,6 +49,26 @@ def tiny_seg(self) -> torch.Tensor: seg[:, :, 6:18, 6:18, 6:18] = 1.0 return seg + def multi_label_seg(self) -> torch.Tensor: + """Three label values, one of them in two disconnected pieces. + + `tiny_seg` cannot express either new granularity: a single binary cube has one + label value and one connected component, so `seg_unit="label"` and + `seg_unit="component"` are both indistinguishable from `"whole"` on it. Label 1 + here is two separated blocks, which is the only way to tell "component" from + "label" apart, and the three values are what makes "label" more than one unit. + + Blocks rather than noise, so a test can name the voxels it expects to change. + Laid out with background between them, so dilation and feathering have somewhere + to go and the components stay separate under a one-voxel structuring element. + """ + seg = torch.zeros(*VOLUME_SHAPE, dtype=torch.float32) + seg[:, :, 2:6, 2:6, 2:6] = 1.0 # label 1, first component + seg[:, :, 2:6, 2:6, 16:20] = 1.0 # label 1, second component + seg[:, :, 10:14, 10:14, 10:14] = 2.0 + seg[:, :, 16:20, 16:20, 4:8] = 3.0 + return seg + def empty_seg(self) -> torch.Tensor: """A mask with no foreground at all -- every voxel background. diff --git a/unit_tests/test_connected_components.py b/unit_tests/test_connected_components.py new file mode 100644 index 0000000..0be50f5 --- /dev/null +++ b/unit_tests/test_connected_components.py @@ -0,0 +1,225 @@ +"""`connected_components_3d` against scipy, which is the only reference that matters. + +This is the one new algorithm in the region-selection feature, so it is checked against +`scipy.ndimage.label` rather than against its own output. scipy is an unconditional +runtime dependency (see pyproject.toml), so there is no skip guard. + +Two things the tests pin that are easy to get wrong and silent when wrong: + +* **The connectivity is 6, not 26.** The obvious GPU implementation is one + `max_pool3d(kernel=3)`, which is a 3x3x3 neighbourhood and therefore merges blobs that + touch only at a corner. `scipy.ndimage.label`'s default structuring element is + `generate_binary_structure(3, 1)`, the six-neighbour cross, and that is also the + conventional reading of "connected components of a label". +* **Components do not cross label boundaries**, and that holds in a single propagation + pass rather than one pass per label value. A spine patch carries five to twenty-five + labels, so the difference is not academic. + +Partitions are compared, not ids: the numbering is an implementation detail and only the +grouping is the contract. +""" + +from __future__ import annotations + +import numpy as np +import torch +from scipy import ndimage +from torch.nn import functional as F + +from smauglab.transforms.gpu.regions import connected_components_3d +from unit_tests.helpers import SmaugLabTestCase + +#: scipy's default structuring element, spelled out because the default is the whole point. +SIX_CONNECTED = ndimage.generate_binary_structure(3, 1) + + +def as_batch(mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """A `[D, H, W]` bool mask as the `(labels, foreground)` pair the function takes.""" + labels = mask.long().view(1, 1, *mask.shape) + return labels, labels > 0 + + +def blobs(size: int = 20, seed: int = 0) -> torch.Tensor: + """A smoothed random volume thresholded at its mean: many components, odd shapes.""" + torch.manual_seed(seed) + smoothed = F.avg_pool3d(F.pad(torch.rand(1, 1, size, size, size), [3] * 6, mode="reflect"), 7, 1, 0) + return smoothed[0, 0] > smoothed.mean() + + +def same_partition(ours: np.ndarray, theirs: np.ndarray, mask: np.ndarray) -> bool: + """Whether two id maps group the masked voxels identically. + + True when the (ours, theirs) pairs present form a bijection: every one of our + components maps to exactly one of theirs and vice versa. Comparing ids directly would + only test that two numberings happen to agree. + """ + pairs = set(zip(ours[mask].tolist(), theirs[mask].tolist())) + ours_seen = {a for a, _ in pairs} + theirs_seen = {b for _, b in pairs} + return len(pairs) == len(ours_seen) == len(theirs_seen) + + +class TestAgreesWithScipy(SmaugLabTestCase): + def test_random_blobs_partition_identically(self): + for seed in range(6): + with self.subTest(seed=seed): + mask = blobs(seed=seed) + labels, foreground = as_batch(mask) + + ids, count = connected_components_3d(labels, foreground) + reference, expected = ndimage.label(mask.numpy(), structure=SIX_CONNECTED) + + self.assertEqual(count, expected, "a different number of components than scipy found") + self.assertTrue( + same_partition(ids[0, 0].numpy(), reference, mask.numpy()), + "the same count but a different grouping", + ) + + def test_a_multi_label_volume_matches_scipy_run_per_label(self): + """One propagation pass must equal the sum of scipy's per-label runs.""" + for seed in range(4): + with self.subTest(seed=seed): + torch.manual_seed(seed) + smoothed = F.avg_pool3d(F.pad(torch.rand(1, 1, 20, 20, 20), [3] * 6, mode="reflect"), 7, 1, 0)[0, 0] + labels = torch.zeros(20, 20, 20, dtype=torch.long) + labels[smoothed > smoothed.mean()] = 1 + labels[smoothed > smoothed.mean() + 0.6 * smoothed.std()] = 2 + batched = labels.view(1, 1, 20, 20, 20) + + _ids, count = connected_components_3d(batched, batched > 0) + expected = sum(ndimage.label((labels == value).numpy(), structure=SIX_CONNECTED)[1] for value in (1, 2)) + + self.assertEqual(count, expected) + + +class TestTheConnectivityIsSix(SmaugLabTestCase): + def test_two_voxels_touching_at_a_corner_are_two_components(self): + """The test that pins 6- over 26-connectivity. A single max_pool3d merges these.""" + labels = torch.zeros(1, 1, 4, 4, 4, dtype=torch.long) + labels[0, 0, 1, 1, 1] = 1 + labels[0, 0, 2, 2, 2] = 1 + + _ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 2, "corner-touching voxels were merged, so the neighbourhood is 26 not 6") + self.assertEqual(ndimage.label((labels[0, 0] > 0).numpy(), structure=SIX_CONNECTED)[1], 2) + + def test_two_voxels_touching_at_a_face_are_one_component(self): + labels = torch.zeros(1, 1, 4, 4, 4, dtype=torch.long) + labels[0, 0, 1, 1, 1] = 1 + labels[0, 0, 1, 1, 2] = 1 + + _ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 1) + + +class TestLabelBoundariesAreRespected(SmaugLabTestCase): + def test_two_adjacent_slabs_of_different_labels_stay_apart(self): + """They touch across a face, so only the label gate keeps them separate.""" + labels = torch.zeros(1, 1, 8, 8, 8, dtype=torch.long) + labels[0, 0, :, :, :4] = 1 + labels[0, 0, :, :, 4:] = 2 + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 2, "the propagation crossed a label boundary") + self.assertEqual(ids[0, 0, :, :, :4].unique().numel(), 1) + self.assertEqual(ids[0, 0, :, :, 4:].unique().numel(), 1) + self.assertNotEqual(int(ids[0, 0, 0, 0, 0]), int(ids[0, 0, 0, 0, 7])) + + +class TestLongThinStructuresConverge(SmaugLabTestCase): + """The iteration-count test: pointer jumping is what makes these logarithmic. + + A serpentine ribbon has a geodesic length far beyond its Euclidean extent, so plain + neighbour propagation would need one round per voxel along it. + """ + + def test_a_serpentine_ribbon_is_one_component(self): + size = 24 + labels = torch.zeros(1, 1, size, size, size, dtype=torch.long) + row = 0 + for plane in range(0, size, 2): + labels[0, 0, plane, row, :] = 1 + labels[0, 0, plane : min(plane + 2, size), row, size - 1 if plane % 4 == 0 else 0] = 1 + row = min(row + 2, size - 1) + + ids, count = connected_components_3d(labels, labels > 0) + reference, expected = ndimage.label((labels[0, 0] > 0).numpy(), structure=SIX_CONNECTED) + + self.assertEqual(count, expected) + self.assertTrue(same_partition(ids[0, 0].numpy(), reference, (labels[0, 0] > 0).numpy())) + + +class TestDegenerateMasks(SmaugLabTestCase): + def test_an_empty_mask_has_no_components_and_does_not_raise(self): + labels = torch.zeros(1, 1, 5, 5, 5, dtype=torch.long) + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertTrue(bool((ids == -1).all()), "background must be -1, not a unit id") + self.assertGreaterEqual(count, 1, "the count sizes a draw and must never be zero") + + def test_a_full_volume_is_one_component(self): + labels = torch.ones(1, 1, 5, 5, 5, dtype=torch.long) + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 1) + self.assertEqual(ids.unique().tolist(), [0]) + + def test_a_single_voxel_is_one_component(self): + labels = torch.zeros(1, 1, 5, 5, 5, dtype=torch.long) + labels[0, 0, 2, 2, 2] = 1 + + _ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 1) + + +class TestPerSampleNumbering(SmaugLabTestCase): + """Ids are 0-based within each sample, because the keep draw is per sample.""" + + def test_each_sample_is_numbered_from_zero(self): + labels = torch.zeros(2, 1, 6, 6, 6, dtype=torch.long) + labels[0, 0, 0, 0, 0] = 1 + labels[1, 0, 0, 0, 0] = 1 + labels[1, 0, 2, 2, 2] = 1 + labels[1, 0, 4, 4, 4] = 1 + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 3, "the count is the largest per-sample count, which sizes the draw") + self.assertEqual(ids[0, 0][labels[0, 0] > 0].tolist(), [0]) + self.assertEqual(sorted(ids[1, 0][labels[1, 0] > 0].tolist()), [0, 1, 2]) + + def test_two_samples_rooted_at_the_same_index_are_not_merged(self): + """Seeds are globally unique for exactly this reason.""" + labels = torch.zeros(2, 1, 4, 4, 4, dtype=torch.long) + labels[:, 0, 0, 0, 0] = 1 + + ids, count = connected_components_3d(labels, labels > 0) + + self.assertEqual(count, 1) + self.assertEqual(int(ids[0, 0, 0, 0, 0]), 0) + self.assertEqual(int(ids[1, 0, 0, 0, 0]), 0) + + +class TestIdsAreUsableAsIndices(SmaugLabTestCase): + """`region_alpha` gathers a per-unit draw with these, so they must be dense.""" + + def test_ids_are_contiguous_from_zero(self): + mask = blobs(seed=3) + labels, foreground = as_batch(mask) + + ids, count = connected_components_3d(labels, foreground) + + present = ids[foreground].unique().tolist() + self.assertEqual(present, list(range(count))) + + +if __name__ == "__main__": + import unittest + + unittest.main() diff --git a/unit_tests/test_region_selection.py b/unit_tests/test_region_selection.py new file mode 100644 index 0000000..a66f24d --- /dev/null +++ b/unit_tests/test_region_selection.py @@ -0,0 +1,682 @@ +"""Where a transform applies, at what granularity, and how sharply it stops. + +`in_seg` and `out_seg` predate this; what they did not express is most of what they +looked like they expressed. The parameters and the four defects they replace are set out +in `smauglab/transforms/gpu/regions.py`; this file is what holds the behaviour. + +The registry-driven cases matter most. The surface has to reach every transform that +reads a segmentation, and a transform added later must not be able to arrive without it +-- or, worse, with the parameters wired into its constructor and never read. + +Three invariants here are about the **RNG stream** rather than about pixels, and all +three are properties the code this replaces broke: + +* a knob that is off spends no random numbers, so a sweep is comparable with its own + baseline at a fixed seed; +* the number of uniforms a draw spends depends on the configuration and never on the + data, so a four-label patch and a five-label one leave the stream in the same place; +* the EM background split is forked, so a cache hit and a cache miss agree. +""" + +from __future__ import annotations + +import gc + +import torch + +from smauglab import registry +from smauglab.config import validate_section +from smauglab.registry import AugType, Backend +from smauglab.transforms.build import build_gpu_pipeline +from smauglab.transforms.gpu import regions +from smauglab.transforms.gpu.base import AugmentationSequentialCustom +from smauglab.transforms.gpu.contrast import RandomScharrGPU +from smauglab.transforms.gpu.regions import ( + DEFAULT_BG_EM_CLUSTERS, + DEFAULT_BG_SPLIT, + DEFAULT_SEG_FEATHER, + DEFAULT_SEG_UNIT, + DEFAULT_SEG_UNIT_PROB, + UnitSpec, + apply_region_selection, + configure_regions, + region_units, +) +from unit_tests.helpers import SmaugLabTestCase, domain_bank_missing, first_output + +BATCH = 2 +SHAPE = (BATCH, 1, 24, 24, 24) + +#: `RandomRedistributeSegGPU` is the one transform that deliberately does not take +#: `in_seg` / `out_seg`: it used to accept `in_seg` with a different meaning, and had it +#: kept the name while changing it, twenty-eight shipped configs would have flipped +#: silently. See its class docstring. +NO_MODE_PARAMETERS = {"RandomRedistributeSegGPU"} + +#: The five that are not about the in/out mode and therefore reach every segmentation-aware +#: transform, including the one above. +GRANULARITY_PARAMETERS = ("seg_unit", "seg_unit_prob", "bg_split", "bg_em_clusters", "seg_feather") + + +def region_aware(): + """Every registered GPU transform that takes the region surface.""" + return [entry for entry in registry.entries(Backend.GPU) if "seg_unit" in registry.accepted_params(entry)] + + +def transfer_augmentations(): + return registry.entries(Backend.GPU, group=AugType.TA) + + +class _Stub: + """Stands in for a transform. The helpers only read these attributes.""" + + def __init__(self, **overrides): + values = { + "in_seg": 0.0, + "out_seg": 0.0, + "seg_unit": DEFAULT_SEG_UNIT, + "seg_unit_prob": 1.0, + "bg_split": DEFAULT_BG_SPLIT, + "bg_em_clusters": DEFAULT_BG_EM_CLUSTERS, + "seg_feather": (DEFAULT_SEG_FEATHER, DEFAULT_SEG_FEATHER), + } + values.update(overrides) + for name, value in values.items(): + setattr(self, name, value) + + +class TestTheSurfaceReachesEveryTransform(SmaugLabTestCase): + """Registry-driven, so a transform added later is covered when it is registered.""" + + def test_every_transfer_augmentation_takes_the_granularity_parameters(self): + for entry in transfer_augmentations(): + accepted = registry.accepted_params(entry) + with self.subTest(entry=entry.name): + for name in GRANULARITY_PARAMETERS: + self.assertIn(name, accepted, f"{entry.name} cannot be given {name}") + + def test_every_transfer_augmentation_but_one_takes_the_mode_parameters(self): + for entry in transfer_augmentations(): + accepted = registry.accepted_params(entry) + expected = entry.name not in NO_MODE_PARAMETERS + with self.subTest(entry=entry.name): + self.assertEqual("in_seg" in accepted, expected) + self.assertEqual("out_seg" in accepted, expected) + + def test_mix_in_out_is_gone_everywhere(self): + for entry in registry.entries(Backend.GPU): + with self.subTest(entry=entry.name): + self.assertNotIn("mix_in_out", registry.accepted_params(entry)) + + def test_the_defaults_mean_off(self): + """A config that sets none of these must get the behaviour it had before.""" + for entry in region_aware(): + accepted = registry.accepted_params(entry) + with self.subTest(entry=entry.name): + self.assertEqual(accepted["bg_split"].default, DEFAULT_BG_SPLIT) + self.assertEqual(accepted["seg_feather"].default, DEFAULT_SEG_FEATHER) + self.assertEqual(tuple(accepted["bg_em_clusters"].default), DEFAULT_BG_EM_CLUSTERS) + if entry.name in NO_MODE_PARAMETERS: + # Its own defaults reproduce the region set it has always used. + self.assertEqual(accepted["seg_unit"].default, "label") + self.assertEqual(accepted["seg_unit_prob"].default, 1.0) + else: + self.assertEqual(accepted["in_seg"].default, 0.0) + self.assertEqual(accepted["out_seg"].default, 0.0) + self.assertEqual(accepted["seg_unit"].default, DEFAULT_SEG_UNIT) + self.assertEqual(accepted["seg_unit_prob"].default, DEFAULT_SEG_UNIT_PROB) + + def test_the_knobs_are_read_and_not_merely_declared(self): + """The guard against a constructor that accepts a parameter and ignores it. + + Several seeds, because a single draw can legitimately keep every label and then + `seg_unit="label"` *is* `"whole"`. A transform that leaves the volume alone on + this input is skipped rather than failed: there is nothing to confine, which is + a statement about the input and not about the knob. + """ + image = torch.rand(*SHAPE) + seg = self.multi_label_seg().expand(BATCH, -1, -1, -1, -1).contiguous() + for entry in region_aware(): + with self.subTest(entry=entry.name): + skip = domain_bank_missing(entry) if entry.external_asset else None + if skip: + self.skipTest(skip) + extra = dict(entry.smoke_kwargs) + mode = {} if entry.name in NO_MODE_PARAMETERS else {"in_seg": 1.0} + + def run(seed, entry=entry, extra=extra, **kwargs): + torch.manual_seed(seed) + transform = entry.cls(p=1.0, **extra, **kwargs) + torch.manual_seed(seed) + return transform.apply_transform(image.clone(), {"seg": seg.clone()}, transform.flags) + + untouched = all(bool(torch.equal(run(seed, **mode), image)) for seed in range(4)) + if untouched: + self.skipTest(f"{entry.name} leaves this input unchanged, so there is nothing to confine") + + differed = any( + not bool( + torch.equal( + run(seed, **mode, seg_unit="whole", seg_unit_prob=0.5), + run(seed, **mode, seg_unit="label", seg_unit_prob=0.5), + ) + ) + for seed in range(6) + ) + self.assertTrue( + differed, + f"{entry.name} accepts seg_unit but its output does not depend on it", + ) + + +class TestWhereEachModeLands(SmaugLabTestCase): + """`seg_unit_prob=1.0` keeps every unit, so these are assertions about *where*.""" + + def setUp(self): + super().setUp() + self.orig = torch.zeros(1, 24, 24, 24) + self.transformed = torch.ones(1, 24, 24, 24) + self.seg = self.multi_label_seg() + self.foreground = self.seg[:, 0] > 0 + + def changed(self, **kwargs): + out = apply_region_selection(_Stub(**kwargs), self.orig, self.transformed, self.seg.clone()) + return out != self.orig + + def test_in_covers_exactly_the_labels(self): + self.assertTrue(bool(torch.equal(self.changed(in_seg=1.0), self.foreground))) + + def test_out_covers_exactly_the_background(self): + self.assertTrue(bool(torch.equal(self.changed(out_seg=1.0), ~self.foreground))) + + def test_both_covers_the_patch_when_every_unit_is_kept(self): + self.assertTrue(bool(self.changed(in_seg=1.0, out_seg=1.0).all())) + + def test_neither_is_the_whole_patch_untouched_by_the_mask(self): + out = apply_region_selection(_Stub(), self.orig, self.transformed, self.seg) + self.assertTrue(bool(torch.equal(out, self.transformed))) + + def test_both_is_a_mixture_and_not_merely_all(self): + """At a probability below 1 the two sides are sampled independently. + + This is the one semantic change to `in_seg` / `out_seg`: both firing used to fall + through to "the whole patch", so `in_seg=1, out_seg=1` was indistinguishable from + `in_seg=0, out_seg=0`. + """ + torch.manual_seed(4) + partials = 0 + for _ in range(12): + changed = self.changed(in_seg=1.0, out_seg=1.0, seg_unit="label", seg_unit_prob=0.5) + touches_both = bool((changed & self.foreground).any()) and bool((changed & ~self.foreground).any()) + if touches_both and not bool(changed.all()): + partials += 1 + self.assertGreater(partials, 0, "'both' never produced a mixture, only the whole patch") + + +class TestGranularityKeepsWholeUnits(SmaugLabTestCase): + """A unit is applied or it is not; it is never partially applied.""" + + def setUp(self): + super().setUp() + self.orig = torch.zeros(1, 24, 24, 24) + self.transformed = torch.ones(1, 24, 24, 24) + self.seg = self.multi_label_seg() + + def assert_whole_units(self, seg_unit): + spec = UnitSpec(seg_unit, DEFAULT_BG_SPLIT, DEFAULT_BG_EM_CLUSTERS) + for seed in range(8): + torch.manual_seed(seed) + seg = self.seg.clone() + units = region_units(seg, self.orig.unsqueeze(1), spec) + out = apply_region_selection(_Stub(in_seg=1.0, seg_unit=seg_unit, seg_unit_prob=0.5), self.orig, self.transformed, seg) + changed = (out != self.orig)[0] + for unit in range(units.n_background, units.n_units): + cells = units.ids[0, 0] == unit + if not bool(cells.any()): + continue + values = changed[cells] + with self.subTest(seg_unit=seg_unit, seed=seed, unit=unit): + self.assertTrue( + bool(values.all()) or bool((~values).all()), + f"{seg_unit} unit {unit} was partially applied", + ) + + def test_label_units_are_whole(self): + self.assert_whole_units("label") + + def test_component_units_are_whole(self): + self.assert_whole_units("component") + + def test_components_split_a_label_that_labels_does_not(self): + """`multi_label_seg`'s label 1 is two separated blocks.""" + by_label = region_units(self.seg.clone(), None, UnitSpec("label", "none", DEFAULT_BG_EM_CLUSTERS)) + by_component = region_units(self.seg.clone(), None, UnitSpec("component", "none", DEFAULT_BG_EM_CLUSTERS)) + self.assertEqual(by_label.n_units, 4, "background plus three label values") + self.assertEqual(by_component.n_units, 5, "the split label contributes one more") + + def test_keeping_every_unit_reduces_to_whole(self): + """`seg_unit_prob=1.0` is how granularity is turned off, and it must be exact.""" + for seg_unit in ("label", "component"): + with self.subTest(seg_unit=seg_unit): + whole = apply_region_selection( + _Stub(in_seg=1.0, seg_unit="whole", seg_unit_prob=1.0), self.orig, self.transformed, self.seg.clone() + ) + granular = apply_region_selection( + _Stub(in_seg=1.0, seg_unit=seg_unit, seg_unit_prob=1.0), self.orig, self.transformed, self.seg.clone() + ) + self.assertTrue(bool(torch.equal(whole, granular))) + + def test_a_single_label_mask_makes_label_the_same_as_whole(self): + """Documented rather than discovered: there is nothing to mix on a binary mask.""" + binary = self.tiny_seg() + whole = apply_region_selection(_Stub(in_seg=1.0, seg_unit="whole"), self.orig, self.transformed, binary.clone()) + by_label = apply_region_selection(_Stub(in_seg=1.0, seg_unit="label"), self.orig, self.transformed, binary.clone()) + self.assertTrue(bool(torch.equal(whole, by_label))) + + +class TestTheBackgroundSplit(SmaugLabTestCase): + """`bg_split="em"` is what makes the background something other than all-or-nothing.""" + + def setUp(self): + super().setUp() + self.seg = self.tiny_seg() + self.foreground = self.seg[:, 0] > 0 + # Bimodal background, so the EM has two modes to find. + torch.manual_seed(1) + self.image = torch.zeros(1, 24, 24, 24) + self.image[0, :12] = 0.2 + self.image[0, 12:] = 0.9 + self.image += torch.rand(1, 24, 24, 24) * 0.02 + self.transformed = torch.full((1, 24, 24, 24), 5.0) + + def changed_background(self, **kwargs): + out = apply_region_selection(_Stub(**kwargs), self.image, self.transformed, self.seg.clone()) + return int(((out != self.image) & ~self.foreground).sum()) + + def test_without_the_split_the_background_is_all_or_nothing(self): + torch.manual_seed(5) + total = int((~self.foreground).sum()) + seen = {self.changed_background(out_seg=1.0, seg_unit_prob=0.5) for _ in range(8)} + self.assertTrue(seen <= {0, total}, f"a partial background without bg_split: {sorted(seen)}") + + def test_with_the_split_part_of_the_background_can_be_used(self): + torch.manual_seed(5) + total = int((~self.foreground).sum()) + counts = [self.changed_background(out_seg=1.0, bg_split="em", seg_unit_prob=0.5) for _ in range(8)] + self.assertTrue( + any(0 < count < total for count in counts), + f"bg_split='em' never produced a partial background: {counts}", + ) + + def test_the_foreground_is_still_untouched_in_mode_out(self): + torch.manual_seed(5) + out = apply_region_selection(_Stub(out_seg=1.0, bg_split="em", seg_unit_prob=0.5), self.image, self.transformed, self.seg) + self.assertTrue(bool(torch.equal(out[self.foreground], self.image[self.foreground]))) + + def test_a_patch_with_no_background_is_not_split(self): + """`em_subdivide_labels` falls back to the largest label when 0 is absent.""" + units = region_units(torch.ones(1, 1, 8, 8, 8), torch.rand(1, 1, 8, 8, 8), UnitSpec("whole", "em", (3, 10))) + self.assertEqual(units.n_background, 1, "the EM fallback split a foreground label") + + def test_a_constant_volume_gives_one_background_unit(self): + """The alternative is clusters whose boundaries are decided by float noise.""" + seg = torch.zeros(1, 1, 10, 10, 10) + seg[0, 0, 4:6, 4:6, 4:6] = 1 + units = region_units(seg, self.constant_volume()[:, :, :10, :10, :10], UnitSpec("whole", "em", (3, 10))) + self.assertEqual(units.n_background, 1) + self.assertTrue(bool(torch.isfinite(units.ids.float()).all())) + + +class TestTheBoundary(SmaugLabTestCase): + """`seg_feather` turns the composite weight from a step into a Gaussian ramp.""" + + def setUp(self): + super().setUp() + self.orig = torch.zeros(1, 24, 24, 24) + self.transformed = torch.ones(1, 24, 24, 24) + self.seg = self.tiny_seg() + + def composite(self, sigma): + return apply_region_selection(_Stub(in_seg=1.0, seg_feather=(sigma, sigma)), self.orig, self.transformed, self.seg.clone()) + + def test_zero_is_a_hard_step(self): + self.assertEqual(sorted(set(self.composite(0.0).flatten().tolist())), [0.0, 1.0]) + + def test_a_positive_sigma_puts_voxels_between_the_two(self): + ramp = self.composite(1.0) + between = (ramp > 1e-6) & (ramp < 1.0 - 1e-6) + self.assertGreater(int(between.sum()), 0) + self.assertTrue(bool((ramp >= 0.0).all() and (ramp <= 1.0).all()), "the weight left [0, 1]") + + def test_the_ramp_is_monotone_across_the_boundary(self): + ramp = self.composite(1.5)[0, 12, 12, :] + rising = ramp[:12] + self.assertTrue(bool((rising[1:] - rising[:-1] >= -1e-6).all()), rising.tolist()) + + def test_the_interior_of_a_unit_stays_at_full_strength(self): + """The ramp is centred on the boundary, so it does not reach a large unit's middle.""" + self.assertGreater(float(self.composite(1.0)[0, 11:13, 11:13, 11:13].min()), 0.999) + + def test_in_and_out_still_sum_to_the_transform(self): + """blur(m) + blur(1 - m) == blur(1) == 1 for a normalised kernel and reflect padding.""" + inside = self.composite(1.0) + outside = apply_region_selection(_Stub(out_seg=1.0, seg_feather=(1.0, 1.0)), self.orig, self.transformed, self.seg.clone()) + self.assertTrue(torch.allclose(inside + outside, self.transformed, atol=1e-5)) + + def test_mode_all_is_unaffected(self): + """The blur of a constant weight is that constant.""" + out = apply_region_selection(_Stub(in_seg=1e-12, seg_feather=(2.0, 2.0)), self.orig, self.transformed, self.seg) + self.assertTrue(torch.allclose(out, self.transformed, atol=1e-6)) + + def test_an_empty_mask_is_unaffected(self): + out = apply_region_selection(_Stub(in_seg=1.0, seg_feather=(2.0, 2.0)), self.orig, self.transformed, self.empty_seg()) + self.assertTrue(torch.allclose(out, self.orig, atol=1e-6)) + + +class TestDegenerateMasks(SmaugLabTestCase): + """An empty mask is ~18% of real nnU-Net samples; see `helpers.empty_seg`.""" + + def setUp(self): + super().setUp() + self.orig = torch.zeros(1, 24, 24, 24) + self.transformed = torch.ones(1, 24, 24, 24) + + def test_in_is_a_no_op_and_out_is_everywhere(self): + for seg_unit in ("whole", "label", "component"): + with self.subTest(seg_unit=seg_unit): + inside = apply_region_selection( + _Stub(in_seg=1.0, seg_unit=seg_unit, seg_unit_prob=0.5), self.orig, self.transformed, self.empty_seg() + ) + outside = apply_region_selection( + _Stub(out_seg=1.0, seg_unit=seg_unit, seg_unit_prob=0.5), self.orig, self.transformed, self.empty_seg() + ) + self.assertTrue(bool(torch.equal(inside, self.orig)), "'in' applied something with nothing labelled") + self.assertTrue(bool(torch.equal(outside, self.transformed)), "'out' did not cover the whole patch") + + def test_the_unit_count_is_never_zero(self): + """It sizes a draw and a gather; zero is a device-side assert, not a failed batch.""" + for seg_unit in ("whole", "label", "component"): + with self.subTest(seg_unit=seg_unit): + units = region_units(self.empty_seg(), None, UnitSpec(seg_unit, "none", DEFAULT_BG_EM_CLUSTERS)) + self.assertGreaterEqual(units.n_units, 1) + self.assertGreaterEqual(units.n_background, 1) + + +class TestTheDecompositionIsSharedAcrossAPass(SmaugLabTestCase): + """Without this the feature is not trainable: a pipeline runs twenty-odd of these.""" + + def setUp(self): + super().setUp() + self.calls = 0 + self.real = regions._build_units + + def counting(*args, **kwargs): + self.calls += 1 + return self.real(*args, **kwargs) + + regions._build_units = counting + self.addCleanup(setattr, regions, "_build_units", self.real) + + def pipeline(self, *transforms): + return AugmentationSequentialCustom(*transforms, data_keys=["input", "mask"], same_on_batch=True) + + def test_two_transforms_in_one_pass_share_one_decomposition(self): + image = torch.rand(1, 1, 24, 24, 24) + seg = self.multi_label_seg() + pipeline = self.pipeline( + RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit="component"), + RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit="component"), + ) + + first_output(pipeline(image.clone(), seg.clone())) + + self.assertEqual(self.calls, 1, "the decomposition was rebuilt per transform") + + def test_a_different_spec_is_a_different_entry(self): + image = torch.rand(1, 1, 24, 24, 24) + seg = self.multi_label_seg() + pipeline = self.pipeline( + RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit="component"), + RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit="label"), + ) + + first_output(pipeline(image.clone(), seg.clone())) + + self.assertEqual(self.calls, 2) + + def test_a_new_mask_tensor_is_a_new_decomposition(self): + """How invalidation works: a geometric transform returns a different object.""" + image = torch.rand(1, 1, 24, 24, 24) + spec = UnitSpec("component", "none", DEFAULT_BG_EM_CLUSTERS) + region_units(self.multi_label_seg(), image, spec) + region_units(self.multi_label_seg(), image, spec) + self.assertEqual(self.calls, 2) + + def test_the_cache_dies_with_the_mask(self): + seg = self.multi_label_seg() + region_units(seg, torch.rand(1, 1, 24, 24, 24), UnitSpec("label", "none", DEFAULT_BG_EM_CLUSTERS)) + self.assertTrue(hasattr(seg, "_smauglab_region_units")) + del seg + gc.collect() # nothing else holds it, so there is nothing left to leak + + +class TestTheRngStream(SmaugLabTestCase): + """Three properties without which a seeded sweep is not comparable to its baseline.""" + + def scharr_then_a_draw(self, seed=11, **kwargs): + image = torch.rand(*SHAPE) + seg = self.multi_label_seg().expand(BATCH, -1, -1, -1, -1).contiguous() + transform = RandomScharrGPU(p=1.0, **kwargs) + torch.manual_seed(seed) + transform.apply_transform(image.clone(), {"seg": seg}, transform.flags) + return torch.rand(4) + + def test_region_selection_off_spends_nothing(self): + self.assertTrue( + bool(torch.equal(self.scharr_then_a_draw(), self.scharr_then_a_draw(in_seg=0.0, out_seg=0.0, seg_unit="component"))), + "the region parameters moved the stream with both probabilities at zero", + ) + + def test_the_draw_width_does_not_depend_on_the_data(self): + """Two labels or five, the stream has to land in the same place.""" + two = torch.zeros(*SHAPE) + two[:, :, 2:6, 2:6, 2:6] = 1 + two[:, :, 10:14, 10:14, 10:14] = 2 + five = two.clone() + for offset, value in enumerate((3, 4, 5)): + five[:, :, 20, 20, offset] = value + + def stream(seg, seg_unit): + transform = RandomScharrGPU(p=1.0, in_seg=1.0, seg_unit=seg_unit, seg_unit_prob=0.5) + torch.manual_seed(11) + transform.apply_transform(torch.rand(*SHAPE), {"seg": seg.clone()}, transform.flags) + return torch.rand(4) + + for seg_unit in ("label", "component"): + with self.subTest(seg_unit=seg_unit): + self.assertTrue( + bool(torch.equal(stream(two, seg_unit), stream(five, seg_unit))), + f"{seg_unit} spent a number of uniforms that depended on the segmentation", + ) + + def test_keeping_every_unit_spends_nothing(self): + """`seg_unit_prob=1` is how granularity is turned off, so it has to be free.""" + self.assertTrue( + bool( + torch.equal( + self.scharr_then_a_draw(in_seg=1.0, seg_unit="whole", seg_unit_prob=1.0), + self.scharr_then_a_draw(in_seg=1.0, seg_unit="component", seg_unit_prob=1.0), + ) + ) + ) + + def test_the_em_split_is_forked_out_of_the_global_stream(self): + """It spends ~100k uniforms on a cache miss and none on a hit.""" + self.assertTrue( + bool( + torch.equal( + self.scharr_then_a_draw(out_seg=1.0, bg_split="none"), + self.scharr_then_a_draw(out_seg=1.0, bg_split="em"), + ) + ), + "the EM draw leaked into the global stream, so cache state changes the augmentation", + ) + + def test_the_feather_range_does_not_shift_the_stream_when_it_is_off(self): + self.assertTrue( + bool( + torch.equal( + self.scharr_then_a_draw(in_seg=1.0, seg_feather=0.0), + self.scharr_then_a_draw(in_seg=1.0), + ) + ) + ) + + def test_a_degenerate_feather_range_still_spends_its_draw(self): + """Otherwise the stream becomes a function of the range and an A/B is not valid.""" + self.assertTrue( + bool( + torch.equal( + self.scharr_then_a_draw(in_seg=1.0, seg_feather=[1.0, 1.0]), + self.scharr_then_a_draw(in_seg=1.0, seg_feather=[0.5, 2.0]), + ) + ) + ) + + +class TestTheModeIsPerSample(SmaugLabTestCase): + """It was one decision for the whole batch at ten of the eleven call sites. + + So `in_seg=0.5` read as "half the batches" where `p` and `mix_prob` both meant "half + the samples". + """ + + def test_all_four_modes_appear_over_a_batch(self): + torch.manual_seed(0) + seen = set() + for _ in range(16): + modes = regions.choose_region_modes(_Stub(in_seg=0.5, out_seg=0.5), BATCH, torch.device("cpu")) + seen.update(modes.tolist()) + self.assertEqual(seen, {regions.MODE_ALL, regions.MODE_IN, regions.MODE_OUT, regions.MODE_BOTH}) + + def test_two_samples_of_one_batch_can_differ(self): + torch.manual_seed(0) + differed = any( + len(set(regions.choose_region_modes(_Stub(in_seg=0.5, out_seg=0.5), BATCH, torch.device("cpu")).tolist())) > 1 + for _ in range(16) + ) + self.assertTrue(differed, "every sample of the batch got the same mode") + + def test_nothing_is_drawn_when_both_probabilities_are_zero(self): + self.assertIsNone(regions.choose_region_modes(_Stub(), BATCH, torch.device("cpu"))) + + +class TestItSurvivesTheConfigPath(SmaugLabTestCase): + """The constructors are the schema, and JSON hands over lists rather than tuples.""" + + def test_a_section_validates_and_reaches_the_transform(self): + section = { + "RandomScharrGPU": { + "in_seg": 0.5, + "out_seg": 0.25, + "seg_unit": "component", + "seg_unit_prob": 0.3, + "bg_split": "em", + "bg_em_clusters": [2, 5], + "seg_feather": [0.5, 1.5], + } + } + self.assertEqual(validate_section(section, Backend.GPU), []) + + built = {type(t).__name__: t for t in build_gpu_pipeline(section)} + scharr = built["RandomScharrGPU"] + self.assertEqual(scharr.seg_unit, "component") + self.assertEqual(scharr.bg_split, "em") + self.assertEqual(scharr.bg_em_clusters, (2, 5)) + self.assertEqual(scharr.seg_feather, (0.5, 1.5)) + + def test_a_bad_value_in_a_section_fails_the_build(self): + """Not at validation -- that only checks keys -- but before the first batch.""" + section = {"RandomScharrGPU": {"seg_unit": "labl"}} + self.assertEqual(validate_section(section, Backend.GPU), []) + with self.assertRaises(ValueError): + build_gpu_pipeline(section) + + def test_the_renamed_parameters_are_rejected_with_guidance(self): + for section, needle in ( + ({"RandomScharrGPU": {"mix_in_out": True}}, "seg_unit"), + ({"RandomRedistributeSegGPU": {"in_seg": 0.25}}, "confine_prob"), + ): + with self.subTest(section=section): + problems = validate_section(section, Backend.GPU) + self.assertEqual(len(problems), 1) + self.assertIn(needle, problems[0]) + + def test_no_shipped_config_still_sets_the_old_names(self): + import json + + from unit_tests.helpers import configs_dir + + for path in sorted(configs_dir().rglob("*.json")): + payload = json.loads(path.read_text()) + for section in payload.values(): + if not isinstance(section, dict): + continue + for name, params in section.items(): + if not isinstance(params, dict): + continue + with self.subTest(config=path.name, augmentation=name): + self.assertNotIn("mix_in_out", params) + if name in {"RandomRedistributeSegGPU", "RedistributeTransform"}: + self.assertNotIn("in_seg", params) + + +class TestTheValuesAreValidated(SmaugLabTestCase): + """At construction, which is where every value check in this package happens.""" + + def configure(self, **kwargs): + class Owner: + pass + + arguments = {"in_seg": 0.0, "out_seg": 0.0} + arguments.update(kwargs) + configure_regions(Owner(), **arguments) + + def test_a_bad_value_names_the_parameter_and_the_class(self): + for kwargs, needle in ( + ({"in_seg": 1.5}, "in_seg"), + ({"out_seg": -0.1}, "out_seg"), + ({"seg_unit": "labl"}, "seg_unit"), + ({"seg_unit_prob": True}, "seg_unit_prob"), + ({"bg_split": "EM"}, "bg_split"), + ({"bg_em_clusters": 5}, "bg_em_clusters"), + ({"bg_em_clusters": [5, 2]}, "bg_em_clusters"), + ({"seg_feather": -1.0}, "seg_feather"), + ): + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError) as caught: + self.configure(**kwargs) + self.assertIn(needle, str(caught.exception)) + self.assertIn("Owner", str(caught.exception)) + + def test_the_legal_seg_unit_values_are_listed(self): + with self.assertRaises(ValueError) as caught: + self.configure(seg_unit="nope") + for value in regions.SEG_UNITS: + self.assertIn(value, str(caught.exception)) + + def test_a_feather_above_one_is_accepted(self): + """It is a sigma in voxels, not a weight -- the point of `bounds=`.""" + self.configure(seg_feather=[0.5, 4.0]) + + def test_a_mix_alpha_above_one_is_still_rejected(self): + """The same check, with its own bounds. Widening one must not widen the other.""" + from smauglab.transforms.gpu.mixing import normalise_blend_range + + with self.assertRaises(ValueError): + normalise_blend_range([0.5, 4.0]) + + +if __name__ == "__main__": + import unittest + + unittest.main()