batchgenerators 0.24__tar.gz → 0.25.1__tar.gz

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

@@ -1,14 +1,19 @@
 Metadata-Version: 2.1
 Name: batchgenerators
-Version: 0.24
+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
 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
-Platform: UNKNOWN
 License-File: LICENSE
-
-UNKNOWN
-
+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.24 → 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
 
@@ -298,7 +300,7 @@ 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: int = 1):
+                      p_rot_per_axis: float = 1, p_independent_scale_per_axis: float = 1):
     """
 
     :param data:
@@ -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.24 → 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.24 → 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():
@@ -89,9 +89,8 @@ def results_loop(in_queues: List[Queue], out_queue: thrQueue, abort_event: Event
             # check if all workers are still alive
             if not all([i.is_alive() for i in worker_list]):
                 abort_event.set()
-                raise RuntimeError("Abort event was set. So someone died and we should end this madness. \nIMPORTANT: "
-                                   "This is not the actual error message! Look further up to see what caused the "
-                                   "error. Please also check whether your RAM was full")
+                raise RuntimeError("One or more background workers are no longer alive. Exiting. Please check the print"
+                                   " statements above for the actual error message")
 
             # if we don't have an item we need to fetch it first. If the queue we want to get it from it empty, try
             # again later
@@ -187,9 +186,8 @@ class MultiThreadedAugmenter(object):
         while item is None:
             if self.abort_event.is_set():
                 self._finish()
-                raise RuntimeError("MultiThreadedAugmenter.abort_event was set, something went wrong. Maybe one of "
-                                   "your workers crashed. This is not the actual error message! Look further up your "
-                                   "stdout to see what caused the error. Please also check whether your RAM was full")
+                raise RuntimeError("One or more background workers are no longer alive. Exiting. Please check the "
+                                   "print statements above for the actual error message")
 
             if not self.pin_memory_queue.empty():
                 item = self.pin_memory_queue.get()

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

@@ -38,39 +38,40 @@ except ImportError:
 def producer(queue: Queue, data_loader, transform, thread_id: int, seed,
              abort_event: Event, wait_time: float = 0.02):
     # the producer will set the abort event if something happens
-    np.random.seed(seed)
-    data_loader.set_thread_id(thread_id)
-    item = None
-
-    try:
-        while True:
+    with threadpool_limits(1, None):
+        np.random.seed(seed)
+        data_loader.set_thread_id(thread_id)
+        item = None
 
-            if abort_event.is_set():
-                return
-            else:
-                if item is None:
-                    item = next(data_loader)
-                    if transform is not None:
-                        item = transform(**item)
+        try:
+            while True:
 
                 if abort_event.is_set():
                     return
-
-                if not queue.full():
-                    queue.put(item)
-                    item = None
                 else:
-                    sleep(wait_time)
+                    if item is None:
+                        item = next(data_loader)
+                        if transform is not None:
+                            item = transform(**item)
 
-    except KeyboardInterrupt:
-        abort_event.set()
-        return
+                    if abort_event.is_set():
+                        return
 
-    except Exception as e:
-        print("Exception in background worker %d:\n" % thread_id, e)
-        traceback.print_exc()
-        abort_event.set()
-        return
+                    if not queue.full():
+                        queue.put(item)
+                        item = None
+                    else:
+                        sleep(wait_time)
+
+        except KeyboardInterrupt:
+            abort_event.set()
+            return
+
+        except Exception as e:
+            print("Exception in background worker %d:\n" % thread_id, e)
+            traceback.print_exc()
+            abort_event.set()
+            return
 
 
 def pin_memory_of_all_eligible_items_in_dict(result_dict):
@@ -99,9 +100,8 @@ def results_loop(in_queue: Queue, out_queue: thrQueue, abort_event: Event,
             # check if all workers are still alive
             if not all([i.is_alive() for i in worker_list]):
                 abort_event.set()
-                raise RuntimeError("Abort event was set. So someone died and we should end this madness. \nIMPORTANT: "
-                                   "This is not the actual error message! Look further up to see what caused the "
-                                   "error. Please also check whether your RAM was full")
+                raise RuntimeError("One or more background workers are no longer alive. Exiting. Please check the "
+                                   "print statements above for the actual error message")
 
             if item is None:
                 if not in_queue.empty():
@@ -178,9 +178,8 @@ class NonDetMultiThreadedAugmenter(object):
             if self.abort_event.is_set():
                 # self.communication_thread handles checking for dead workers and will set the abort event if necessary
                 self._finish()
-                raise RuntimeError("MultiThreadedAugmenter.abort_event was set, something went wrong. Maybe one of "
-                                   "your workers crashed. This is not the actual error message! Look further up your "
-                                   "stdout to see what caused the error. Please also check whether your RAM was full")
+                raise RuntimeError("One or more background workers are no longer alive. Exiting. Please check the "
+                                   "print statements above for the actual error message")
 
             if not self.results_loop_queue.empty():
                 item = self.results_loop_queue.get()
@@ -209,13 +208,18 @@ class NonDetMultiThreadedAugmenter(object):
             if isinstance(self.generator, DataLoader):
                 self.generator.was_initialized = False
 
+            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]
+                _ = [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.24 → 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: