batchgenerators 0.25.3__tar.gz → 0.25.5__tar.gz

{batchgenerators-0.25.3 → batchgenerators-0.25.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batchgenerators
-Version: 0.25.3
+Version: 0.25.5
 Summary: Data augmentation toolkit
 Author-email: "Division of Medical Image Computing, German Cancer Research Center AND Applied Computer Vision Lab, Helmholtz Imaging Platform" <f.isensee@dkfz-heidelberg.de>
 License:                                  Apache License

{batchgenerators-0.25.3 → batchgenerators-0.25.5}/batchgenerators/augmentations/spatial_transformations.py

@@ -45,7 +45,7 @@ def augment_rot90(sample_data, sample_seg, num_rot=(1, 2, 3), axes=(0, 1, 2)):
     return sample_data, sample_seg
 
 
-def augment_resize(sample_data, sample_seg, target_size, order=3, order_seg=1):
+def augment_resize(sample_data, sample_seg, target_size, order=3, order_seg=1, *, seg_tiebreak='nearest'):
     """
     Reshapes data (and seg) to target_size
     :param sample_data: np.ndarray or list/tuple of np.ndarrays, must be (c, x, y(, z))) (if list/tuple then each entry
@@ -71,14 +71,15 @@ def augment_resize(sample_data, sample_seg, target_size, order=3, order_seg=1):
     if sample_seg is not None:
         target_seg = np.ones([sample_seg.shape[0]] + target_size_here)
         for c in range(sample_seg.shape[0]):
-            target_seg[c] = resize_segmentation(sample_seg[c], target_size_here, order_seg)
+            target_seg[c] = resize_segmentation(sample_seg[c], target_size_here, order_seg,
+                                                seg_tiebreak=seg_tiebreak)
     else:
         target_seg = None
 
     return sample_data, target_seg
 
 
-def augment_zoom(sample_data, sample_seg, zoom_factors, order=3, order_seg=1):
+def augment_zoom(sample_data, sample_seg, zoom_factors, order=3, order_seg=1, *, seg_tiebreak='nearest'):
     """
     zooms data (and seg) by factor zoom_factors
     :param sample_data: np.ndarray or list/tuple of np.ndarrays, must be (c, x, y(, z))) (if list/tuple then each entry
@@ -107,7 +108,8 @@ def augment_zoom(sample_data, sample_seg, zoom_factors, order=3, order_seg=1):
     if sample_seg is not None:
         target_seg = np.ones([sample_seg.shape[0]] + target_shape_here)
         for c in range(sample_seg.shape[0]):
-            target_seg[c] = resize_segmentation(sample_seg[c], target_shape_here, order_seg)
+            target_seg[c] = resize_segmentation(sample_seg[c], target_shape_here, order_seg,
+                                                seg_tiebreak=seg_tiebreak)
     else:
         target_seg = None
 
@@ -192,7 +194,8 @@ def augment_spatial(data, seg, patch_size, patch_center_dist_from_border=30,
                     do_scale=True, scale=(0.75, 1.25), border_mode_data='nearest', border_cval_data=0, order_data=3,
                     border_mode_seg='constant', border_cval_seg=0, order_seg=0, random_crop=True, p_el_per_sample=1,
                     p_scale_per_sample=1, p_rot_per_sample=1, independent_scale_for_each_axis=False,
-                    p_rot_per_axis: float = 1, p_independent_scale_per_axis: int = 1):
+                    p_rot_per_axis: float = 1, p_independent_scale_per_axis: int = 1,
+                    *, seg_tiebreak: str = 'nearest'):
     dim = len(patch_size)
     seg_result = None
     if seg is not None:
@@ -277,7 +280,7 @@ def augment_spatial(data, seg, patch_size, patch_center_dist_from_border=30,
                 for channel_id in range(seg.shape[1]):
                     seg_result[sample_id, channel_id] = interpolate_img(seg[sample_id, channel_id], coords, order_seg,
                                                                         border_mode_seg, cval=border_cval_seg,
-                                                                        is_seg=True)
+                                                                        is_seg=True, seg_tiebreak=seg_tiebreak)
         else:
             if seg is None:
                 s = None
@@ -300,7 +303,8 @@ def augment_spatial_2(data, seg, patch_size, patch_center_dist_from_border=30,
                       do_scale=True, scale=(0.75, 1.25), border_mode_data='nearest', border_cval_data=0, order_data=3,
                       border_mode_seg='constant', border_cval_seg=0, order_seg=0, random_crop=True, p_el_per_sample=1,
                       p_scale_per_sample=1, p_rot_per_sample=1, independent_scale_for_each_axis=False,
-                      p_rot_per_axis: float = 1, p_independent_scale_per_axis: float = 1):
+                      p_rot_per_axis: float = 1, p_independent_scale_per_axis: float = 1,
+                      *, seg_tiebreak: str = 'nearest'):
     """
 
     :param data:
@@ -444,7 +448,7 @@ def augment_spatial_2(data, seg, patch_size, patch_center_dist_from_border=30,
                 for channel_id in range(seg.shape[1]):
                     seg_result[sample_id, channel_id] = interpolate_img(seg[sample_id, channel_id], coords, order_seg,
                                                                         border_mode_seg, cval=border_cval_seg,
-                                                                        is_seg=True)
+                                                                        is_seg=True, seg_tiebreak=seg_tiebreak)
         else:
             if seg is None:
                 s = None
@@ -549,7 +553,7 @@ def augment_misalign(data, seg, data_size,
                      do_transl=False,
                      border_mode_data='constant', border_cval_data=0,
                      border_mode_seg='constant', border_cval_seg=0,
-                     order_data=3, order_seg=0):
+                     order_data=3, order_seg=0, *, seg_tiebreak='nearest'):
 
     dim = len(data_size)
 
@@ -587,7 +591,7 @@ def augment_misalign(data, seg, data_size,
                     if channel_id in im_channels_2_misalign:
                         seg[sample_id, channel_id] = interpolate_img(seg[sample_id, channel_id], coords, order_seg,
                                                                      border_mode_seg, cval=border_cval_seg,
-                                                                     is_seg=True)
+                                                                     is_seg=True, seg_tiebreak=seg_tiebreak)
 
         if do_rotation and np.random.uniform() < p_rot_per_sample:
             coords = create_zero_centered_coordinate_mesh(data_size)
@@ -620,7 +624,7 @@ def augment_misalign(data, seg, data_size,
                     if channel_id in im_channels_2_misalign:
                         seg[sample_id, channel_id] = interpolate_img(seg[sample_id, channel_id], coords, order_seg,
                                                                      border_mode_seg, cval=border_cval_seg,
-                                                                     is_seg=True)
+                                                                     is_seg=True, seg_tiebreak=seg_tiebreak)
 
         if do_transl and np.random.uniform() < p_transl_per_sample:
             coords = create_zero_centered_coordinate_mesh(data_size)
@@ -654,5 +658,5 @@ def augment_misalign(data, seg, data_size,
                     if channel_id in label_channels_2_misalign:
                         seg[sample_id, channel_id] = interpolate_img(seg[sample_id, channel_id], coords, order_seg,
                                                                      border_mode_seg, cval=border_cval_seg,
-                                                                     is_seg=True)
+                                                                     is_seg=True, seg_tiebreak=seg_tiebreak)
     return data, seg

{batchgenerators-0.25.3 → batchgenerators-0.25.5}/batchgenerators/augmentations/utils.py

@@ -16,11 +16,15 @@
 import random
 import numpy as np
 from copy import deepcopy
-from scipy.ndimage import map_coordinates, fourier_gaussian
-from scipy.ndimage.filters import gaussian_filter, gaussian_gradient_magnitude
-from scipy.ndimage.morphology import grey_dilation
+from scipy.ndimage import (
+    fourier_gaussian,
+    gaussian_filter,
+    gaussian_gradient_magnitude,
+    grey_dilation,
+    map_coordinates,
+)
+from scipy.ndimage import label as lb
 from skimage.transform import resize
-from scipy.ndimage.measurements import label as lb
 import pandas as pd
 
 
@@ -43,6 +47,24 @@ def create_zero_centered_coordinate_mesh(shape):
     return coords
 
 
+def unique_labels(arr: np.ndarray) -> np.ndarray:
+    """
+    np.unique(arr): the distinct values, ascending.
+
+    Which implementation is cheaper depends on the dtype. pd.unique hashes where np.unique sorts, and
+    on a label map hashing wins for integers (1.3x at int8) but loses for floats (4.5x slower at
+    float32, 2.9x at float64). Segmentations travel through this library as float32 (augment_spatial
+    allocates seg_result as float32), so the float branch is the common one. Either way this is
+    milliseconds against one map_coordinates per label - consistency, not a speedup.
+
+    Order matters to the callers: it fixes the channel order of a one hot encoding, and
+    _resample_seg_by_argmax resolves a strict `>` chain to the first label attaining the maximum.
+    """
+    if arr.dtype.kind in 'iub':
+        return np.sort(pd.unique(arr.ravel()))
+    return np.unique(arr)
+
+
 def convert_seg_image_to_one_hot_encoding(image, classes=None):
     '''
     image must be either (x, y, z) or (x, y)
@@ -50,7 +72,7 @@ def convert_seg_image_to_one_hot_encoding(image, classes=None):
     Example (3D): if input is of shape (x, y, z), the output will ne of shape (n_classes, x, y, z)
     '''
     if classes is None:
-        classes = np.unique(image)
+        classes = unique_labels(image)
     out_image = np.zeros([len(classes)]+list(image.shape), dtype=image.dtype)
     for i, c in enumerate(classes):
         out_image[i][image == c] = 1
@@ -62,7 +84,7 @@ def convert_seg_image_to_one_hot_encoding_batched(image, classes=None):
     same as convert_seg_image_to_one_hot_encoding, but expects image to be (b, x, y, z) or (b, x, y)
     '''
     if classes is None:
-        classes = np.unique(image)
+        classes = unique_labels(image)
     output_shape = [image.shape[0]] + [len(classes)] + list(image.shape[1:])
     out_image = np.zeros(output_shape, dtype=image.dtype)
     for b in range(image.shape[0]):
@@ -142,14 +164,165 @@ def uncenter_coords(coords):
     return coords
 
 
-def interpolate_img(img, coords, order=3, mode='nearest', cval=0.0, is_seg=False):
+SEG_TIEBREAKS = ('nearest', 'lowest', 'highest')
+
+
+def _resample_seg_by_argmax(labels, score_fn, nearest_fn, out_shape, dtype, seg_tiebreak,
+                            empty_value=None):
+    """
+    Resample a segmentation by interpolating each label's indicator and taking the per-voxel argmax over
+    labels. The argmax is accumulated incrementally, so the (n_labels, *out_shape) score stack is never
+    materialized.
+
+    This replaces the `interpolated >= 0.5` assignment these functions used to do, which had two problems:
+
+      * it can leave a voxel unwritten. Where more than two labels meet, no single indicator has to reach
+        0.5 (four labels meeting at a sample point give 0.25 each), and for spline orders >= 2 skimage's
+        clip=True destroys the sum-to-one property outright - on random four-label 2d data at order=3,
+        6.2 % of output voxels had no label at 0.5 or above. Those voxels kept whatever the result array
+        was initialized with (0), which need not be a label that occurs in the input at all.
+      * where several labels do pass 0.5 it assigned in ascending label order, so the largest label won by
+        being written last. That is a decision about label values, not about geometry.
+
+    seg_tiebreak decides the voxels where the top score is shared. That is not a rare case: whenever a
+    sample point falls exactly between two labels - which an even integer resize factor does at every
+    boundary voxel - the two indicators are exactly 0.5.
+
+      'nearest'  take the nearest neighbour label, if it is one of the tied labels. Symmetric in the labels,
+                 and the only option that is a function of the geometry rather than of the label values.
+                 Where three or more labels meet the nearest neighbour can be a label that lost; such a
+                 voxel keeps the smallest of the tied labels, as with 'lowest'. Matches nnU-Net's
+                 resample_torch seg_tiebreak='nearest'.
+      'lowest'   keep the smallest label, i.e. plain argmax semantics.
+      'highest'  keep the largest label. This is the direction the old code leaned, but it is NOT
+                 bit-identical to it: the old code compared each label against 0.5, not against the other
+                 labels, so at orders >= 2 the two disagree wherever clipping broke sum-to-one.
+
+    Only exact ties are resolved by seg_tiebreak; near ties go to the larger score, as they should.
+
+    empty_value, if given, is written where every label scores zero. That happens only outside the image,
+    and only for padding modes that put something constant there, so this is how the caller's cval keeps
+    reaching those voxels now that the result is no longer zero-initialized.
+    """
+    if seg_tiebreak not in SEG_TIEBREAKS:
+        raise ValueError('unknown seg_tiebreak: %s. Must be one of %s' % (seg_tiebreak, str(SEG_TIEBREAKS)))
+    out_shape = tuple(out_shape)
+    if len(labels) == 0:
+        # an empty input: there is nothing to sample, so every output voxel is outside the image. Returning here
+        # also keeps the nearest neighbour sample below from being asked of an empty array, which fails.
+        return np.full(out_shape, 0 if empty_value is None else empty_value, dtype=dtype)
+    if len(labels) == 1 and empty_value is None:
+        # nothing to interpolate, and this also keeps the loop below from having to handle a missing `best`
+        return np.full(out_shape, labels[0], dtype=dtype)
+
+    # Memory: the scores are float64 (that is what the interpolators return, and ties are decided by exact
+    # equality on them), so `best` and the current `score` are the two large buffers and everything else is
+    # kept small. The running winner and the nearest neighbour are held as int8/int16 indices into `labels`
+    # and mapped to label values once at the end, the per-label masks are three reused bool buffers, and the
+    # previous score is released before the next one is allocated. Without that, peak memory was about twice
+    # the bytes per output voxel.
+    idx_dtype = _label_index_dtype(len(labels))
+    win = np.zeros(out_shape, dtype=idx_dtype)
+    best = None
+    better = np.empty(out_shape, dtype=bool)
+    # 'nearest': the nearest neighbour label settles a tie only if it is one of the tied labels. Where three
+    # or more labels meet it need not be - a 2x2x2 block [3,1,1,1,2,2,2,3] sampled at its centre scores
+    # 1 and 2 at 0.375 each and 3 at 0.25, and its nearest neighbour is a 3. So instead of marking tied
+    # voxels, track whether the nearest neighbour's label currently has the top score: a label that takes
+    # the lead strictly decides that afresh, one that draws level can only add to it. Where it holds at the
+    # end the nearest neighbour is the answer (it changes nothing where its label won outright), and where
+    # it does not the argmax winner stands.
+    nearest = seg_tiebreak == 'nearest'
+    if nearest:
+        nn_idx = _nearest_label_index(labels, nearest_fn(), idx_dtype)
+        nn_top = np.empty(out_shape, dtype=bool)
+        is_nn = np.empty(out_shape, dtype=bool)
+        level = np.empty(out_shape, dtype=bool)
+    for i, c in enumerate(labels):
+        score = score_fn(c)
+        if best is None:
+            # the first label is the running winner everywhere until something beats it. Inside the image the
+            # indicators sum to ~1, so this is only ever a starting point; outside it every label scores zero
+            # and empty_value (if the caller gave one) has the final say below.
+            best = score
+            if nearest:
+                np.equal(nn_idx, 0, out=nn_top)
+                np.greater(score, 0, out=level)
+                nn_top &= level
+            del score
+            continue
+        if nearest:
+            np.equal(nn_idx, i, out=is_nn)
+            # `score > 0` guard: two labels that both score zero at a voxel are not tied for the win there,
+            # they are simply both absent, and some other label will claim it. `better` is free until it is
+            # computed below, so it holds that guard meanwhile.
+            np.equal(score, best, out=level)
+            np.greater(score, 0, out=better)
+            level &= better
+            level &= is_nn
+            nn_top |= level
+        # '>=' lets the later (larger) label take ties, '>' leaves them with the earlier (smaller) one
+        (np.greater_equal if seg_tiebreak == 'highest' else np.greater)(score, best, out=better)
+        if nearest:
+            np.copyto(nn_top, is_nn, where=better)
+        np.copyto(win, i, where=better)
+        np.maximum(best, score, out=best)
+        del score
+
+    if nearest:
+        # copyto rather than win[nn_top] = nn_idx[nn_top]: nn_top holds nearly everywhere (also wherever the
+        # nearest neighbour's label won outright), and the indexed form copies all of those out first
+        np.copyto(win, nn_idx, where=nn_top)
+        del nn_idx, nn_top, is_nn, level
+    del better
+    empty = best <= 0 if empty_value is not None else None
+    del best
+    result = labels.astype(dtype, copy=False)[win]
+    if empty is not None:
+        result[empty] = empty_value
+    return result
+
+
+def _label_index_dtype(n_labels):
+    # signed, so that -1 can mean 'not a label'
+    for t in (np.int8, np.int16, np.int32):
+        if n_labels <= np.iinfo(t).max:
+            return t
+    return np.int64
+
+
+def _nearest_label_index(labels, nn, idx_dtype, chunk=1 << 22):
+    """
+    The index into `labels` of each voxel of the nearest neighbour sample `nn` (float64), or -1 where it is not
+    a label (cval outside the image). `labels` is sorted, so this is a searchsorted, verified with the same
+    `==` that compared `nn` with each label before. Chunked, so that its int64 temporaries stay small.
+    """
+    out = np.empty(nn.shape, dtype=idx_dtype)
+    flat_nn, flat_out = nn.reshape(-1), out.reshape(-1)
+    last = len(labels) - 1
+    for s in range(0, flat_nn.size, chunk):
+        v = flat_nn[s:s + chunk]
+        idx = np.minimum(np.searchsorted(labels, v), last)
+        idx[labels[idx] != v] = -1
+        flat_out[s:s + chunk] = idx
+    return out
+
+
+def interpolate_img(img, coords, order=3, mode='nearest', cval=0.0, is_seg=False, *, seg_tiebreak='nearest'):
     if is_seg and order != 0:
-        unique_labels = np.unique(img)
-        result = np.zeros(coords.shape[1:], img.dtype)
-        for i, c in enumerate(unique_labels):
-            res_new = map_coordinates((img == c).astype(float), coords, order=order, mode=mode, cval=cval)
-            result[res_new >= 0.5] = c
-        return result
+        return _resample_seg_by_argmax(
+            unique_labels(img),
+            # the indicators are padded with 0, not cval: outside the image no label is present. Padding them
+            # with cval raised every label's score by the same amount there, so for cval > 0 a voxel entirely
+            # outside was a tie among all labels rather than empty, and empty_value never reached it.
+            lambda c: map_coordinates((img == c).astype(float), coords, order=order, mode=mode, cval=0.0),
+            lambda: map_coordinates(img.astype(float), coords, order=0, mode=mode, cval=cval),
+            coords.shape[1:], img.dtype, seg_tiebreak,
+            # outside the image the caller asked for cval, and with the old zero-initialized result that is
+            # what those voxels happened to get whenever cval was 0. Keep it, explicitly and for any cval.
+            # Converted the way order 0 converts it, so a cval the dtype cannot hold (-1 in uint8) comes out
+            # the same at every order instead of raising here.
+            empty_value=np.array(cval, dtype=float).astype(img.dtype) if mode == 'constant' else None)
     else:
         return map_coordinates(img.astype(float), coords, order=order, mode=mode, cval=cval).astype(img.dtype)
 
@@ -581,29 +754,34 @@ def transpose_channels(batch):
         raise ValueError("wrong dimensions in transpose_channel generator!")
 
 
-def resize_segmentation(segmentation, new_shape, order=3):
+def resize_segmentation(segmentation, new_shape, order=3, *, seg_tiebreak='nearest'):
     '''
-    Resizes a segmentation map. Supports all orders (see skimage documentation). Will transform segmentation map to one
-    hot encoding which is resized and transformed back to a segmentation map.
-    This prevents interpolation artifacts ([0, 0, 2] -> [0, 1, 2])
+    Resizes a segmentation map. Supports all orders (see skimage documentation). Each label's indicator is
+    resized on its own and the per-voxel winner is the label with the largest interpolated indicator, so no
+    intermediate label value can be produced ([0, 0, 2] -> [0, 1, 2]).
     :param segmentation:
     :param new_shape:
     :param order:
+    :param seg_tiebreak: how to settle voxels where two labels share the top score ('nearest', 'lowest' or
+    'highest'). See _resample_seg_by_argmax. 'nearest' is the only geometric choice and is the default.
     :return:
     '''
     tpe = segmentation.dtype
     assert len(segmentation.shape) == len(new_shape), "new shape must have same dimensionality as segmentation"
+    if segmentation.size == 0:
+        # skimage cannot resize an empty array (scipy's zoom divides 0 by 0), at any order. Zeros of new_shape is
+        # what orders >= 1 returned before the argmax, from their zero-initialized result.
+        return np.zeros(new_shape, dtype=tpe)
     if order == 0:
         return resize(segmentation.astype(float), new_shape, order, mode="edge", clip=True, anti_aliasing=False).astype(tpe)
     else:
-        reshaped = np.zeros(new_shape, dtype=segmentation.dtype)
-
-        unique_labels = np.sort(pd.unique(segmentation.ravel()))
-        for i, c in enumerate(unique_labels):
-            mask = segmentation == c
-            reshaped_multihot = resize(mask.astype(float), new_shape, order, mode="edge", clip=True, anti_aliasing=False)
-            reshaped[reshaped_multihot >= 0.5] = c
-        return reshaped
+        return _resample_seg_by_argmax(
+            unique_labels(segmentation),
+            lambda c: resize((segmentation == c).astype(float), new_shape, order, mode="edge", clip=True,
+                             anti_aliasing=False),
+            lambda: resize(segmentation.astype(float), new_shape, 0, mode="edge", clip=True,
+                           anti_aliasing=False),
+            new_shape, tpe, seg_tiebreak)
 
 
 def resize_multichannel_image(multichannel_image, new_shape, order=3):
@@ -616,11 +794,18 @@ def resize_multichannel_image(multichannel_image, new_shape, order=3):
     :return:
     '''
     tpe = multichannel_image.dtype
+    is_int = np.issubdtype(tpe, np.integer)
     new_shp = [multichannel_image.shape[0]] + list(new_shape)
-    result = np.zeros(new_shp, dtype=multichannel_image.dtype)
+    # Integer input is accumulated in float and rounded at the end. Writing the interpolated float
+    # straight into an integer array truncates towards zero, which shifts every voxel that did not land
+    # exactly on a sample point down by up to one intensity step. Floating point input keeps its own
+    # dtype, so nothing about that path changes and the buffer does not grow.
+    result = np.zeros(new_shp, dtype=float if is_int else tpe)
     for i in range(multichannel_image.shape[0]):
         result[i] = resize(multichannel_image[i].astype(float), new_shape, order, clip=True, anti_aliasing=False)
-    return result.astype(tpe)
+    if is_int:
+        np.rint(result, out=result)
+    return result.astype(tpe, copy=False)
 
 
 def get_range_val(value, rnd_type="uniform"):

{batchgenerators-0.25.3 → batchgenerators-0.25.5}/batchgenerators/transforms/spatial_transforms.py

@@ -61,7 +61,7 @@ class Rot90Transform(AbstractTransform):
 
 class ZoomTransform(AbstractTransform):
     def __init__(self, zoom_factors=1, order=3, order_seg=1, concatenate_list=False, data_key="data",
-                 label_key="seg"):
+                 label_key="seg", *, seg_tiebreak='nearest'):
         """
         Zooms 'data' (and 'seg') by zoom_factors
         :param zoom_factors: int or list/tuple of int
@@ -81,6 +81,7 @@ class ZoomTransform(AbstractTransform):
         self.data_key = data_key
         self.label_key = label_key
         self.order = order
+        self.seg_tiebreak = seg_tiebreak
         self.zoom_factors = zoom_factors
 
     def __call__(self, **data_dict):
@@ -105,7 +106,8 @@ class ZoomTransform(AbstractTransform):
             sample_seg = None
             if seg is not None:
                 sample_seg = seg[b]
-            res_data, res_seg = augment_zoom(data[b], sample_seg, self.zoom_factors, self.order, self.order_seg)
+            res_data, res_seg = augment_zoom(data[b], sample_seg, self.zoom_factors, self.order, self.order_seg,
+                                             seg_tiebreak=self.seg_tiebreak)
             results.append((res_data, res_seg))
 
         if concatenate:
@@ -123,7 +125,7 @@ class ZoomTransform(AbstractTransform):
 class ResizeTransform(AbstractTransform):
 
     def __init__(self, target_size, order=3, order_seg=1, concatenate_list=False, data_key="data",
-                 label_key="seg"):
+                 label_key="seg", *, seg_tiebreak='nearest'):
         """
         Reshapes 'data' (and 'seg') to target_size
         :param target_size: int or list/tuple of int
@@ -143,6 +145,7 @@ class ResizeTransform(AbstractTransform):
         self.data_key = data_key
         self.label_key = label_key
         self.order = order
+        self.seg_tiebreak = seg_tiebreak
         self.target_size = target_size
 
     def __call__(self, **data_dict):
@@ -167,7 +170,8 @@ class ResizeTransform(AbstractTransform):
             sample_seg = None
             if seg is not None:
                 sample_seg = seg[b]
-            res_data, res_seg = augment_resize(data[b], sample_seg, self.target_size, self.order, self.order_seg)
+            res_data, res_seg = augment_resize(data[b], sample_seg, self.target_size, self.order, self.order_seg,
+                                               seg_tiebreak=self.seg_tiebreak)
             results.append((res_data, res_seg))
 
         if concatenate:
@@ -303,7 +307,8 @@ class SpatialTransform(AbstractTransform):
                  do_scale=True, scale=(0.75, 1.25), border_mode_data='nearest', border_cval_data=0, order_data=3,
                  border_mode_seg='constant', border_cval_seg=0, order_seg=0, random_crop=True, data_key="data",
                  label_key="seg", p_el_per_sample=1, p_scale_per_sample=1, p_rot_per_sample=1,
-                 independent_scale_for_each_axis=False, p_rot_per_axis:float=1, p_independent_scale_per_axis: int=1):
+                 independent_scale_for_each_axis=False, p_rot_per_axis:float=1, p_independent_scale_per_axis: int=1,
+                 *, seg_tiebreak: str='nearest'):
         self.independent_scale_for_each_axis = independent_scale_for_each_axis
         self.p_rot_per_sample = p_rot_per_sample
         self.p_scale_per_sample = p_scale_per_sample
@@ -330,6 +335,7 @@ class SpatialTransform(AbstractTransform):
         self.random_crop = random_crop
         self.p_rot_per_axis = p_rot_per_axis
         self.p_independent_scale_per_axis = p_independent_scale_per_axis
+        self.seg_tiebreak = seg_tiebreak
 
     def __call__(self, **data_dict):
         data = data_dict.get(self.data_key)
@@ -358,7 +364,8 @@ class SpatialTransform(AbstractTransform):
                                   p_rot_per_sample=self.p_rot_per_sample,
                                   independent_scale_for_each_axis=self.independent_scale_for_each_axis,
                                   p_rot_per_axis=self.p_rot_per_axis, 
-                                  p_independent_scale_per_axis=self.p_independent_scale_per_axis)
+                                  p_independent_scale_per_axis=self.p_independent_scale_per_axis,
+                                  seg_tiebreak=self.seg_tiebreak)
         data_dict[self.data_key] = ret_val[0]
         if seg is not None:
             data_dict[self.label_key] = ret_val[1]
@@ -419,7 +426,8 @@ class SpatialTransform_2(AbstractTransform):
                  do_scale=True, scale=(0.75, 1.25), border_mode_data='nearest', border_cval_data=0, order_data=3,
                  border_mode_seg='constant', border_cval_seg=0, order_seg=0, random_crop=True, data_key="data",
                  label_key="seg", p_el_per_sample=1, p_scale_per_sample=1, p_rot_per_sample=1,
-                 independent_scale_for_each_axis=False, p_rot_per_axis:float=1, p_independent_scale_per_axis: float=1):
+                 independent_scale_for_each_axis=False, p_rot_per_axis:float=1, p_independent_scale_per_axis: float=1,
+                 *, seg_tiebreak: str='nearest'):
         self.p_rot_per_sample = p_rot_per_sample
         self.p_scale_per_sample = p_scale_per_sample
         self.p_el_per_sample = p_el_per_sample
@@ -445,6 +453,7 @@ class SpatialTransform_2(AbstractTransform):
         self.p_independent_scale_per_axis = p_independent_scale_per_axis
         self.independent_scale_for_each_axis = independent_scale_for_each_axis
         self.p_rot_per_axis = p_rot_per_axis
+        self.seg_tiebreak = seg_tiebreak
 
     def __call__(self, **data_dict):
         data = data_dict.get(self.data_key)
@@ -473,7 +482,8 @@ class SpatialTransform_2(AbstractTransform):
                                     p_rot_per_sample=self.p_rot_per_sample,
                                   independent_scale_for_each_axis=self.independent_scale_for_each_axis,
                                   p_rot_per_axis=self.p_rot_per_axis,
-                                  p_independent_scale_per_axis=self.p_independent_scale_per_axis)
+                                  p_independent_scale_per_axis=self.p_independent_scale_per_axis,
+                                  seg_tiebreak=self.seg_tiebreak)
 
         data_dict[self.data_key] = ret_val[0]
         if seg is not None:
@@ -633,8 +643,9 @@ class MisalignTransform(AbstractTransform):
                  p_transl_per_sample=0.1, p_transl_per_dir=1.0,
                  border_mode_data='constant', border_cval_data=0,
                  border_mode_seg='constant', border_cval_seg=0,
-                 order_data=3, order_seg=0):
+                 order_data=3, order_seg=0, *, seg_tiebreak='nearest'):
 
+        self.seg_tiebreak = seg_tiebreak
         self.data_key = data_key
         self.label_key = label_key
 
@@ -709,7 +720,8 @@ class MisalignTransform(AbstractTransform):
                                    border_cval_data=self.border_cval_data,
                                    order_seg=self.order_seg,
                                    border_mode_seg=self.border_mode_seg,
-                                   border_cval_seg=self.border_cval_seg)
+                                   border_cval_seg=self.border_cval_seg,
+                                   seg_tiebreak=self.seg_tiebreak)
 
         data_dict[self.data_key] = ret_val[0]
         if seg is not None:

{batchgenerators-0.25.3 → batchgenerators-0.25.5}/batchgenerators.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batchgenerators
-Version: 0.25.3
+Version: 0.25.5
 Summary: Data augmentation toolkit
 Author-email: "Division of Medical Image Computing, German Cancer Research Center AND Applied Computer Vision Lab, Helmholtz Imaging Platform" <f.isensee@dkfz-heidelberg.de>
 License:                                  Apache License

{batchgenerators-0.25.3 → batchgenerators-0.25.5}/batchgenerators.egg-info/SOURCES.txt

@@ -58,4 +58,5 @@ tests/test_normalizations.py
 tests/test_random_crop.py
 tests/test_resample_augmentations.py
 tests/test_sanity.py
+tests/test_seg_resampling.py
 tests/test_spatial_transformations.py

{batchgenerators-0.25.3 → batchgenerators-0.25.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "batchgenerators"
-version = "0.25.3"
+version = "0.25.5"
 description = "Data augmentation toolkit"
 readme = "README.md"
 license = { file = "LICENSE" }

batchgenerators-0.25.5/tests/test_seg_resampling.py

@@ -0,0 +1,282 @@
+import unittest
+
+import numpy as np
+from skimage.transform import resize
+
+from batchgenerators.augmentations.utils import (interpolate_img, resize_multichannel_image,
+                                                resize_segmentation)
+
+
+def nearest(seg, new_shape):
+    return resize(seg.astype(float), new_shape, 0, mode="edge", clip=True, anti_aliasing=False).astype(seg.dtype)
+
+
+def label_scores(seg, new_shape, order):
+    labels = np.sort(np.unique(seg))
+    return labels, np.stack([resize((seg == c).astype(float), new_shape, order, mode="edge", clip=True,
+                                    anti_aliasing=False) for c in labels])
+
+
+class TestResizeSegmentation(unittest.TestCase):
+    """
+    resize_segmentation and interpolate_img resize a segmentation by interpolating each label's indicator and
+    picking a winner per voxel. They used to assign with `interpolated >= 0.5` into a zero-initialized result,
+    which could leave a voxel unwritten (it then kept a 0 that need not be a label of the input) and, where
+    several labels passed 0.5, handed the voxel to whichever label was written last, i.e. the largest.
+    """
+
+    def test_no_label_is_invented(self):
+        """Four labels meeting at an output centre score 0.25 each, so the old rule wrote nothing there."""
+        seg = np.array([[1, 2, 1, 1],
+                        [3, 4, 1, 1],
+                        [1, 1, 1, 1],
+                        [1, 1, 1, 1]], dtype=np.uint8)
+        labels, scores = label_scores(seg, (2, 2), 1)
+        self.assertLess(scores[:, 0, 0].max(), 0.5, 'this test is only meaningful if no label reaches 0.5')
+        for tiebreak in ('nearest', 'lowest', 'highest'):
+            out = resize_segmentation(seg, (2, 2), order=1, seg_tiebreak=tiebreak)
+            self.assertTrue(set(np.unique(out).tolist()) <= set(np.unique(seg).tolist()),
+                            f'{np.unique(out)} contains a label that is not in the input ({tiebreak})')
+
+    def test_no_label_is_invented_randomized(self):
+        """At order 3 skimage's clip=True breaks sum-to-one outright, so the old rule missed whole voxels."""
+        rng = np.random.RandomState(3)
+        missed_by_the_old_rule = 0
+        for _ in range(60):
+            shape = tuple(rng.randint(6, 14) for _ in range(2))
+            seg = rng.randint(1, 5, shape).astype(np.uint8)  # labels 1..4, the input never contains 0
+            new_shape = tuple(max(int(round(s / f)), 2)
+                              for s, f in zip(shape, rng.choice([1.5, 2., 2.5, 3.], 2)))
+            labels, scores = label_scores(seg, new_shape, 3)
+            missed_by_the_old_rule += int(((scores >= 0.5).sum(0) == 0).sum())
+            for tiebreak in ('nearest', 'lowest', 'highest'):
+                out = resize_segmentation(seg, new_shape, order=3, seg_tiebreak=tiebreak)
+                self.assertTrue(np.isin(out, labels).all(),
+                                f'{np.unique(out)} vs input labels {labels} ({tiebreak})')
+        self.assertGreater(missed_by_the_old_rule, 0, 'the old rule should have left voxels unwritten here')
+
+    def test_nearest_tiebreak_is_symmetric_in_the_label_values(self):
+        """
+        An even integer factor puts an output centre exactly between two labels at every boundary voxel, and
+        an exact tie carries no information about which label belongs there. 'nearest' answers it from the
+        geometry, so relabelling moves 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 = np.full((8, 2), lo, dtype=np.uint8)
+            forward[5:] = hi                                  # boundary lands on the output centre x = 4.5
+            swapped = np.full((8, 2), hi, dtype=np.uint8)
+            swapped[5:] = lo
+            remap = {lo: hi, hi: lo}
+
+            a = resize_segmentation(forward, (4, 2), order=1, seg_tiebreak='nearest')
+            b = resize_segmentation(swapped, (4, 2), order=1, seg_tiebreak='nearest')
+            self.assertTrue(np.array_equal(np.vectorize(remap.get)(a), b),
+                            f"'nearest' is not symmetric for ({lo}, {hi}): {a[:, 0]} vs {b[:, 0]}")
+
+            # the tie is at output index 2; 'lowest' and 'highest' take it by label value
+            self.assertEqual(resize_segmentation(forward, (4, 2), 1, seg_tiebreak='lowest')[2, 0], lo)
+            self.assertEqual(resize_segmentation(forward, (4, 2), 1, seg_tiebreak='highest')[2, 0], hi)
+
+    def test_nearest_tiebreak_matches_a_nearest_neighbour_resize(self):
+        seg = np.full((8, 2), 1, dtype=np.uint8)
+        seg[5:] = 2
+        self.assertTrue(np.array_equal(resize_segmentation(seg, (4, 2), order=1, seg_tiebreak='nearest'),
+                                       nearest(seg, (4, 2))))
+
+    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: halving this 2x2x2
+        block scores 1 and 2 at 0.375 each and 3 at 0.25, and the nearest neighbour of the centre is a 3.
+        """
+        block = np.array([3, 1, 1, 1, 2, 2, 2, 3], dtype=np.int16).reshape(2, 2, 2)
+        seg = np.tile(block, (4, 4, 4))
+        labels, scores = label_scores(seg, (4, 4, 4), 1)
+        self.assertTrue(np.allclose(scores[:, 0, 0, 0], [0.375, 0.375, 0.25]))
+        self.assertEqual(int(nearest(seg, (4, 4, 4))[0, 0, 0]), 3, 'the nearest neighbour has to be the loser')
+        for tiebreak, expected in (('nearest', 1), ('lowest', 1), ('highest', 2)):
+            out = resize_segmentation(seg, (4, 4, 4), order=1, seg_tiebreak=tiebreak)
+            self.assertTrue((out == expected).all(), f'{np.unique(out)} ({tiebreak})')
+
+    def test_winner_has_the_top_score_randomized(self):
+        """whatever settles a tie, the label a voxel gets must be one that attains the maximum score there"""
+        rng = np.random.RandomState(4)
+        for _ in range(40):
+            shape = tuple(rng.randint(4, 9) for _ in range(3))
+            seg = rng.randint(0, rng.randint(2, 6), shape).astype(np.int16)
+            new_shape = tuple(max(s // 2, 1) for s in shape)
+            labels, scores = label_scores(seg, new_shape, 1)
+            top = scores.max(0)
+            for tiebreak in ('nearest', 'lowest', 'highest'):
+                out = resize_segmentation(seg, new_shape, order=1, seg_tiebreak=tiebreak)
+                own = np.take_along_axis(scores, np.searchsorted(labels, out)[None], 0)[0]
+                self.assertTrue((own == top).all(), f'{int((own != top).sum())} voxels ({tiebreak})')
+
+    def test_block_aligned_roundtrip_is_exact(self):
+        """00001111 -> 0011 -> 00001111: down and up share the block grid, so nothing may move."""
+        seg = np.array([0, 0, 0, 0, 1, 1, 1, 1], dtype=np.uint8)[:, None].repeat(2, axis=1)
+        for tiebreak in ('nearest', 'lowest', 'highest'):
+            down = resize_segmentation(seg, (4, 2), order=1, seg_tiebreak=tiebreak)
+            back = resize_segmentation(down, (8, 2), order=1, seg_tiebreak=tiebreak)
+            self.assertTrue(np.array_equal(back, seg), f'tiebreak={tiebreak}')
+
+    def test_single_label_is_passed_through(self):
+        seg = np.full((8, 6), 3, dtype=np.uint8)
+        out = resize_segmentation(seg, (4, 3), order=3)
+        self.assertTrue(np.array_equal(out, np.full((4, 3), 3, dtype=np.uint8)))
+
+    def test_order_zero_is_untouched(self):
+        rng = np.random.RandomState(0)
+        seg = rng.randint(0, 5, (9, 7)).astype(np.uint8)
+        self.assertTrue(np.array_equal(resize_segmentation(seg, (5, 4), order=0), nearest(seg, (5, 4))))
+
+    def test_dtype_is_preserved(self):
+        for dtype in (np.uint8, np.int16, np.int32):
+            seg = np.array([1, 1, 1, 1, 2, 2, 2, 2], dtype=dtype)[:, None].repeat(2, axis=1)
+            self.assertEqual(resize_segmentation(seg, (4, 2), order=1).dtype, dtype)
+
+    def test_unknown_tiebreak_is_rejected(self):
+        seg = np.zeros((8, 8), dtype=np.uint8)
+        seg[4:] = 1
+        with self.assertRaises(ValueError):
+            resize_segmentation(seg, (4, 4), order=1, seg_tiebreak='bogus')
+
+    def test_empty_input_gives_zeros(self):
+        """
+        An empty segmentation came back as zeros of new_shape (the zero-initialized result) at orders >= 1. The
+        argmax broke that: it takes the nearest neighbour sample up front, which skimage cannot compute for an
+        empty array (scipy's zoom divides 0 by 0), and 'lowest'/'highest' indexed an empty label list. Order 0
+        failed the same way even before.
+        """
+        for order in (0, 1, 3):
+            for tiebreak in ('nearest', 'lowest', 'highest'):
+                for in_shape, out_shape in (((0, 4, 4), (0, 8, 8)), ((0, 4, 4), (2, 8, 8)), ((0, 5), (0, 10))):
+                    out = resize_segmentation(np.zeros(in_shape, dtype=np.int16), out_shape, order=order,
+                                              seg_tiebreak=tiebreak)
+                    msg = f'{in_shape} -> {out_shape}, order {order}, {tiebreak}'
+                    self.assertEqual(out.shape, out_shape, msg)
+                    self.assertEqual(out.dtype, np.int16, msg)
+                    self.assertFalse(out.any(), msg)
+                out = resize_segmentation(np.ones((4, 4, 4), dtype=np.int16), (0, 8, 8), order=order,
+                                          seg_tiebreak=tiebreak)
+                self.assertEqual(out.shape, (0, 8, 8))
+
+
+class TestInterpolateImg(unittest.TestCase):
+    """interpolate_img is the same rule on the augment_spatial path, and had the same two problems."""
+
+    @staticmethod
+    def _coords(xs, ys):
+        return np.stack(np.meshgrid(np.array(xs, dtype=float), np.array(ys, dtype=float), indexing='ij'))
+
+    def test_no_label_is_invented(self):
+        img = np.array([[1, 2, 1, 1],
+                        [3, 4, 1, 1],
+                        [1, 1, 1, 1],
+                        [1, 1, 1, 1]], dtype=np.uint8)
+        coords = self._coords([0.5, 2.5], [0.5, 2.5])
+        for tiebreak in ('nearest', 'lowest', 'highest'):
+            out = interpolate_img(img, coords, order=1, is_seg=True, seg_tiebreak=tiebreak)
+            self.assertTrue(set(np.unique(out).tolist()) <= set(np.unique(img).tolist()),
+                            f'{np.unique(out)} contains a label that is not in the input ({tiebreak})')
+
+    def test_nearest_tiebreak_is_symmetric_in_the_label_values(self):
+        coords = self._coords([4.5], [0.0])
+        for lo, hi in ((1, 2), (2, 7)):
+            forward = np.full((8, 2), lo, dtype=np.uint8)
+            forward[5:] = hi
+            swapped = np.full((8, 2), hi, dtype=np.uint8)
+            swapped[5:] = lo
+            a = interpolate_img(forward, coords, order=1, is_seg=True, seg_tiebreak='nearest')
+            b = interpolate_img(swapped, coords, order=1, is_seg=True, seg_tiebreak='nearest')
+            self.assertEqual({lo: hi, hi: lo}[int(a.ravel()[0])], int(b.ravel()[0]))
+            self.assertEqual(int(interpolate_img(forward, coords, 1, is_seg=True, seg_tiebreak='lowest').ravel()[0]), lo)
+            self.assertEqual(int(interpolate_img(forward, coords, 1, is_seg=True, seg_tiebreak='highest').ravel()[0]), hi)
+
+    def test_constant_padding_still_uses_cval_outside_the_image(self):
+        """
+        The old code zero-initialized the result, so voxels sampled outside the image happened to come back
+        as 0 - right whenever cval was 0, and only then. Now that the result carries a real label from the
+        start, cval has to be written there deliberately.
+        """
+        img = np.array([[1, 1, 2, 2]] * 4, dtype=np.uint8)  # the input does not contain 0
+        coords = self._coords([-5., 1.5, 9.], [2.5])        # outside, inside (clear of the 1/2 tie), outside
+        for cval in (0.0, 7.0):
+            for tiebreak in ('nearest', 'lowest', 'highest'):
+                out = interpolate_img(img, coords, order=1, mode='constant', cval=cval, is_seg=True,
+                                      seg_tiebreak=tiebreak).ravel()
+                self.assertEqual(int(out[0]), int(cval), tiebreak)
+                self.assertEqual(int(out[2]), int(cval), tiebreak)
+                self.assertEqual(int(out[1]), 2, tiebreak)
+        # a cval the dtype cannot hold comes out as order 0 converts it, instead of raising
+        for order in (0, 1, 3):
+            out = interpolate_img(img, coords, order=order, mode='constant', cval=-1, is_seg=True).ravel()
+            self.assertEqual(int(out[0]), int(np.array(-1.).astype(np.uint8)), f'order {order}')
+        # padding modes with no 'outside' must be left alone
+        self.assertTrue((interpolate_img(img, coords, order=1, mode='nearest', is_seg=True) == 2).all())
+
+    def test_empty_input_gives_the_outside_value(self):
+        """With no image to sample from every point is outside it: cval for 'constant', 0 otherwise."""
+        img = np.zeros((0, 4), dtype=np.uint8)
+        coords = self._coords([0.5, 1.5], [0.5, 2.5])
+        for order in (1, 3):
+            for tiebreak in ('nearest', 'lowest', 'highest'):
+                out = interpolate_img(img, coords, order=order, mode='nearest', is_seg=True, seg_tiebreak=tiebreak)
+                self.assertEqual(out.shape, (2, 2))
+                self.assertEqual(out.dtype, np.uint8)
+                self.assertFalse(out.any(), f'order {order}, {tiebreak}')
+                for cval in (0., 7.):
+                    out = interpolate_img(img, coords, order=order, mode='constant', cval=cval, is_seg=True,
+                                          seg_tiebreak=tiebreak)
+                    self.assertTrue((out == cval).all(), f'order {order}, {tiebreak}, cval {cval}')
+
+    def test_order_zero_and_non_seg_are_untouched(self):
+        img = np.arange(16, dtype=np.uint8).reshape(4, 4)
+        coords = self._coords([0.5, 2.5], [0.5, 2.5])
+        self.assertTrue(np.array_equal(interpolate_img(img, coords, order=0, is_seg=True),
+                                       interpolate_img(img, coords, order=0, is_seg=False)))
+
+
+class TestResizeMultichannelImage(unittest.TestCase):
+    """
+    The image counterpart in the same module. It has no per-label reduction and therefore none of the
+    tie-break problem, but it used to allocate its output with the *input* dtype and assign the
+    interpolated float into it, which truncates towards zero for integer input.
+    """
+
+    def test_integer_input_is_rounded_not_truncated(self):
+        img = (np.arange(8, dtype=np.uint8) * 30)[None, :, None].repeat(2, axis=2)
+        exact = resize(img[0].astype(float), (4, 2), 3, clip=True, anti_aliasing=False)
+        got = resize_multichannel_image(img, (4, 2), 3)
+        self.assertTrue(np.array_equal(got[0], np.rint(exact).astype(np.uint8)))
+        self.assertFalse(np.array_equal(got[0], np.trunc(exact).astype(np.uint8)),
+                         'this test needs a case where rounding and truncation disagree')
+
+    def test_integer_input_has_no_downward_bias(self):
+        """truncation costs about half an intensity step per voxel, everywhere, in one direction"""
+        rng = np.random.RandomState(0)
+        bias = []
+        for _ in range(50):
+            a = rng.randint(0, 256, (1, 16, 16)).astype(np.uint8)
+            exact = resize(a[0].astype(float), (9, 9), 3, clip=True, anti_aliasing=False)
+            bias.append((resize_multichannel_image(a, (9, 9), 3)[0].astype(float) - exact).mean())
+        self.assertLess(abs(float(np.mean(bias))), 0.02)
+
+    def test_float_input_is_unchanged(self):
+        """the only path our pipelines use, and it must stay bit-identical"""
+        rng = np.random.RandomState(0)
+        for dtype in (np.float32, np.float64):
+            a = rng.rand(2, 16, 16).astype(dtype)
+            expected = np.stack([resize(a[i].astype(float), (9, 9), 3, clip=True,
+                                        anti_aliasing=False).astype(dtype) for i in range(2)])
+            self.assertTrue(np.array_equal(resize_multichannel_image(a, (9, 9), 3), expected))
+
+    def test_dtype_is_preserved(self):
+        rng = np.random.RandomState(0)
+        for dtype in (np.uint8, np.int16, np.int32, np.float32, np.float64):
+            a = (rng.rand(1, 8, 8) * 100).astype(dtype)
+            self.assertEqual(resize_multichannel_image(a, (4, 4), 3).dtype, dtype)
+
+
+if __name__ == '__main__':
+    unittest.main()