batchgeneratorsv2 0.3.5__tar.gz → 0.3.6__tar.gz

{batchgeneratorsv2-0.3.5 → batchgeneratorsv2-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batchgeneratorsv2
-Version: 0.3.5
+Version: 0.3.6
 Summary: Batchgenerators but better
 Author: Helmholtz Imaging Applied Computer Vision Lab
 Author-email: Fabian Isensee <f.isensee@dkfz-heidelberg.de>

{batchgeneratorsv2-0.3.5 → batchgeneratorsv2-0.3.6}/batchgeneratorsv2/transforms/spatial/spatial.py

@@ -13,6 +13,9 @@ from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
 from batchgeneratorsv2.transforms.utils.cropping import crop_tensor
 
 
+SEG_TIEBREAKS = ('nearest', 'lowest', 'highest')
+
+
 class SpatialTransform(BasicTransform):
     def __init__(self,
                  patch_size: Tuple[int, ...],
@@ -28,7 +31,7 @@ class SpatialTransform(BasicTransform):
                  p_scaling: float = 0,
                  scaling: RandomScalar = (0.7, 1.3),
                  p_synchronize_scaling_across_axes: float = 0,
-                 bg_style_seg_sampling: bool = True,
+                 bg_style_seg_sampling: bool = False,
                  mode_seg: str = 'bilinear',
                  border_mode_seg: str = "zeros",
                  center_deformation: bool = True,
@@ -36,17 +39,29 @@ class SpatialTransform(BasicTransform):
                  padding_mode_image: str = "zeros",
                  padding_value_seg: float = 0,
                  padding_value_image: float = 0,
-                 align_corners: bool = False
+                 align_corners: bool = False,
+                 *,
+                 seg_tiebreak: str = 'nearest'
                  ):
         """
         magnitude must be given in pixels!
         deformation scale is given as a paercentage of the edge length
 
         padding_mode_image: see torch grid_sample documentation. This currently applies to image and regression target
-        because both call self._apply_to_image. Can be "zeros", "constant", "reflection", "border"
+        because both call self._apply_to_image. Can be "zeros", "constant", "reflection", "border". "constant" pads
+        with padding_value_image and resamples the padded image, exactly as "zeros" does with 0.
 
         border_mode_seg: can be "zeros", "constant", "reflection", "border". padding values are only considered for
-        the corresponding "constant" modes.
+        the corresponding "constant" modes. "zeros" and "constant" pad the segmentation with a label (0 or
+        padding_value_seg) and resample the padded segmentation, so the image ends at the outer edge of its
+        border pixels: a sample point up to half a pixel past the last pixel centre still gets that pixel's
+        label, and one further out gets the padding label wherever the padding outweighs the image.
+
+        seg_tiebreak: how to settle voxels where two labels share the top interpolated score, which is what
+        every boundary voxel becomes when a sample point falls exactly between two labels. 'nearest' takes
+        the nearest neighbour label (symmetric in the labels, and the only option that is a function of the
+        geometry rather than of the label values), 'lowest' keeps the smallest label (plain argmax, matches
+        nnU-Net's resample_torch), 'highest' keeps the largest. Only used when mode_seg != 'nearest'.
         """
         super().__init__()
         self.patch_size = patch_size
@@ -65,6 +80,9 @@ class SpatialTransform(BasicTransform):
         self.p_synchronize_scaling_across_axes = p_synchronize_scaling_across_axes
         self.p_synchronize_def_scale_across_axes = p_synchronize_def_scale_across_axes
         self.bg_style_seg_sampling = bg_style_seg_sampling
+        if seg_tiebreak not in SEG_TIEBREAKS:
+            raise ValueError(f'unknown seg_tiebreak: {seg_tiebreak}. Must be one of {SEG_TIEBREAKS}')
+        self.seg_tiebreak = seg_tiebreak
         self.mode_seg = mode_seg
         self.border_mode_seg = border_mode_seg
         self.center_deformation = center_deformation
@@ -103,19 +121,24 @@ class SpatialTransform(BasicTransform):
             return padding_mode
         raise RuntimeError(f'Unknown pad mode: {padding_mode}')
 
-    @staticmethod
-    def _requires_constant_padding_fixup(padding_mode: str, padding_value: float) -> bool:
-        return padding_mode == 'constant' and padding_value != 0
-
-    def _compute_out_of_bounds_mask(self, grid: torch.Tensor, spatial_shape: Tuple[int, ...]) -> torch.Tensor:
+    def _kernel_bound(self, spatial_shape: Tuple[int, ...]) -> np.ndarray:
+        """
+        Per grid axis, the largest |grid coordinate| at which the self.mode_seg kernel still draws on pixels
+        inside the image only (bilinear reaches one pixel, bicubic two). Axis k of the grid is spatial axis
+        dim - 1 - k, which is how grid_sample reads it (x is the last spatial axis).
+        """
+        margin = {'bilinear': 0, 'bicubic': 1}[self.mode_seg]
+        size = np.array(spatial_shape[::-1], dtype=float)
+        # pixel coordinate x < margin or x > size - 1 - margin, as a bound on |grid coordinate| (symmetric)
         if self.align_corners:
-            lo = grid.new_tensor(-1.)
-            hi = grid.new_tensor(1.)
-        else:
-            size = grid.new_tensor(spatial_shape)
-            lo = -1 + 1 / size
-            hi = 1 - 1 / size
-        return ((grid < lo) | (grid > hi)).any(dim=-1)
+            return 1 - 2 * margin / np.maximum(size - 1, 1e-12)
+        return 1 - (2 * margin + 1) / size
+
+    def _may_reach_outside(self, grid: torch.Tensor, spatial_shape: Tuple[int, ...]) -> bool:
+        """False guarantees that no kernel reaches outside. One contiguous pass over the grid, under 1 ms at 128^3."""
+        bound = self._kernel_bound(spatial_shape).min()
+        gmin, gmax = torch.aminmax(grid)
+        return not (-gmin.item() <= bound and gmax.item() <= bound)
 
     def get_parameters(self, **data_dict) -> dict:
         dim = data_dict['image'].ndim - 1
@@ -260,16 +283,19 @@ class SpatialTransform(BasicTransform):
             )
 
         grid = params['grid']
+        # 'constant' pads with padding_value_image and resamples the padded image, like 'zeros' does with 0 and
+        # like the segmentation path does with its padding label. grid_sample only pads with zeros, and it is
+        # linear in its input, so sampling (img - v) and adding v back is exactly constant padding with v.
+        shift = float(self.padding_value_image) if self.padding_mode_image == 'constant' else 0.
         result = grid_sample(
-            img[None],
+            (img - shift)[None] if shift != 0 else img[None],
             grid[None],
             mode=self.mode_image,
             padding_mode=self._get_grid_sample_padding_mode(self.padding_mode_image),
             align_corners=self.align_corners,
         )[0]
-        if self._requires_constant_padding_fixup(self.padding_mode_image, self.padding_value_image):
-            out_of_bounds_mask = self._compute_out_of_bounds_mask(grid, img.shape[1:])
-            result.masked_fill_(out_of_bounds_mask.unsqueeze(0), self.padding_value_image)
+        if shift != 0:
+            result += shift
         return result
 
     def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
@@ -289,81 +315,176 @@ class SpatialTransform(BasicTransform):
 
         grid = params['grid']
         grid_sample_padding_mode = self._get_grid_sample_padding_mode(self.border_mode_seg)
+        # 'zeros' and 'constant' pad the segmentation with a label: 0, or padding_value_seg. Outside the image is
+        # then treated exactly like pixels of that label - pad first, then resample - so the image ends at the
+        # edge of its outermost pixels, not at their centres. A sample point half a pixel past the last pixel
+        # centre still belongs to that pixel, and an augmented patch reaching past the crop does not cut the
+        # segmentation short. That is also what the image path does with zeros padding. 'border' and
+        # 'reflection' never leave the image.
+        pads = grid_sample_padding_mode == 'zeros'
+        pad_label = self.padding_value_seg if self.border_mode_seg == 'constant' else 0
 
         if self.mode_seg == 'nearest':
-            result_seg = grid_sample(
-                segmentation[None].float(),
+            # grid_sample only pads with zeros; shifting by the padding label makes that padding the label
+            shift = float(pad_label) if pads else 0.
+            result_seg = (grid_sample(
+                segmentation[None].float() - shift,
                 grid[None],
                 mode=self.mode_seg,
                 padding_mode=grid_sample_padding_mode,
                 align_corners=self.align_corners
-            )[0].to(segmentation.dtype)
+            )[0] + shift).to(segmentation.dtype)
         else:
+            # Both branches compute the same thing: the per-voxel argmax over the interpolated one-hot label
+            # channels, with self.seg_tiebreak settling voxels where the top score is shared. They are
+            # bit-identical (tests/test_seg_sampling.py pins that) and differ only in time and memory:
+            #
+            #   bg_style_seg_sampling=True   materializes the (n_labels, *patch_size) float16 score stack and
+            #                                argmaxes it in one go.
+            #   bg_style_seg_sampling=False  keeps a single running score volume and folds the argmax into the
+            #                                label loop.
+            #
+            # False is the default because it is the better of the two on both counts. Measured on a 64^3
+            # patch: 0.51x the time at 3 labels, 0.68x at 12, 0.71x at 40, against 0.5 MB of score buffer
+            # instead of 1.5 / 6 / 20 MB. The stack only looked cheaper while the old code was allowed to
+            # skip the background label and to answer the two-label case with a single grid_sample - the
+            # first is not compatible with an argmax, and the second now sits in front of both branches,
+            # where it makes them identical at two labels. True is kept only so callers that pass it
+            # explicitly keep working; there is no longer a reason to choose it.
+            #
+            # The float16 score representation and scale_factor are load-bearing for that bit-identity: float16
+            # has a spacing of 0.5 at this magnitude, so it collapses near ties into exact ones, and both
+            # branches have to collapse the same ones.
+            #
+            # Neither branch uses the old `interpolated >= 0.5` assignment. That rule could leave a voxel
+            # unwritten - where three or more labels meet, no single indicator has to reach 0.5 - and the voxel
+            # then kept the zero result_seg was initialized with, which need not be a label of the input at all.
+            scale_factor = 1000
+            half = 0.5 * scale_factor
             result_seg = torch.zeros((segmentation.shape[0], *self.patch_size), dtype=segmentation.dtype)
-            if self.bg_style_seg_sampling:
-                for c in range(segmentation.shape[0]):
-                    labels = torch.from_numpy(np.sort(pd.unique(segmentation[c].numpy().ravel())))
-                    # result_seg is zero-initialized, so when the lowest label is the 0 background we never have to
-                    # write it: any voxel it would claim is already 0, and being the lowest label it can only be
-                    # overwritten by (never overwrite) the foreground labels processed after it. Skipping it saves a
-                    # full grid_sample per channel. Output is bit-identical.
-                    bg_is_zero = bool(labels[0] == 0)
-                    # if we only have 2 labels then we can save compute time
-                    if len(labels) == 2:
-                        out = grid_sample(
-                            ((segmentation[c] == labels[1]).float())[None, None],
+            nn_seg = None  # nearest neighbour sample, computed on first use
+
+            def _nearest():
+                # the nearest neighbour of the padded segmentation, i.e. what mode_seg='nearest' returns. Where it
+                # is not one of the tied labels it is not used (see use_nn), so it can never invent a label.
+                nonlocal nn_seg
+                if nn_seg is None:
+                    shift = float(pad_label) if pads else 0.
+                    nn_seg = (grid_sample(
+                        segmentation[None].float() - shift,
+                        grid[None],
+                        mode='nearest',
+                        padding_mode=grid_sample_padding_mode,
+                        align_corners=self.align_corners
+                    )[0] + shift).to(segmentation.dtype)
+                return nn_seg
+
+            # The padding label enters the argmax only for patches whose kernel reaches outside the image, which
+            # _may_reach_outside establishes cheaply; the usual patch stays inside, and its scores are computed
+            # exactly as they would be without padding.
+            may_reach_outside = pads and self._may_reach_outside(grid, segmentation.shape[1:])
+            # every label's indicator is built into the same input-sized float buffer (torch.eq writes 0. / 1.
+            # straight into it) rather than allocating a bool and a float tensor afresh per label: one pass over
+            # the input per label instead of two. Allocated on first use, so single-label channels never need it
+            ind_buf = []
+
+            for c in range(segmentation.shape[0]):
+                labels = torch.from_numpy(np.sort(pd.unique(segmentation[c].numpy().ravel())))
+                if may_reach_outside:
+                    labels = torch.unique(torch.cat((labels, torch.tensor([pad_label]).to(labels.dtype))))
+                if len(labels) == 1:
+                    result_seg[c] = labels[0].item()
+                    continue
+
+                def _score(u):
+                    if not ind_buf:
+                        ind_buf.append(torch.empty(segmentation.shape[1:], dtype=torch.float32))
+                    ind = ind_buf[0]
+                    torch.eq(segmentation[c], u, out=ind)
+                    if may_reach_outside and u == pad_label:
+                        # the padding label's indicator is 1 outside the image: sample (ind - 1) with zeros padding
+                        # and add the 1 back
+                        return grid_sample(
+                            ind.sub_(1).mul_(scale_factor)[None, None],
                             grid[None],
                             mode=self.mode_seg,
                             padding_mode=grid_sample_padding_mode,
                             align_corners=self.align_corners
-                        )[0][0] >= 0.5
-                        result_seg[c][out] = labels[1]
-                        if not bg_is_zero:
-                            result_seg[c][~out] = labels[0]
+                        )[0][0].add_(scale_factor).to(torch.float16)
+                    return grid_sample(
+                        ind.mul_(scale_factor)[None, None],
+                        grid[None],
+                        mode=self.mode_seg,
+                        padding_mode=grid_sample_padding_mode,
+                        align_corners=self.align_corners
+                    )[0][0].to(torch.float16)
+
+                if len(labels) == 2:
+                    # Shared by both branches, so they cannot drift apart here. grid_sample is linear in its
+                    # input and does not clip, and with the padding label counted the indicators partition the
+                    # padded image, so the two scores sum to scale_factor everywhere: the second label is the
+                    # argmax exactly where its own score passes half of it, and one grid_sample answers the whole
+                    # channel instead of two.
+                    hi = _score(labels[1])
+                    better = (hi >= half) if self.seg_tiebreak == 'highest' else (hi > half)
+                    # with two labels the nearest neighbour is always one of them, so every tie is its to settle
+                    use_nn = (hi == half) if self.seg_tiebreak == 'nearest' else None
+                    result_seg[c] = labels[0].item()
+                    result_seg[c][better] = labels[1]
+                elif self.bg_style_seg_sampling:
+                    scores = torch.empty((len(labels), *self.patch_size), dtype=torch.float16)
+                    for i, u in enumerate(labels):
+                        scores[i] = _score(u)
+                    if self.seg_tiebreak == 'highest':
+                        # argmax reports the first maximum, so reverse the stack to get the last one
+                        winner = len(labels) - 1 - scores.flip(0).argmax(0)
                     else:
-                        for u in (labels[1:] if bg_is_zero else labels):
-                            result_seg[c][
-                                grid_sample(
-                                    ((segmentation[c] == u).float())[None, None],
-                                    grid[None],
-                                    mode=self.mode_seg,
-                                    padding_mode=grid_sample_padding_mode,
-                                    align_corners=self.align_corners
-                                )[0][0] >= 0.5] = u
-            else:
-                # Per-voxel argmax over the interpolated one-hot label channels, computed incrementally so the
-                # (num_labels, *patch_size) stack is never materialized: only a running best value and the winning
-                # label (written straight into result_seg) are kept. This equals threshold-then-argmax sampling
-                # because the interpolated channels form a convex combination (they sum to <= scale_factor per
-                # voxel, so at most one label can exceed 0.7 * scale_factor). Load-bearing for bit-identical
-                # results: the running comparison is done in float16, and the strict > keeps the lowest label
-                # index on ties, exactly like torch.argmax(0) over a float16 stack.
-                scale_factor = 1000
-                for c in range(segmentation.shape[0]):
-                    labels = torch.from_numpy(np.sort(pd.unique(segmentation[c].numpy().ravel())))
+                        winner = scores.argmax(0)
+                    result_seg[c] = labels[winner].to(result_seg.dtype)
+                    if self.seg_tiebreak == 'nearest':
+                        # The nearest neighbour settles a tie only if its label is one of the tied ones: where
+                        # three or more labels meet it can be a label that lost.
+                        own = scores.gather(0, torch.searchsorted(labels, _nearest()[c])[None])[0]
+                        use_nn = own == scores.max(0).values
+                        del own
+                    else:
+                        use_nn = None
+                    del scores
+                else:
                     best_val = None
+                    if self.seg_tiebreak == 'nearest':
+                        # The nearest neighbour's label settles a voxel where it has the top score: where it
+                        # won outright that changes nothing, where it is tied it settles the tie. Where three
+                        # or more labels meet it can be a label that lost, and then the argmax winner stands.
+                        # So keep the score of the nearest neighbour's label per voxel and compare it with the
+                        # top score once, at the end. -inf never equals a score: a voxel whose nearest
+                        # neighbour is not among the labels keeps the argmax.
+                        nn_c = _nearest()[c]
+                        nn_score = torch.full(nn_c.shape, float('-inf'), dtype=torch.float16)
+                        is_nn = torch.empty(nn_c.shape, dtype=torch.bool)
+                    else:
+                        nn_c = None
+                    better = torch.empty(tuple(self.patch_size), dtype=torch.bool)
                     for u in labels:
-                        onehot = (segmentation[c] == u).float()
-                        onehot *= scale_factor
-                        cur = grid_sample(
-                            onehot[None, None],
-                            grid[None],
-                            mode=self.mode_seg,
-                            padding_mode=grid_sample_padding_mode,
-                            align_corners=self.align_corners
-                        )[0][0].to(torch.float16)
+                        cur = _score(u)
+                        if nn_c is not None:
+                            torch.eq(nn_c, u, out=is_nn)
+                            torch.where(is_nn, cur, nn_score, out=nn_score)
                         if best_val is None:
                             best_val = cur
-                            if u != 0:  # result_seg is zero-initialized, so filling with 0 would be a no-op
-                                result_seg[c] = u
-                        else:
-                            better = cur > best_val
-                            torch.maximum(best_val, cur, out=best_val)
-                            result_seg[c][better] = u
-
-        if self._requires_constant_padding_fixup(self.border_mode_seg, self.padding_value_seg):
-            out_of_bounds_mask = self._compute_out_of_bounds_mask(grid, segmentation.shape[1:])
-            result_seg.masked_fill_(out_of_bounds_mask.unsqueeze(0), self.padding_value_seg)
+                            result_seg[c] = u.item()
+                            continue
+                        # '>=' lets the later (larger) label take ties, '>' leaves them with the earlier one
+                        (torch.ge if self.seg_tiebreak == 'highest' else torch.gt)(cur, best_val, out=better)
+                        torch.maximum(best_val, cur, out=best_val)
+                        result_seg[c][better] = u
+                    use_nn = (nn_score == best_val) if nn_c is not None else None
+
+                if use_nn is not None and torch.any(use_nn):
+                    # where the nearest neighbour's label won outright this changes nothing. torch.where rather
+                    # than indexing: use_nn holds nearly everywhere, and indexing materializes an index per voxel
+                    torch.where(use_nn, _nearest()[c], result_seg[c], out=result_seg[c])
+
         del grid
         return result_seg.contiguous()
 

{batchgeneratorsv2-0.3.5 → batchgeneratorsv2-0.3.6}/batchgeneratorsv2.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batchgeneratorsv2
-Version: 0.3.5
+Version: 0.3.6
 Summary: Batchgenerators but better
 Author: Helmholtz Imaging Applied Computer Vision Lab
 Author-email: Fabian Isensee <f.isensee@dkfz-heidelberg.de>

{batchgeneratorsv2-0.3.5 → batchgeneratorsv2-0.3.6}/batchgeneratorsv2.egg-info/SOURCES.txt

@@ -58,4 +58,5 @@ batchgeneratorsv2/transforms/utils/nnunet_masking.py
 batchgeneratorsv2/transforms/utils/pseudo2d.py
 batchgeneratorsv2/transforms/utils/random.py
 batchgeneratorsv2/transforms/utils/remove_label.py
-batchgeneratorsv2/transforms/utils/seg_to_regions.py
+batchgeneratorsv2/transforms/utils/seg_to_regions.py
+tests/test_seg_sampling.py

{batchgeneratorsv2-0.3.5 → batchgeneratorsv2-0.3.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "batchgeneratorsv2"
-version = "0.3.5"
+version = "0.3.6"
 requires-python = ">=3.9"
 description = "Batchgenerators but better"
 readme = "readme.md"

batchgeneratorsv2-0.3.6/tests/test_seg_sampling.py

@@ -0,0 +1,383 @@
+import unittest
+
+import numpy as np
+import torch
+from torch.nn.functional import grid_sample
+
+from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform
+
+
+def grid_coord(x, n):
+    """normalized grid coordinate that samples exactly at input coordinate x (align_corners=False)"""
+    return (2.0 * x + 1.0) / n - 1.0
+
+
+def blobs(rng, shape, n_labels):
+    v = sum(np.roll(rng.rand(*shape), rng.randint(0, 5), axis=rng.randint(0, len(shape))) for _ in range(3))
+    q = np.quantile(v, np.linspace(0, 1, n_labels + 1)[1:-1]) if n_labels > 1 else []
+    return np.digitize(v, q).astype(np.int16)
+
+
+def transform(patch_size, tiebreak, bg_style, **kwargs):
+    kwargs.setdefault('random_crop', False)
+    return SpatialTransform(patch_size=patch_size, patch_center_dist_from_border=0,
+                            bg_style_seg_sampling=bg_style, seg_tiebreak=tiebreak, mode_seg='bilinear', **kwargs)
+
+
+def sample_row(row, tiebreak, bg_style):
+    """sample a length-8 label row at x = 0.5, 2.5, 4.5, 6.5, i.e. straight through the 2x downsampling grid"""
+    seg = torch.tensor([[row]], dtype=torch.int16)
+    grid = torch.tensor([[[grid_coord(x, 8), 0.0] for x in (0.5, 2.5, 4.5, 6.5)]])
+    t = transform((1, 4), tiebreak, bg_style)
+    return t._apply_to_segmentation(seg, grid=grid, center_location_in_pixels=None).numpy().ravel()
+
+
+def padded_argmax(seg, grid, mode, pad_label, tiebreak='lowest', margin=8):
+    """
+    The definition of 'zeros' / 'constant' padding for a segmentation, computed independently of
+    SpatialTransform: pad the channel with `margin` pixels of pad_label, move the grid to the padded image,
+    and take the argmax there, where no sample point is outside any more.
+    """
+    shape = seg.shape[1:]
+    padded = torch.nn.functional.pad(seg[0][None].float(), [margin] * (2 * len(shape)), value=float(pad_label))[0]
+    g = grid.clone()
+    for k, n in enumerate(shape[::-1]):  # grid axis k is spatial axis dim - 1 - k
+        x = ((g[..., k] + 1) * n - 1) / 2 + margin
+        g[..., k] = (2 * x + 1) / (n + 2 * margin) - 1
+    labels = torch.unique(padded).to(seg.dtype)
+    scores = torch.stack([grid_sample(((padded == u).float() * 1000)[None, None], g[None], mode=mode,
+                                      padding_mode='border', align_corners=False)[0, 0].to(torch.float16)
+                          for u in labels])
+    if tiebreak == 'highest':
+        return labels[len(labels) - 1 - scores.flip(0).argmax(0)]
+    return labels[scores.argmax(0)]
+
+
+class TestSegSampling(unittest.TestCase):
+    """
+    SpatialTransform resamples a segmentation by interpolating each label's indicator and picking a winner
+    per voxel. Both bg_style_seg_sampling branches must implement the same rule - the per-voxel argmax, with
+    seg_tiebreak settling exact ties - and must never write a label that is not in the input.
+    """
+
+    def test_branches_agree_exactly(self):
+        for i in range(30):
+            rng = np.random.RandomState(1000 + i)
+            seg = blobs(rng, (32, 32), rng.randint(1, 8))
+            if rng.rand() < 0.4:  # label values need not start at zero
+                seg = seg + rng.randint(1, 4)
+            for tiebreak in ('nearest', 'lowest', 'highest'):
+                out = []
+                for bg_style in (True, False):
+                    torch.manual_seed(i)
+                    np.random.seed(i)
+                    t = transform((24, 24), tiebreak, bg_style, random_crop=True, p_elastic_deform=1.0,
+                                  p_rotation=1.0, rotation=(0, 2 * np.pi), p_scaling=1.0, scaling=(0.7, 1.4))
+                    s = torch.from_numpy(seg)[None]
+                    out.append(t(**{'image': s.float(), 'segmentation': s})['segmentation'].numpy())
+                self.assertTrue(np.array_equal(out[0], out[1]),
+                                f'branches disagree on {np.sum(out[0] != out[1])} voxels (i={i}, {tiebreak})')
+
+    def test_no_label_is_invented(self):
+        """
+        The regression that motivated the rewrite. The old rule assigned with `interpolated >= 0.5` into a
+        zero-initialized result, so a voxel that no label claimed kept the 0 - which need not occur in the
+        input. Four labels meeting at a sample point give 0.25 each and no label claims it.
+        """
+        seg = torch.tensor([[[1, 2], [3, 4]]], dtype=torch.int16)
+        grid = torch.tensor([[[0.0, 0.0]]])  # the exact centre of the 2x2 block
+        for bg_style in (True, False):
+            for tiebreak in ('nearest', 'lowest', 'highest'):
+                out = transform((1, 1), tiebreak, bg_style)._apply_to_segmentation(
+                    seg, grid=grid, center_location_in_pixels=None).numpy()
+                self.assertTrue(set(np.unique(out).tolist()) <= {1, 2, 3, 4},
+                                f'produced a label that is not in the input: {np.unique(out)} '
+                                f'(bg_style={bg_style}, {tiebreak})')
+
+    def test_no_label_is_invented_under_augmentation(self):
+        # 'border' padding, so no voxel is outside the image: with zeros padding the outside is label 0 by
+        # definition, and test_zeros_padding_is_label_zero covers that
+        for i in range(20):
+            rng = np.random.RandomState(2000 + i)
+            seg = blobs(rng, (32, 32), rng.randint(2, 8)) + rng.randint(1, 4)  # never contains 0
+            present = set(np.unique(seg).tolist())
+            for bg_style in (True, False):
+                for tiebreak in ('nearest', 'lowest', 'highest'):
+                    torch.manual_seed(i)
+                    np.random.seed(i)
+                    t = transform((24, 24), tiebreak, bg_style, random_crop=True, p_elastic_deform=1.0,
+                                  p_rotation=1.0, rotation=(0, 2 * np.pi), p_scaling=1.0, scaling=(0.7, 1.4),
+                                  border_mode_seg='border')
+                    s = torch.from_numpy(seg)[None]
+                    out = t(**{'image': s.float(), 'segmentation': s})['segmentation'].numpy()
+                    self.assertTrue(set(np.unique(out).tolist()) <= present,
+                                    f'invented {sorted(set(np.unique(out).tolist()) - present)} '
+                                    f'(i={i}, bg_style={bg_style}, {tiebreak})')
+
+    def test_nearest_tiebreak_is_symmetric_in_the_label_values(self):
+        """
+        Sampling exactly between two labels is an exact tie, and it is not a rare case: it is what an even
+        integer downsampling factor does at every boundary voxel. 'nearest' has to answer it from the
+        geometry, so relabelling must move the same decision with the labels. 'lowest' and 'highest' answer
+        it from the label values, so they do not.
+        """
+        for lo, hi in ((1, 2), (2, 7), (3, 4)):
+            forward = [lo] * 5 + [hi] * 3
+            swapped = [hi] * 5 + [lo] * 3
+            remap = {lo: hi, hi: lo}
+            for bg_style in (True, False):
+                a = sample_row(forward, 'nearest', bg_style)
+                b = sample_row(swapped, 'nearest', bg_style)
+                self.assertTrue(np.array_equal([remap[v] for v in a], list(b)),
+                                f"'nearest' is not symmetric for ({lo}, {hi}): {a} vs {b}")
+                # the tied voxel is the third one (x = 4.5, exactly between index 4 and index 5)
+                self.assertEqual(sample_row(forward, 'lowest', bg_style)[2], lo)
+                self.assertEqual(sample_row(forward, 'highest', bg_style)[2], hi)
+
+    def test_nearest_tiebreak_matches_a_nearest_neighbour_sample(self):
+        row = [1] * 5 + [2] * 3
+        seg = torch.tensor([[row]], dtype=torch.int16)
+        grid = torch.tensor([[[grid_coord(x, 8), 0.0] for x in (0.5, 2.5, 4.5, 6.5)]])
+        nn = grid_sample(seg[None].float(), grid[None], mode='nearest', padding_mode='border',
+                         align_corners=False)[0].numpy().ravel()
+        for bg_style in (True, False):
+            self.assertTrue(np.array_equal(sample_row(row, 'nearest', bg_style), nn.astype(np.int16)))
+
+    def test_single_label_is_passed_through(self):
+        for bg_style in (True, False):
+            self.assertTrue(np.array_equal(sample_row([7] * 8, 'nearest', bg_style), np.full(4, 7)))
+
+    def test_3d(self):
+        rng = np.random.RandomState(7)
+        seg = torch.from_numpy(blobs(rng, (20, 20, 20), 5))[None]
+        out = []
+        for bg_style in (True, False):
+            torch.manual_seed(0)
+            np.random.seed(0)
+            t = transform((16, 16, 16), 'nearest', bg_style, random_crop=True, p_rotation=1.0,
+                          rotation=(0, 2 * np.pi), p_scaling=1.0, scaling=(0.7, 1.4))
+            out.append(t(**{'image': seg.float(), 'segmentation': seg})['segmentation'].numpy())
+        self.assertTrue(np.array_equal(out[0], out[1]))
+        self.assertEqual(out[0].shape, (1, 16, 16, 16))
+
+    def test_zeros_padding_is_label_zero(self):
+        """
+        With zeros padding the outside of the image is label 0, whatever labels the channel holds, as it is on
+        the crop path and with mode_seg='nearest'. The argmax alone would hand a point entirely outside to
+        labels[0] (labels[-1] for 'highest'), i.e. to whichever label happens to be the smallest in this
+        channel, so the result would depend on which labels are present. Points only partly outside still
+        take the argmax over the labels they do see.
+        """
+        n = 8
+        # x = -3 and x = 10 lie entirely outside, x = -0.4 and x = 7.4 in the band that is partly outside
+        xs = (-3.0, -0.4, 3.5, 7.4, 10.0)
+        grid = torch.tensor([[[grid_coord(x, n), 0.0] for x in xs]])
+        for row in ([1] * 4 + [2] * 4, [1] * 3 + [2] * 2 + [3] * 3, [5] * 4 + [9] * 4):
+            seg = torch.tensor([[row]], dtype=torch.int16)
+            for border_mode in ('zeros', 'constant'):
+                for bg_style in (True, False):
+                    for tiebreak in ('nearest', 'lowest', 'highest'):
+                        t = transform((1, len(xs)), tiebreak, bg_style, border_mode_seg=border_mode,
+                                      padding_value_seg=0)
+                        out = t._apply_to_segmentation(seg, grid=grid, center_location_in_pixels=None).numpy().ravel()
+                        msg = f'{row} {border_mode} bg_style={bg_style} {tiebreak}: {out}'
+                        self.assertEqual((out[0], out[4]), (0, 0), msg)
+                        self.assertEqual((out[1], out[3]), (row[0], row[-1]), msg)
+            t = SpatialTransform(patch_size=(1, len(xs)), patch_center_dist_from_border=0, random_crop=False,
+                                 mode_seg='nearest')
+            out = t._apply_to_segmentation(seg, grid=grid, center_location_in_pixels=None).numpy().ravel()
+            self.assertEqual((out[0], out[4]), (0, 0))
+
+    def test_two_labels_match_the_full_argmax(self):
+        """
+        Two labels take a shortcut (one grid_sample, `hi > half`) that is only the argmax where the two scores
+        sum to scale_factor. Check it against an explicit argmax with sample points inside, across and beyond
+        the border, for every padding mode and for bicubic, whose kernel reaches one pixel further. For 'zeros'
+        and 'constant' the reference pads the segmentation for real and resamples the padded image.
+        """
+        rng = np.random.RandomState(3)
+        for i in range(10):
+            seg = torch.from_numpy(blobs(rng, (16, 16), 2) + rng.randint(0, 3))[None]
+            labels = torch.unique(seg)
+            # sample points everywhere, including a margin outside the image
+            grid = torch.from_numpy(rng.uniform(-1.3, 1.3, size=(12, 12, 2))).float()
+            for padding, value in (('zeros', 0), ('constant', -1), ('border', 0), ('reflection', 0)):
+                for mode in ('bilinear', 'bicubic'):
+                    t = transform((12, 12), 'lowest', False, border_mode_seg=padding, padding_value_seg=value)
+                    t.mode_seg = mode
+                    out = t._apply_to_segmentation(seg, grid=grid, center_location_in_pixels=None)[0]
+                    if padding in ('zeros', 'constant'):
+                        ref = padded_argmax(seg, grid, mode, value).to(out.dtype)
+                    else:
+                        scores = torch.stack([
+                            grid_sample(((seg[0] == u).float() * 1000)[None, None], grid[None], mode=mode,
+                                        padding_mode=padding, align_corners=False)[0, 0].to(torch.float16)
+                            for u in labels])
+                        ref = labels[scores.argmax(0)].to(out.dtype)
+                    self.assertTrue(torch.equal(out, ref),
+                                    f'{int((out != ref).sum())} voxels differ (i={i}, {padding}, {mode})')
+
+    def test_nearest_tiebreak_only_picks_among_the_tied_labels(self):
+        """
+        Where three or more labels meet, the nearest neighbour can be a label that lost: the centre of this
+        2x2x2 block scores 1 and 2 at 0.375 each and 3 at 0.25, and its nearest neighbour is a 3.
+        """
+        block = torch.tensor([3, 1, 1, 1, 2, 2, 2, 3], dtype=torch.int16).reshape(2, 2, 2)
+        seg = block.repeat(4, 4, 4)[None]
+        centres = [grid_coord(2 * i + 0.5, 8) for i in range(4)]
+        grid = torch.stack(torch.meshgrid(*[torch.tensor(centres)] * 3, indexing='ij')[::-1], dim=-1).float()
+        for bg_style in (True, False):
+            for tiebreak, expected in (('nearest', 1), ('lowest', 1), ('highest', 2)):
+                t = transform((4, 4, 4), tiebreak, bg_style)
+                out = t._apply_to_segmentation(seg, grid=grid, center_location_in_pixels=None)
+                self.assertTrue((out == expected).all(), f'{torch.unique(out).tolist()} ({tiebreak}, {bg_style})')
+
+    def test_winner_has_the_top_score(self):
+        """
+        Whatever settles a tie, the label a voxel gets must attain the maximum (float16) score there. Random
+        labels sampled at the centres of 2x2x2 blocks put three or more labels into an exact tie at many voxels;
+        random rotations and scalings of smooth blobs cover the general case.
+        """
+        rng = np.random.RandomState(8)
+        centres = torch.tensor([grid_coord(2 * i + 0.5, 12) for i in range(6)])
+        block_grid = torch.stack(torch.meshgrid(*[centres] * 3, indexing='ij')[::-1], dim=-1).float()
+        for i in range(12):
+            if i % 2:
+                seg = torch.from_numpy(rng.randint(0, rng.randint(3, 6), (12, 12, 12)).astype(np.int16))[None]
+            else:
+                seg = torch.from_numpy(blobs(rng, (12, 12, 12), rng.randint(3, 7)))[None]
+            labels = torch.unique(seg)
+            for bg_style in (True, False):
+                for tiebreak in ('nearest', 'lowest', 'highest'):
+                    if i % 2:
+                        t = transform((6, 6, 6), tiebreak, bg_style, border_mode_seg='border')
+                        grid = block_grid
+                    else:
+                        torch.manual_seed(i)
+                        np.random.seed(i)
+                        t = transform((10, 10, 10), tiebreak, bg_style, border_mode_seg='border', p_rotation=1.0,
+                                      rotation=(0, 2 * np.pi), p_scaling=1.0, scaling=(0.7, 1.4))
+                        grid = t.get_parameters(image=seg.float(), segmentation=seg)['grid']
+                    out = t._apply_to_segmentation(seg, grid=grid, center_location_in_pixels=None)[0]
+                    scores = torch.stack([grid_sample(((seg[0] == u).float() * 1000)[None, None], grid[None],
+                                                      mode='bilinear', padding_mode='border',
+                                                      align_corners=False)[0, 0].to(torch.float16)
+                                          for u in labels])
+                    own = scores.gather(0, torch.searchsorted(labels, out)[None])[0]
+                    self.assertTrue(torch.equal(own, scores.max(0).values),
+                                    f'{int((own != scores.max(0).values).sum())} voxels ({tiebreak}, {bg_style}, i={i})')
+
+    def test_may_reach_outside(self):
+        """
+        _may_reach_outside decides whether the padding label has to enter the argmax at all. False must
+        guarantee that no kernel draws on a pixel outside the image, on a non-cubic image so that pairing a
+        grid axis with the wrong spatial axis shows; bicubic reaches one pixel further than bilinear.
+        """
+        rng = np.random.RandomState(5)
+        for shape in ((6, 11), (5, 9, 14)):
+            for reach in (0.3, 0.6, 0.8, 0.9, 1.0, 1.2):
+                grid = torch.from_numpy(rng.uniform(-reach, reach, size=(7, 8, 9)[:len(shape)] + (len(shape),))).float()
+                for align_corners in (False, True):
+                    for mode in ('bilinear', 'bicubic'):
+                        t = SpatialTransform(patch_size=grid.shape[:-1], patch_center_dist_from_border=0,
+                                             random_crop=False, mode_seg=mode, align_corners=align_corners)
+                        if len(shape) == 3 and mode == 'bicubic':
+                            continue  # grid_sample has no 3D bicubic
+                        # a kernel that draws on the padding shows up as weight missing from a sampled ones image
+                        w_in = grid_sample(torch.ones((1, 1, *shape)), grid[None], mode=mode, padding_mode='zeros',
+                                           align_corners=align_corners)[0, 0]
+                        reaches = bool((torch.abs(w_in - 1) > 1e-5).any())
+                        if not t._may_reach_outside(grid, shape):
+                            self.assertFalse(reaches, f'{shape} {reach} {align_corners} {mode}')
+                        if reach > 1.0:  # past +-1 is outside even with align_corners, where +-1 is a pixel centre
+                            self.assertTrue(t._may_reach_outside(grid, shape))
+
+    def test_padding_is_a_label(self):
+        """
+        'zeros' and 'constant' pad the segmentation with a label and resample the padded segmentation. On a row
+        of four pixels, x = 3.3 lies inside the last pixel (past its centre), x = 3.7 is 0.2 px beyond the image
+        with 70 % of a linear kernel in the padding, and x = 5 is entirely outside. Single-label channels are
+        padded too.
+        """
+        xs = [3.3, 3.7, 5.0, -3.0]
+        grid = torch.tensor([[[grid_coord(x, 4), grid_coord(0.5, 2)] for x in xs]]).float()
+        for row in ([1, 1, 2, 2], [1, 3, 2, 2], [2, 2, 2, 2]):
+            seg = torch.tensor([row] * 2, dtype=torch.int16)[None]
+            for border_mode, pad in (('zeros', 0), ('constant', -1)):
+                for mode_seg in ('nearest', 'bilinear'):
+                    for bg_style in (True, False):
+                        t = SpatialTransform((1, 4), 0, False, mode_seg=mode_seg, border_mode_seg=border_mode,
+                                             padding_value_seg=pad, bg_style_seg_sampling=bg_style)
+                        out = t._apply_to_segmentation(seg, grid=grid, center_location_in_pixels=None).ravel()
+                        self.assertEqual(out.tolist(), [2, pad, pad, pad], f'{row} {border_mode} {mode_seg} {bg_style}')
+
+    def test_transform_padding_equals_dataloader_padding(self):
+        """
+        Padding by the transform must be indistinguishable from the segmentation having been padded with that
+        label beforehand, which is what nnU-Net's dataloader does at real image borders: rotated and scaled
+        patches that reach past the crop must not cut the segmentation short.
+        """
+        rng = np.random.RandomState(12)
+        for i in range(8):
+            shape = (10, 12, 14)
+            seg = torch.from_numpy(blobs(rng, shape, rng.randint(2, 5)))[None]
+            grid = torch.from_numpy(rng.uniform(-1.25, 1.25, size=(6, 7, 8, 3))).float()
+            for pad in (0, -1):
+                for bg_style in (True, False):
+                    for tiebreak in ('lowest', 'highest'):
+                        t = transform((6, 7, 8), tiebreak, bg_style, border_mode_seg='constant', padding_value_seg=pad)
+                        out = t._apply_to_segmentation(seg, grid=grid, center_location_in_pixels=None)[0]
+                        ref = padded_argmax(seg, grid, 'bilinear', pad, tiebreak).to(out.dtype)
+                        self.assertTrue(torch.equal(out, ref), f'{int((out != ref).sum())} voxels (i={i}, {pad}, '
+                                                               f'{bg_style}, {tiebreak})')
+
+    def test_constant_image_padding_is_padding(self):
+        """
+        'constant' must be the image padded with padding_value_image and resampled, like 'zeros' with 0: checked
+        against an image that is padded for real, with sample points inside, across and beyond the border.
+        """
+        rng = np.random.RandomState(6)
+        margin = 8
+        for shape in ((6, 11), (5, 9, 14)):
+            img = torch.from_numpy(rng.rand(2, *shape)).float()
+            grid = torch.from_numpy(rng.uniform(-1.3, 1.3, size=(7, 8, 9)[:len(shape)] + (len(shape),))).float()
+            g = grid.clone()
+            for k, n in enumerate(shape[::-1]):  # grid axis k is spatial axis dim - 1 - k
+                g[..., k] = (2 * (((g[..., k] + 1) * n - 1) / 2 + margin) + 1) / (n + 2 * margin) - 1
+            for value in (0., -1., 5.):
+                padded = torch.nn.functional.pad(img, [margin] * (2 * len(shape)), value=value)
+                for mode in ('bilinear', 'nearest') + (('bicubic',) if len(shape) == 2 else ()):
+                    t = SpatialTransform(patch_size=grid.shape[:-1], patch_center_dist_from_border=0, random_crop=False,
+                                         mode_image=mode, padding_mode_image='constant', padding_value_image=value)
+                    out = t._apply_to_image(img, grid=grid, center_location_in_pixels=None)
+                    ref = grid_sample(padded[None], g[None], mode=mode, padding_mode='border', align_corners=False)[0]
+                    self.assertTrue(torch.allclose(out, ref, atol=1e-5), f'{shape} {value} {mode}: '
+                                                                          f'{float((out - ref).abs().max())}')
+            # 'zeros' must not have moved at all
+            t = SpatialTransform(patch_size=grid.shape[:-1], patch_center_dist_from_border=0, random_crop=False)
+            self.assertTrue(torch.equal(t._apply_to_image(img, grid=grid, center_location_in_pixels=None),
+                                        grid_sample(img[None], grid[None], mode='bilinear', padding_mode='zeros',
+                                                    align_corners=False)[0]))
+
+    def test_constant_padding_keeps_the_image(self):
+        """
+        With nnU-Net's border_mode_seg='constant' and a non-zero padding value, sampling a non-cubic
+        segmentation exactly at its own pixel centres must return it unchanged - no padding value inside.
+        """
+        rng = np.random.RandomState(7)
+        for shape in ((6, 20), (4, 9, 24)):
+            seg = torch.from_numpy(blobs(rng, shape, 3))[None]
+            axes = [torch.tensor([grid_coord(x, n) for x in range(n)]) for n in shape]
+            grid = torch.stack(torch.meshgrid(*axes, indexing='ij')[::-1], dim=-1).float()
+            t = transform(shape, 'nearest', False, border_mode_seg='constant', padding_value_seg=-1)
+            out = t._apply_to_segmentation(seg, grid=grid, center_location_in_pixels=None)
+            self.assertTrue(torch.equal(out, seg), f'{int((out != seg).sum())} voxels changed ({shape})')
+
+    def test_unknown_tiebreak_is_rejected(self):
+        with self.assertRaises(ValueError):
+            SpatialTransform(patch_size=(4, 4), patch_center_dist_from_border=0, random_crop=False,
+                             seg_tiebreak='bogus')
+
+
+if __name__ == '__main__':
+    unittest.main()