batchgenerators 0.25__tar.gz → 0.25.1__tar.gz

{batchgenerators-0.25 → batchgenerators-0.25.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: batchgenerators
-Version: 0.25
+Version: 0.25.1
 Summary: Data augmentation toolkit
 Home-page: https://github.com/MIC-DKFZ/batchgenerators
 Author: Division of Medical Image Computing, German Cancer Research Center AND Applied Computer Vision Lab, Helmholtz Imaging Platform
@@ -8,3 +8,12 @@ Author-email: f.isensee@dkfz-heidelberg.de
 License: Apache License Version 2.0, January 2004
 Keywords: data augmentation,deep learning,image segmentation,image classification,medical image analysis,medical image segmentation
 License-File: LICENSE
+Requires-Dist: pillow>=7.1.2
+Requires-Dist: numpy>=1.10.2
+Requires-Dist: scipy
+Requires-Dist: scikit-image
+Requires-Dist: scikit-learn
+Requires-Dist: future
+Requires-Dist: pandas
+Requires-Dist: unittest2
+Requires-Dist: threadpoolctl

{batchgenerators-0.25 → batchgenerators-0.25.1}/batchgenerators/augmentations/spatial_transformations.py

@@ -16,10 +16,12 @@
 from builtins import range
 
 import numpy as np
+from scipy.ndimage import map_coordinates
 from batchgenerators.augmentations.utils import create_zero_centered_coordinate_mesh, elastic_deform_coordinates, \
     interpolate_img, \
     rotate_coords_2d, rotate_coords_3d, scale_coords, resize_segmentation, resize_multichannel_image, \
-    elastic_deform_coordinates_2
+    elastic_deform_coordinates_2, \
+    get_organ_gradient_field, ignore_anatomy
 from batchgenerators.augmentations.crop_and_pad_augmentations import random_crop as random_crop_aug
 from batchgenerators.augmentations.crop_and_pad_augmentations import center_crop as center_crop_aug
 
@@ -484,3 +486,173 @@ def augment_transpose_axes(data_sample, seg_sample, axes=(0, 1, 2)):
     if seg_sample is not None:
         seg_sample = seg_sample.transpose(*static_axes)
     return data_sample, seg_sample
+
+def augment_anatomy_informed(data, seg,
+                             active_organs, dilation_ranges, directions_of_trans, modalities,
+                             spacing_ratio=0.3125/3.0, blur=32, anisotropy_safety= True,
+                             max_annotation_value=1, replace_value=0):
+    if sum(active_organs) > 0:
+        data_shape = data.shape
+        coords = create_zero_centered_coordinate_mesh(data_shape[-3:])
+
+        for organ_idx, active in reversed(list(enumerate(active_organs))):
+            if active:
+                dil_magnitude = np.random.uniform(low=dilation_ranges[organ_idx][0], high=dilation_ranges[organ_idx][1])
+
+                t, u, v = get_organ_gradient_field(seg == organ_idx + 2,
+                                                   spacing_ratio=spacing_ratio,
+                                                   blur=blur)
+
+                if directions_of_trans[organ_idx][0]:
+                    coords[0, :, :, :] = coords[0, :, :, :] + t * dil_magnitude * spacing_ratio
+                if directions_of_trans[organ_idx][1]:
+                    coords[1, :, :, :] = coords[1, :, :, :] + u * dil_magnitude
+                if directions_of_trans[organ_idx][2]:
+                    coords[2, :, :, :] = coords[2, :, :, :] + v * dil_magnitude
+
+        for d in range(3):
+            ctr = data.shape[d+1] / 2  # !!!
+            coords[d] += ctr - 0.5  # !!!
+
+        if anisotropy_safety:
+            coords[0, 0, :, :][coords[0, 0, :, :] < 0] = 0.0
+            coords[0, 1, :, :][coords[0, 1, :, :] < 0] = 0.0
+            coords[0, -1, :, :][coords[0, -1, :, :] > (data_shape[-3] - 1)] = data_shape[-3] - 1
+            coords[0, -2, :, :][coords[0, -2, :, :] > (data_shape[-3] - 1)] = data_shape[-3] - 1
+
+
+        for modality in modalities:
+            data[modality, :, :, :] = map_coordinates(data[modality, :, :, :], coords, order=1, mode='constant')
+
+        seg[:, :, :] = ignore_anatomy(seg[:, :, :], max_annotation_value=max_annotation_value, replace_value=replace_value)
+        seg[:, :, :] = map_coordinates(seg[:, :, :], coords, order=0, mode='constant')
+
+    else:
+        seg[:, :, :] = ignore_anatomy(seg[:, :, :], max_annotation_value=max_annotation_value, replace_value=replace_value)
+
+    return data, seg
+
+
+def augment_misalign(data, seg, data_size,
+                     im_channels_2_misalign=[0, ],
+                     label_channels_2_misalign=[0, ],
+                     do_squeeze=False,
+                     sq_x=[1.0, 1.0], sq_y=[0.9, 1.1], sq_z=[1.0, 1.0],
+                     p_sq_per_sample=0.1, p_sq_per_dir=1.0,
+                     do_rotation=False,
+                     angle_x=(-0 / 360. * 2 * np.pi, 0 / 360. * 2 * np.pi),
+                     angle_y=(-0 / 360. * 2 * np.pi, 0 / 360. * 2 * np.pi),
+                     angle_z=(-15 / 360. * 2 * np.pi, 15 / 360. * 2 * np.pi),
+                     p_rot_per_sample=0.1, p_rot_per_axis=1.0,
+                     tr_x=[-32, 32], tr_y=[-32, 32], tr_z=[-2, 2],
+                     p_transl_per_sample=0.1, p_transl_per_dir=1.0,
+                     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):
+
+    dim = len(data_size)
+
+    for sample_id in range(data.shape[0]):
+
+        if do_squeeze and np.random.uniform() < p_sq_per_sample:
+            coords = create_zero_centered_coordinate_mesh(data_size)
+            sq = []
+            if dim == 3:
+                if np.random.uniform() <= p_sq_per_dir:
+                    sq.append(np.random.uniform(sq_z[0], sq_z[1]))
+                else:
+                    sq.append(1.0)
+
+            if np.random.uniform() <= p_sq_per_dir:
+                sq.append(np.random.uniform(sq_y[0], sq_y[1]))
+            else:
+                sq.append(1.0)
+
+            if np.random.uniform() <= p_sq_per_dir:
+                sq.append(np.random.uniform(sq_x[0], sq_x[1]))
+            else:
+                sq.append(1.0)
+            coords = scale_coords(coords, sq)
+            for d in range(dim):
+                ctr = data.shape[d + 2] / 2. - 0.5
+                coords[d] += ctr
+
+            for channel_id in range(data.shape[1]):
+                if channel_id in im_channels_2_misalign:
+                    data[sample_id, channel_id] = interpolate_img(data[sample_id, channel_id], coords, order_data,
+                                                                  border_mode_data, cval=border_cval_data)
+            if seg is not None:
+                for channel_id in range(seg.shape[1]):
+                    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)
+
+        if do_rotation and np.random.uniform() < p_rot_per_sample:
+            coords = create_zero_centered_coordinate_mesh(data_size)
+            if np.random.uniform() <= p_rot_per_axis:
+                a_z = np.random.uniform(angle_z[0], angle_z[1])
+            else:
+                a_z = 0
+            if dim == 3:
+                if np.random.uniform() <= p_rot_per_axis:
+                    a_y = np.random.uniform(angle_y[0], angle_y[1])
+                else:
+                    a_y = 0
+                if np.random.uniform() <= p_rot_per_axis:
+                    a_x = np.random.uniform(angle_x[0], angle_x[1])
+                else:
+                    a_x = 0
+                coords = rotate_coords_3d(coords, a_z, a_y, a_x)
+            else:
+                coords = rotate_coords_2d(coords, a_z)
+            for d in range(dim):
+                ctr = data.shape[d + 2] / 2. - 0.5
+                coords[d] += ctr
+
+            for channel_id in range(data.shape[1]):
+                if channel_id in im_channels_2_misalign:
+                    data[sample_id, channel_id] = interpolate_img(data[sample_id, channel_id], coords, order_data,
+                                                                  border_mode_data, cval=border_cval_data)
+            if seg is not None:
+                for channel_id in range(seg.shape[1]):
+                    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)
+
+        if do_transl and np.random.uniform() < p_transl_per_sample:
+            coords = create_zero_centered_coordinate_mesh(data_size)
+            tr = []
+            if dim == 3:
+                if np.random.uniform() <= p_transl_per_dir:
+                    tr.append(np.random.uniform(tr_z[0], tr_z[1]))
+                else:
+                    tr.append(1.0)
+
+            if np.random.uniform() <= p_transl_per_dir:
+                tr.append(np.random.uniform(tr_y[0], tr_y[1]))
+            else:
+                tr.append(1.0)
+
+            if np.random.uniform() <= p_transl_per_dir:
+                tr.append(np.random.uniform(tr_x[0], tr_x[1]))
+            else:
+                tr.append(1.0)
+
+            for d in range(dim):
+                ctr = data.shape[d + 2] / 2. - 0.5 + tr[d]
+                coords[d] += ctr
+
+            for channel_id in range(data.shape[1]):
+                if channel_id in im_channels_2_misalign:
+                    data[sample_id, channel_id] = interpolate_img(data[sample_id, channel_id], coords, order_data,
+                                                                  border_mode_data, cval=border_cval_data)
+            if seg is not None:
+                for channel_id in range(seg.shape[1]):
+                    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)
+    return data, seg

{batchgenerators-0.25 → batchgenerators-0.25.1}/batchgenerators/augmentations/utils.py

@@ -21,6 +21,7 @@ from scipy.ndimage.filters import gaussian_filter, gaussian_gradient_magnitude
 from scipy.ndimage.morphology import grey_dilation
 from skimage.transform import resize
 from scipy.ndimage.measurements import label as lb
+import pandas as pd
 
 
 def generate_elastic_transform_coordinates(shape, alpha, sigma):
@@ -591,13 +592,13 @@ def resize_segmentation(segmentation, new_shape, order=3):
     :return:
     '''
     tpe = segmentation.dtype
-    unique_labels = np.unique(segmentation)
     assert len(segmentation.shape) == len(new_shape), "new shape must have same dimensionality as segmentation"
     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)
@@ -773,3 +774,25 @@ def mask_random_squares(img, square_size, n_squares, n_val, channel_wise_n_val=F
         img = mask_random_square(img, square_size, n_val, channel_wise_n_val=channel_wise_n_val,
                                  square_pos=square_pos)
     return img
+
+def get_organ_gradient_field(organ, spacing_ratio=0.3125/3.0, blur=32):
+    """
+    Calculates the gradient field around the organ segmentations for the anatomy-informed augmentation
+
+    :param organ: binary organ segmentation
+    :param spacing_ratio: ratio of the axial spacing and the slice thickness, needed for the right vector field calculation
+    :param blur: kernel constant
+    """
+    organ_blurred = gaussian_filter(organ.astype(float),
+                                    sigma=(blur * spacing_ratio, blur, blur),
+                                    order=0,
+                                    mode='nearest')
+
+    t, u, v = np.gradient(organ_blurred)
+    t = t * spacing_ratio
+
+    return t, u, v
+
+def ignore_anatomy(segm, max_annotation_value=1, replace_value=0):
+    segm[segm > max_annotation_value] = replace_value
+    return segm

{batchgenerators-0.25 → batchgenerators-0.25.1}/batchgenerators/dataloading/multi_threaded_augmenter.py

@@ -80,7 +80,7 @@ def results_loop(in_queues: List[Queue], out_queue: thrQueue, abort_event: Event
     end_ctr = 0
 
     while True:
-        # if abort_event is set we need to clean up. This is where it hangs sometimes so it makes sense to drain all
+        # if abort_event is set we need to clean up. This is where it hangs sometimes, so it makes sense to drain all
         # the incoming queues and ignore all the errors occuring during this process.
         try:
             if abort_event.is_set():

{batchgenerators-0.25 → batchgenerators-0.25.1}/batchgenerators/dataloading/nondet_multi_threaded_augmenter.py

@@ -208,12 +208,18 @@ class NonDetMultiThreadedAugmenter(object):
             if isinstance(self.generator, DataLoader):
                 self.generator.was_initialized = False
 
-            for i in range(self.num_processes):
-                self._processes.append(Process(target=producer, args=(
-                    self._queue, self.generator, self.transform, i, self.seeds[i], self.abort_event, self.wait_time
-                )))
-                self._processes[-1].daemon = True
-            _ = [i.start() for i in self._processes]
+            if torch is not None:
+                torch_nthreads = torch.get_num_threads()
+                torch.set_num_threads(1)
+            with threadpool_limits(limits=1, user_api=None):
+                for i in range(self.num_processes):
+                    self._processes.append(Process(target=producer, args=(
+                        self._queue, self.generator, self.transform, i, self.seeds[i], self.abort_event, self.wait_time
+                    )))
+                    self._processes[-1].daemon = True
+                _ = [i.start() for i in self._processes]
+            if torch is not None:
+                torch.set_num_threads(torch_nthreads)
 
             if torch is not None and torch.cuda.is_available():
                 gpu = torch.cuda.current_device()

{batchgenerators-0.25 → batchgenerators-0.25.1}/batchgenerators/transforms/noise_transforms.py

@@ -133,6 +133,20 @@ class BlankSquareNoiseTransform(AbstractTransform):
                                                                          self.channel_wise_n_val, self.square_pos)
         return data_dict
 
+class ColorFunctionExtractor:
+    def __init__(self, rectangle_value):
+        self.rectangle_value = rectangle_value
+
+    def __call__(self, x):
+        if np.isscalar(self.rectangle_value):
+            return self.rectangle_value
+        elif callable(self.rectangle_value):
+            return self.rectangle_value(x)
+        elif isinstance(self.rectangle_value, (tuple, list)):
+            return np.random.uniform(*self.rectangle_value)
+        else:
+            raise RuntimeError("unrecognized format for rectangle_value")
+
 
 class BlankRectangleTransform(AbstractTransform):
     def __init__(self, rectangle_size, rectangle_value, num_rectangles, force_square=False, p_per_sample=0.5,
@@ -179,16 +193,7 @@ class BlankRectangleTransform(AbstractTransform):
         self.p_per_sample = p_per_sample
         self.p_per_channel = p_per_channel
         self.apply_to_keys = apply_to_keys
-
-        # intensity value
-        if np.isscalar(rectangle_value):
-            self.color_fn = lambda x: rectangle_value
-        elif callable(rectangle_value):
-            self.color_fn = lambda x: rectangle_value(x)
-        elif isinstance(rectangle_value, (tuple, list)):
-            self.color_fn = lambda x: np.random.uniform(*rectangle_value)
-        else:
-            raise RuntimeError("unrecognized format for rectangle_value")
+        self.color_fn = ColorFunctionExtractor(rectangle_value)
 
     def __call__(self, **data_dict):
         for k in self.apply_to_keys:

{batchgenerators-0.25 → batchgenerators-0.25.1}/batchgenerators/transforms/spatial_transforms.py

@@ -16,8 +16,10 @@
 from batchgenerators.transforms.abstract_transforms import AbstractTransform
 from batchgenerators.augmentations.spatial_transformations import augment_spatial, augment_spatial_2, \
     augment_channel_translation, \
-    augment_mirroring, augment_transpose_axes, augment_zoom, augment_resize, augment_rot90
+    augment_mirroring, augment_transpose_axes, augment_zoom, augment_resize, augment_rot90, \
+    augment_anatomy_informed, augment_misalign
 import numpy as np
+from batchgenerators.augmentations.utils import get_organ_gradient_field
 
 
 class Rot90Transform(AbstractTransform):
@@ -524,3 +526,193 @@ class TransposeAxesTransform(AbstractTransform):
             data_dict[self.label_key] = seg
         return data_dict
 
+class AnatomyInformedTransform(AbstractTransform):
+    """
+    The data augmentation is presented at MICCAI 2023 in the proceedings of 'Anatomy-informed Data Augmentation for enhanced Prostate Cancer Detection'.
+    It simulates the distension or evacuation of the bladder or/and rectal space to mimic typical physiological soft tissue deformations of the prostate
+    and generates unique lesion shapes without altering their label.
+    You can find more information here: https://github.com/MIC-DKFZ/anatomy_informed_DA
+    If you use this augmentation please cite it.
+
+    Args:
+        `dil_ranges`: dilation range per organs
+        `modalities`: on which input channels should the transformation be applied
+        `directions_of_trans`: to which directions should the organs be dilated per organs
+        `p_per_sample`: probability of the transformation per organs
+        `spacing_ratio`: ratio of the transversal plane spacing and the slice thickness, in our case it was 0.3125/3
+        `blur`: Gaussian kernel parameter, we used the value 32 for 0.3125mm transversal plane spacing
+        `anisotropy_safety`: it provides a certain protection against transformation artifacts in 2 slices from the image border
+        `max_annotation_value`: the value that should be still relevant for the main task
+        `replace_value`: segmentation values larger than the `max_annotation_value` will be replaced with
+    """
+    def __init__(self, dil_ranges, modalities, directions_of_trans, p_per_sample,
+                 spacing_ratio=0.3125/3.0, blur=32, anisotropy_safety= True,
+                 max_annotation_value=1, replace_value=0):
+        self.dil_ranges = dil_ranges
+        self.modalities = modalities
+
+        self.directions_of_trans = directions_of_trans
+        self.p_per_sample = p_per_sample
+        self.spacing_ratio = spacing_ratio
+        self.blur = blur
+        self.anisotropy_safety = anisotropy_safety
+
+        self.max_annotation_value = max_annotation_value
+        self.replace_value = replace_value
+
+        self.dim = 3
+
+    def __call__(self, **data_dict):
+
+        data_shape = data_dict['data'].shape
+        if len(data_shape) == 5:
+            self.dim = 3
+
+        active_organs = []
+        for prob in self.p_per_sample:
+            if np.random.uniform() < prob:
+                active_organs.append(1)
+            else:
+                active_organs.append(0)
+
+        for b in range(data_shape[0]):
+            data_dict['data'][b, :, :, :, :], data_dict['seg'][b, 0, :, :, :] = augment_anatomy_informed(data=data_dict['data'][b, :, :, :, :],
+                                                                                                         seg=data_dict['seg'][b, 0, :, :, :],
+                                                                                                         active_organs=active_organs,
+                                                                                                         dilation_ranges=self.dil_ranges,
+                                                                                                         directions_of_trans=self.directions_of_trans,
+                                                                                                         modalities=self.modalities,
+                                                                                                         spacing_ratio=self.spacing_ratio,
+                                                                                                         blur=self.blur,
+                                                                                                         anisotropy_safety=self.anisotropy_safety,
+                                                                                                         max_annotation_value=self.max_annotation_value,
+                                                                                                         replace_value=self.replace_value)
+        return data_dict
+
+
+class MisalignTransform(AbstractTransform):
+    """
+    The misalignment data augmentation is introduced in Nature Scientific reports 2023.
+    It simulates additional misalignments/registration errors between multi-channel (multi-modal, longitudinal)
+    data to make neural networks robust for misalignments.
+    Currently the following transformations are supported, but they can be extended easily:
+    - squeezing/scaling (good approximation for misalignments between the T2w and DWI MRI sequences)
+    - rotation
+    - channel shift/translation
+    You can find more information here: https://github.com/MIC-DKFZ/misalignmnet_DA
+    If you use this augmentation please cite it: https://www.nature.com/articles/s41598-023-46747-z
+    Always double check whether the directions/axes are correct!
+
+    Additional Misalignment-related Args to the Spatial Transforms:
+        `im_channels_2_misalign`: on which image channels should the transformation be applied
+        `label_channels_2_misalign`: on which segmentation channels should the transformation be applied
+        `do_squeeze`: whether misalignment resulted from squeezing is necessary
+        `sq_x, sq_y`, `sq_z`: squeezing/scaling ranges per directions, randomly sampled from interval.
+        `p_sq_per_sample`: probability of the transformation per sample
+        `p_sq_per_dir`: probability of the transformation per direction
+        `do_rotation`: whether misalignment resulted from rotation is necessary
+        `angle_x`, `angle_y`, `angle_z`: rotation angels per axes, randomly sampled from interval.
+        `p_rot_per_sample`: probability of the transformation per sample
+        `p_rot_per_axis`: probability of the transformation per axes
+        `do_transl`: whether misalignment resulted from rotation is necessary
+        `tr_x`, `tr_y`, `tr_z`: shift/translation per directions, randomly sampled from interval.
+        `p_transl_per_sample`: probability of the transformation per sample
+        `p_transl_per_dir`: probability of the transformation per direction
+    """
+
+    def __init__(self, data_key="data", label_key="seg",
+                 im_channels_2_misalign=[0, ], label_channels_2_misalign=[0, ],
+                 do_squeeze=True, sq_x=[1.0, 1.0], sq_y=[0.9, 1.1], sq_z=[1.0, 1.0],
+                 p_sq_per_sample=0.1, p_sq_per_dir=1.0,
+                 do_rotation=True,
+                 angle_x=(-0 / 360. * 2 * np.pi, 0 / 360. * 2 * np.pi),
+                 angle_y=(-0 / 360. * 2 * np.pi, 0 / 360. * 2 * np.pi),
+                 angle_z=(-15 / 360. * 2 * np.pi, 15 / 360. * 2 * np.pi),
+                 p_rot_per_sample=0.1, p_rot_per_axis=1.0,
+                 do_transl=True, tr_x=[-32, 32], tr_y=[-32, 32], tr_z=[-2, 2],
+                 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):
+
+        self.data_key = data_key
+        self.label_key = label_key
+
+        self.im_channels_2_misalign = im_channels_2_misalign
+        self.label_channels_2_misalign = label_channels_2_misalign
+
+        self.do_squeeze = do_squeeze
+        self.sq_x = sq_x
+        self.sq_y = sq_y
+        self.sq_z = sq_z
+        self.p_sq_per_sample = p_sq_per_sample
+        self.p_sq_per_dir = p_sq_per_dir
+
+        self.do_rotation = do_rotation
+        self.angle_x = angle_x
+        self.angle_y = angle_y
+        self.angle_z = angle_z
+        self.p_rot_per_sample = p_rot_per_sample
+        self.p_rot_per_axis = p_rot_per_axis
+
+        self.do_transl = do_transl
+        self.tr_x = tr_x
+        self.tr_y = tr_y
+        self.tr_z = tr_z
+        self.p_transl_per_sample = p_transl_per_sample
+        self.p_transl_per_dir = p_transl_per_dir
+
+        self.order_data = order_data
+        self.order_seg = order_seg
+        self.border_mode_data = border_mode_data
+        self.border_cval_data = border_cval_data
+        self.border_mode_seg = border_mode_seg
+        self.border_cval_seg = border_cval_seg
+
+    def __call__(self, **data_dict):
+        data = data_dict.get(self.data_key)
+        seg = data_dict.get(self.label_key)
+
+        if data.shape[1] < 2:
+            raise ValueError("only support multi-modal images")
+        else:
+            if len(data.shape) == 4:
+                data_size = (data.shape[2], data.shape[3])
+            elif len(data.shape) == 5:
+                data_size = (data.shape[2], data.shape[3], data.shape[4])
+            else:
+                raise ValueError("only support 2D/3D batch data.")
+
+        ret_val = augment_misalign(data, seg, data_size=data_size,
+                                   im_channels_2_misalign=self.im_channels_2_misalign,
+                                   label_channels_2_misalign=self.label_channels_2_misalign,
+                                   do_squeeze=self.do_squeeze,
+                                   sq_x=self.sq_x,
+                                   sq_y=self.sq_y,
+                                   sq_z=self.sq_z,
+                                   p_sq_per_sample=self.p_sq_per_sample,
+                                   p_sq_per_dir=self.p_sq_per_dir,
+                                   do_rotation=self.do_rotation,
+                                   angle_x=self.angle_x,
+                                   angle_y=self.angle_y,
+                                   angle_z=self.angle_z,
+                                   p_rot_per_sample=self.p_rot_per_sample,
+                                   p_rot_per_axis=self.p_rot_per_axis,
+                                   do_transl=self.do_transl,
+                                   tr_x=self.tr_x,
+                                   tr_y=self.tr_y,
+                                   tr_z=self.tr_z,
+                                   p_transl_per_sample=self.p_transl_per_sample,
+                                   p_transl_per_dir=self.p_transl_per_dir,
+                                   order_data=self.order_data,
+                                   border_mode_data=self.border_mode_data,
+                                   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)
+
+        data_dict[self.data_key] = ret_val[0]
+        if seg is not None:
+            data_dict[self.label_key] = ret_val[1]
+
+        return data_dict

{batchgenerators-0.25 → batchgenerators-0.25.1}/batchgenerators/transforms/utility_transforms.py

@@ -19,7 +19,8 @@ from typing import List, Type, Union, Tuple
 import numpy as np
 
 from batchgenerators.augmentations.utils import convert_seg_image_to_one_hot_encoding, \
-    convert_seg_to_bounding_box_coordinates, transpose_channels
+    convert_seg_to_bounding_box_coordinates, transpose_channels, \
+    ignore_anatomy
 from batchgenerators.transforms.abstract_transforms import AbstractTransform
 
 
@@ -276,6 +277,21 @@ class RemoveLabelTransform(AbstractTransform):
         return data_dict
 
 
+class IgnoreAnatomy(AbstractTransform):
+    """
+    Replaces every annotation values larger than the max_annotation_value with the replace_value.
+    This transform is used for the anatomy-informed augmentation scheme to remove all additional anatomical annotations.
+    You can find more information here: https://github.com/MIC-DKFZ/anatomy_informed_DA
+    """
+    def __init__(self, max_annotation_value=1, replace_value=0):
+        self.max_annotation_value = max_annotation_value
+        self.replace_value = replace_value
+
+    def __call__(self, **data_dict):
+        data_dict['seg'] = ignore_anatomy(data_dict['seg'], max_annotation_value=self.max_annotation_value, replace_value=self.replace_value)
+        return data_dict
+
+
 class RenameTransform(AbstractTransform):
     '''
     Saves the value of data_dict[in_key] to data_dict[out_key]. Optionally removes data_dict[in_key] from the dict.