batchgeneratorsv2 0.1.1__tar.gz → 0.2__tar.gz

{batchgeneratorsv2-0.1.1 → batchgeneratorsv2-0.2}/PKG-INFO

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

{batchgeneratorsv2-0.1.1 → batchgeneratorsv2-0.2}/batchgeneratorsv2/benchmarks/bg_comparison/nnUNet_pipeline_bg.py

@@ -87,4 +87,4 @@ if __name__ == '__main__':
             compute_times[i].append(time() - st)
 
     for t, ct in zip(tr_transforms, compute_times):
-        print(t.__class__.__name__, np.percentile(ct, 20))
+        print(t.__class__.__name__, np.mean(ct))

{batchgeneratorsv2-0.1.1 → batchgeneratorsv2-0.2}/batchgeneratorsv2/benchmarks/bg_comparison/nnUNet_pipeline_here.py

@@ -135,4 +135,4 @@ if __name__ == '__main__':
                 compute_times[i].append(time() - st)
 
     for t, ct in zip(transforms, compute_times):
-        print(t.__class__.__name__ if not isinstance(t, RandomTransform) else t.transform.__class__.__name__, np.percentile(ct, 20))
+        print(t.__class__.__name__ if not isinstance(t, RandomTransform) else t.transform.__class__.__name__, np.mean(ct))

{batchgeneratorsv2-0.1.1 → batchgeneratorsv2-0.2}/batchgeneratorsv2/transforms/base/basic_transform.py

@@ -53,6 +53,11 @@ class BasicTransform(abc.ABC):
     def get_parameters(self, **data_dict) -> dict:
         return {}
 
+    def __repr__(self):
+        ret_str = str(type(self).__name__) + "( " + ", ".join(
+            [key + " = " + repr(val) for key, val in self.__dict__.items()]) + " )"
+        return ret_str
+
 
 class ImageOnlyTransform(BasicTransform):
     def apply(self, data_dict: dict, **params) -> dict:

{batchgeneratorsv2-0.1.1 → batchgeneratorsv2-0.2}/batchgeneratorsv2/transforms/intensity/contrast.py

@@ -22,6 +22,9 @@ class BGContrast():
     def __call__(self, *args, **kwargs):
         return self.sample_contrast(*args, **kwargs)
 
+    def __repr__(self):
+        return self.__class__.__name__ + f"(contrast_range={self.contrast_range})"
+
 class ContrastTransform(ImageOnlyTransform):
     def __init__(self, contrast_range: RandomScalar, preserve_range: bool, synchronize_channels: bool, p_per_channel: float = 1):
         super().__init__()

{batchgeneratorsv2-0.1.1 → batchgeneratorsv2-0.2}/batchgeneratorsv2/transforms/spatial/spatial.py

@@ -1,6 +1,9 @@
 from copy import deepcopy
 from typing import Tuple, List, Union
 
+import math
+
+import SimpleITK
 import numpy as np
 import pandas as pd
 import torch
@@ -14,14 +17,19 @@ from batchgeneratorsv2.transforms.utils.cropping import crop_tensor
 
 
 class SpatialTransform(BasicTransform):
-    def __init__(self, patch_size: Tuple[int, ...],
+    def __init__(self,
+                 patch_size: Tuple[int, ...],
                  patch_center_dist_from_border: Union[int, List[int], Tuple[int, ...]],
                  random_crop: bool,
-                 p_elastic_deform: float = 0, elastic_deform_scale: RandomScalar = (0, 0.2),
+                 p_elastic_deform: float = 0,
+                 elastic_deform_scale: RandomScalar = (0, 0.2),
                  elastic_deform_magnitude: RandomScalar = (0, 0.2),
-                 p_synchronize_def_scale_across_axes: float = False,
-                 p_rotation: float = 0, rotation: RandomScalar = (0, 2 * np.pi),
-                 p_scaling: float = 0, scaling: RandomScalar = (0.7, 1.3), p_synchronize_scaling_across_axes: float = False,
+                 p_synchronize_def_scale_across_axes: float = 0,
+                 p_rotation: float = 0,
+                 rotation: RandomScalar = (0, 2 * np.pi),
+                 p_scaling: float = 0,
+                 scaling: RandomScalar = (0.7, 1.3),
+                 p_synchronize_scaling_across_axes: float = 0,
                  bg_style_seg_sampling: bool = True,
                  mode_seg: str = 'bilinear'
                  ):
@@ -44,6 +52,7 @@ class SpatialTransform(BasicTransform):
         self.mode_seg = mode_seg
 
     def get_parameters(self, **data_dict) -> dict:
+        # note that we revert the axis order here because grid_sample uses dimensions in reverse order!
         dim = data_dict['image'].ndim - 1
 
         do_rotation = np.random.uniform() < self.p_rotation
@@ -51,21 +60,23 @@ class SpatialTransform(BasicTransform):
         do_deform = np.random.uniform() < self.p_elastic_deform
 
         if do_rotation:
-            angles = [sample_scalar(self.rotation, image=data_dict['image'], dim=i) for i in range(dim)]
+            angles = [sample_scalar(self.rotation, image=data_dict['image'], dim=i) for i in range(dim - 1, -1, -1)]
         else:
             angles = [0] * dim
         if do_scale:
-            scales = [sample_scalar(self.scaling, image=data_dict['image'], dim=i) for i in
-                      range(dim)] if np.random.uniform() < self.p_synchronize_scaling_across_axes else [sample_scalar(
-                self.scaling, image=data_dict['image'], dim=None)] * dim
+            if np.random.uniform() <= self.p_synchronize_scaling_across_axes:
+                scales = [sample_scalar(self.scaling, image=data_dict['image'], dim=None)] * dim
+            else:
+                scales = [sample_scalar(self.scaling, image=data_dict['image'], dim=i) for i in range(dim - 1, -1, -1)]
         else:
             scales = [1] * dim
+
         # affine matrix
         if do_scale or do_rotation:
             if dim == 3:
                 affine = create_affine_matrix_3d(angles, scales)
             elif dim == 2:
-                affine = create_affine_matrix_2d(angles[0], scales)
+                affine = create_affine_matrix_2d(angles[-1], scales)
             else:
                 raise RuntimeError(f'Unsupported dimension: {dim}')
         else:
@@ -74,19 +85,24 @@ class SpatialTransform(BasicTransform):
         # elastic deformation. We need to create the displacement field here
         # we use the method from augment_spatial_2 in batchgenerators
         if do_deform:
-            grid_scale = [i / j for i, j in zip(data_dict['image'].shape[1:], self.patch_size)]
-            deformation_scales = [
-                sample_scalar(self.elastic_deform_scale, image=data_dict['image'], dim=i, patch_size=self.patch_size)
-                for i in
-                range(dim)] if np.random.uniform() < self.p_synchronize_scaling_across_axes else [sample_scalar(
-                self.elastic_deform_scale, image=data_dict['image'], dim=None, patch_size=self.patch_size)] * dim
+            grid_scale = [i / j for i, j in zip(data_dict['image'].shape[1:], self.patch_size)][::-1]
+            if np.random.uniform() <= self.p_synchronize_def_scale_across_axes:
+                deformation_scales = [
+                    sample_scalar(self.elastic_deform_scale, image=data_dict['image'], dim=None, patch_size=self.patch_size)
+                    ] * dim
+            else:
+                deformation_scales = [
+                    sample_scalar(self.elastic_deform_scale, image=data_dict['image'], dim=i, patch_size=self.patch_size)
+                    for i in range(dim - 1, -1, -1)
+                    ]
+
             # sigmas must be in pixels, as this will be applied to the deformation field
-            sigmas = [i * j for i, j in zip(deformation_scales, self.patch_size)]
+            sigmas = [i * j for i, j in zip(deformation_scales, self.patch_size)][::-1]
             # the magnitude of the deformation field must adhere to the torch's value range for grid_sample, i.e. [-1. 1] and not pixel coordinates. Do not use sigmas here
             # we need to correct magnitude by grid_scale to account for the fact that the grid will be wrt to the image size but the magnitude should be wrt the patch size. oof.
             magnitude = [
                 sample_scalar(self.elastic_deform_magnitude, image=data_dict['image'], patch_size=self.patch_size,
-                              dim=i, deformation_scale=deformation_scales[i]) / grid_scale[i] for i in range(dim)]
+                              dim=i, deformation_scale=deformation_scales[i]) / grid_scale[i] for i in range(dim - 1, -1, -1)]
             # doing it like this for better memory layout for blurring
             offsets = torch.normal(mean=0, std=1, size=(dim, *self.patch_size))
 
@@ -98,7 +114,7 @@ class SpatialTransform(BasicTransform):
 
                 # fft numpy, this is faster o.O
                 tmp = np.fft.fftn(offsets[d].numpy())
-                tmp = fourier_gaussian(tmp, sigmas)
+                tmp = fourier_gaussian(tmp, sigmas[d])
                 offsets[d] = torch.from_numpy(np.fft.ifftn(tmp).real)
 
                 mx = torch.max(torch.abs(offsets[d]))
@@ -109,10 +125,10 @@ class SpatialTransform(BasicTransform):
         # grid center must be in [-1, 1] as required by grid_sample
         shape = data_dict['image'].shape[1:]
         if not self.random_crop:
-            center_location_in_pixels = [i / 2 for i in shape]
+            center_location_in_pixels = [i / 2 for i in shape][::-1]
         else:
             center_location_in_pixels = []
-            for d in range(dim):
+            for d in range(dim - 1, -1, -1):
                 mn = self.patch_center_dist_from_border[d]
                 mx = shape[d] - self.patch_center_dist_from_border[d]
                 if mx < mn:
@@ -130,14 +146,14 @@ class SpatialTransform(BasicTransform):
             # No spatial transformation is being done. Round grid_center and crop without having to interpolate.
             # This saves compute.
             # cropping requires the center to be given as integer coordinates
-            img = crop_tensor(img, [round(i) for i in params['center_location_in_pixels']], self.patch_size, pad_mode='constant',
+            img = crop_tensor(img, [math.floor(i) for i in params['center_location_in_pixels']][::-1], self.patch_size, pad_mode='constant',
                               pad_kwargs={'value': 0})
             return img
         else:
             grid = _create_identity_grid(self.patch_size)
 
             # the grid must be scaled. The grid is [-1, 1] in image coordinates, but we want it to represent the smaller patch
-            grid_scale = torch.Tensor([i / j for i, j in zip(img.shape[1:], self.patch_size)])
+            grid_scale = torch.Tensor([i / j for i, j in zip(img.shape[1:], self.patch_size)][::-1])
             grid /= grid_scale
 
             # we deform first, then rotate
@@ -147,9 +163,13 @@ class SpatialTransform(BasicTransform):
                 grid = torch.matmul(grid, torch.from_numpy(params['affine']).float())
 
             # we center the grid around the center_location_in_pixels. We should center the mean of the grid, not the center position
-            mn = grid.mean(dim=list(range(img.ndim - 1)))
+            # only do this if we elastic deform
+            if params['elastic_offsets'] is not None:
+                mn = grid.mean(dim=list(range(img.ndim - 1)))
+            else:
+                mn = 0
             new_center = torch.Tensor(
-                [(j / (i / 2) - 1) for i, j in zip(img.shape[1:], params['center_location_in_pixels'])])
+                [(j / (i / 2) - 1) for i, j in zip(img.shape[1:][::-1], params['center_location_in_pixels'])])
             grid += - mn + new_center
             return grid_sample(img[None], grid[None], mode='bilinear', padding_mode="zeros", align_corners=False)[0]
 
@@ -159,14 +179,17 @@ class SpatialTransform(BasicTransform):
             # No spatial transformation is being done. Round grid_center and crop without having to interpolate.
             # This saves compute.
             # cropping requires the center to be given as integer coordinates
-            segmentation = crop_tensor(segmentation, [round(i) for i in params['center_location_in_pixels']], self.patch_size,
-                                       pad_mode='constant', pad_kwargs={'value': 0})
+            segmentation = crop_tensor(segmentation,
+                                       [math.floor(i) for i in params['center_location_in_pixels']][::-1],
+                                       self.patch_size,
+                                       pad_mode='constant',
+                                       pad_kwargs={'value': 0})
             return segmentation
         else:
             grid = _create_identity_grid(self.patch_size)
 
             # the grid must be scaled. The grid is [-1, 1] in image coordinates, but we want it to represent the smaller patch
-            grid_scale = torch.Tensor([i / j for i, j in zip(segmentation.shape[1:], self.patch_size)])
+            grid_scale = torch.Tensor([i / j for i, j in zip(segmentation.shape[1:], self.patch_size)][::-1])
             grid /= grid_scale
 
             # we deform first, then rotate
@@ -176,9 +199,12 @@ class SpatialTransform(BasicTransform):
                 grid = torch.matmul(grid, torch.from_numpy(params['affine']).float())
 
             # we center the grid around the center_location_in_pixels. We should center the mean of the grid, not the center coordinate
-            mn = grid.mean(dim=list(range(segmentation.ndim - 1)))
+            if params['elastic_offsets'] is not None:
+                mn = grid.mean(dim=list(range(segmentation.ndim - 1)))
+            else:
+                mn = 0
             new_center = torch.Tensor(
-                [(j / (i / 2) - 1) for i, j in zip(segmentation.shape[1:], params['center_location_in_pixels'])])
+                [(j / (i / 2) - 1) for i, j in zip(segmentation.shape[1:][::-1], params['center_location_in_pixels'])])
             grid += - mn + new_center
 
             if self.mode_seg == 'nearest':
@@ -308,58 +334,176 @@ def _create_identity_grid(size: List[int]) -> Tensor:
 
 
 if __name__ == '__main__':
-    torch.set_num_threads(1)
-
-    shape = (128, 128, 128)
-    patch_size = (128, 128, 128)
-    labels = 2
-
-
-    # seg = torch.rand([i // 32 for i in shape]) * labels
-    # seg_up = torch.round(torch.nn.functional.interpolate(seg[None, None], size=shape, mode='trilinear')[0],
-    #                      decimals=0).to(torch.int16)
-    # img = torch.ones((1, *shape))
-    # img[tuple([slice(img.shape[0])] + [slice(i // 4, i // 4 * 2) for i in shape])] = 200
+    # torch.set_num_threads(1)
+    #
+    # shape = (128, 128, 128)
+    # patch_size = (128, 128, 128)
+    # labels = 2
+    #
+    #
+    # # seg = torch.rand([i // 32 for i in shape]) * labels
+    # # seg_up = torch.round(torch.nn.functional.interpolate(seg[None, None], size=shape, mode='trilinear')[0],
+    # #                      decimals=0).to(torch.int16)
+    # # img = torch.ones((1, *shape))
+    # # img[tuple([slice(img.shape[0])] + [slice(i // 4, i // 4 * 2) for i in shape])] = 200
+    #
+    #
+    # import SimpleITK as sitk
+    # # img = camera()
+    # # seg = None
+    # img = sitk.GetArrayFromImage(sitk.ReadImage('/media/isensee/raw_data/nnUNet_raw/Dataset137_BraTS2021/imagesTr/BraTS2021_00000_0000.nii.gz'))
+    # seg = sitk.GetArrayFromImage(sitk.ReadImage('/media/isensee/raw_data/nnUNet_raw/Dataset137_BraTS2021/labelsTr/BraTS2021_00000.nii.gz'))
+    #
+    # patch_size = (192, 192, 192)
+    # sp = SpatialTransform(
+    #     patch_size=(192, 192, 192),
+    #     patch_center_dist_from_border=[i / 2 for i in patch_size],
+    #     random_crop=True,
+    #     p_elastic_deform=0,
+    #     elastic_deform_magnitude=(0.1, 0.1),
+    #     elastic_deform_scale=(0.1, 0.1),
+    #     p_synchronize_def_scale_across_axes=0.5,
+    #     p_rotation=1,
+    #     rotation=(-30 / 360 * np.pi, 30 / 360 * np.pi),
+    #     p_scaling=1,
+    #     scaling=(0.75, 1),
+    #     p_synchronize_scaling_across_axes=0.5,
+    #     bg_style_seg_sampling=True,
+    #     mode_seg='bilinear'
+    # )
+    #
+    # data_dict = {'image': torch.from_numpy(deepcopy(img[None])).float()}
+    # if seg is not None:
+    #     data_dict['segmentation'] = torch.from_numpy(deepcopy(seg[None]))
+    # # out = sp(**data_dict)
+    # #
+    # # view_batch(out['image'], out['segmentation'])
+    #
+    # from time import time
+    # times = []
+    # for _ in range(10):
+    #     data_dict = {'image': torch.from_numpy(deepcopy(img[None])).float()}
+    #     if seg is not None:
+    #         data_dict['segmentation'] = torch.from_numpy(deepcopy(seg[None]))
+    #     st = time()
+    #     out = sp(**data_dict)
+    #     times.append(time() - st)
+    # print(np.median(times))
+
+
+    #################
+    # with this part we can qualitatively test that the correct axes are ebing augmented. Just set one of the probs to 1 and off you go
+    #################
+
+    # def constant_scaling(image, dim, patch_size):
+    #     return 0.1
+    #
+    # def constant_magnitude(image, dim, patch_size, deformation_scale):
+    #     return 0.25 if dim == 2 else 0
+    #
+    # def rot(image, dim):
+    #     return 45/360 * 2 * np.pi if dim == 1 else 0
+    #
+    # sp = SpatialTransform(
+    #     patch_size=(64, 60, 68),
+    #     patch_center_dist_from_border=0,
+    #     random_crop=False,
+    #     p_elastic_deform=0,
+    #     elastic_deform_scale=0,
+    #     elastic_deform_magnitude=0,
+    #     p_synchronize_def_scale_across_axes=0,
+    #     p_rotation=1,
+    #     rotation=rot,
+    #     p_scaling=0,
+    #     scaling=constant_scaling,
+    #     p_synchronize_scaling_across_axes=0,
+    #     bg_style_seg_sampling=False,
+    #     mode_seg='bilinear'
+    # )
+    #
+    # patch = torch.zeros((1, 64, 60, 68))
+    # patch[:, :, 10, 30] = 1
+    # patch[:, 50, :, 30] = 1
+    # patch[:, 40, 20, :] = 1
+    # SimpleITK.WriteImage(SimpleITK.GetImageFromArray(patch[0].numpy()), 'orig.nii.gz')
+    #
+    # params = sp.get_parameters(image=patch)
+    # transformed = sp._apply_to_image(patch, **params)
+    #
+    # SimpleITK.WriteImage(SimpleITK.GetImageFromArray(transformed[0].numpy()), 'transformed.nii.gz')
 
+    # p = torch.zeros((1, 1, 8, 16, 32))
+    # p[:, :, 2:6, 10:16, 10:24] = 1
+    # grid = _create_identity_grid(p.shape[2:])
+    # grid[:, :, :, 0] *= 0.5
+    # out = grid_sample(p, grid[None], mode='bilinear', padding_mode="zeros", align_corners=False)
+    # torch.all(out == p)
+    # SimpleITK.WriteImage(SimpleITK.GetImageFromArray(p[0, 0].numpy()), 'orig.nii.gz')
+    # SimpleITK.WriteImage(SimpleITK.GetImageFromArray(out[0, 0].numpy()), 'transformed.nii.gz')
 
-    import SimpleITK as sitk
-    # img = camera()
-    # seg = None
-    img = sitk.GetArrayFromImage(sitk.ReadImage('/media/isensee/raw_data/nnUNet_raw/Dataset137_BraTS2021/imagesTr/BraTS2021_00000_0000.nii.gz'))
-    seg = sitk.GetArrayFromImage(sitk.ReadImage('/media/isensee/raw_data/nnUNet_raw/Dataset137_BraTS2021/labelsTr/BraTS2021_00000.nii.gz'))
+    #################
+    # with this part I verify that the crop through spatialtransforms grid sample yields the same result as crop_tensor
+    #################
 
-    patch_size = (192, 192, 192)
     sp = SpatialTransform(
-        patch_size=(192, 192, 192),
-        patch_center_dist_from_border=[i / 2 for i in patch_size],
+        patch_size=(48, 52, 54),
+        patch_center_dist_from_border=0,
         random_crop=True,
         p_elastic_deform=0,
-        elastic_deform_magnitude=(0.1, 0.1),
-        elastic_deform_scale=(0.1, 0.1),
-        p_synchronize_def_scale_across_axes=0.5,
         p_rotation=1,
-        rotation=(-30 / 360 * np.pi, 30 / 360 * np.pi),
-        p_scaling=1,
-        scaling=(0.75, 1),
-        p_synchronize_scaling_across_axes=0.5,
-        bg_style_seg_sampling=True,
-        mode_seg='bilinear'
+        p_scaling=0,
+        rotation=0
+    )
+    sp2 = SpatialTransform(
+        patch_size=(48, 52, 54),
+        patch_center_dist_from_border=0,
+        random_crop=True,
+        p_elastic_deform=0,
+        p_rotation=0,
+        p_scaling=0,
     )
 
-    data_dict = {'image': torch.from_numpy(deepcopy(img[None])).float()}
-    if seg is not None:
-        data_dict['segmentation'] = torch.from_numpy(deepcopy(seg[None]))
-    # out = sp(**data_dict)
-    #
-    # view_batch(out['image'], out['segmentation'])
-
-    from time import time
-    times = []
-    for _ in range(10):
-        data_dict = {'image': torch.from_numpy(deepcopy(img[None])).float()}
-        if seg is not None:
-            data_dict['segmentation'] = torch.from_numpy(deepcopy(seg[None]))
-        st = time()
-        out = sp(**data_dict)
-        times.append(time() - st)
-    print(np.median(times))
+    patch = torch.zeros((1, 64, 60, 68))
+    patch[:, :, 10, 30] = 1
+    patch[:, 50, :, 30] = 1
+    patch[:, 40, 20, :] = 1
+    SimpleITK.WriteImage(SimpleITK.GetImageFromArray(patch[0].numpy()), 'orig.nii.gz')
+
+    center_coords = [30, 28, 44]
+    params = sp.get_parameters(image=patch)
+    params['center_location_in_pixels'] = center_coords
+    params2 = sp2.get_parameters(image=patch)
+    params2['center_location_in_pixels'] = center_coords
+    transformed = sp._apply_to_image(patch, **params)
+    transformed2 = sp._apply_to_image(patch, **params)
+
+    SimpleITK.WriteImage(SimpleITK.GetImageFromArray(transformed[0].numpy()), 'transformed.nii.gz')
+    SimpleITK.WriteImage(SimpleITK.GetImageFromArray(transformed2[0].numpy()), 'transformed2.nii.gz')
+
+
+
+####################
+    # This is exploraroty code to check how to retrieve coordinates. I used it to verify that grid_sample does in fact
+    # use coordinates in reversed dimension order (zyx and not xyz)
+    ####################
+    # # create a dummy input which has a unique shape in each exis
+    # p = torch.zeros((1, 1, 8, 16, 32))
+    # # set one pixel to 1
+    # p[:, :, 4, 0, 31] = 1
+    # # now create an identity grid. I have verified that this grid yields the same image as the input when used in grid_sample. So the grid is correct
+    # grid = _create_identity_grid((8, 16, 32)).contiguous() # grid is shape torch.Size([8, 16, 32, 3])
+    # out = grid_sample(p, grid[None], mode='bilinear', padding_mode="zeros", align_corners=False)
+    # assert torch.all(out == p)  # this passes
+    # # reduce the grid to the location we are interested in. That are the coordinates where we placed the 1. The 4:5 etc is only so that we keep the number of dimensions
+    # grid = grid[4:5, 0:1, 31:32]
+    # # What coordinate would we expect? Note that grid is [-1, 1]
+    # # For the first dimension, coordinate 4 out of shape 8 is approximately in the middle, so about 0
+    # # For the second dimension, coordinate 0 out of shape 16 is very low, so we expect -1 ish (remember there is aligned corners and shit)
+    # # For the third dimension, coordinate 31 out of shape 32 is very high, so we expect 1 ish (remember there is aligned corners and shit)
+    # # So we expect [0, -1, 1]
+    # # What do we get?
+    # print(grid)
+    # # > tensor([[[[ 0.9688, -0.9375,  0.1250]]]])
+    # # not what we expect
+    # out = grid_sample(p, grid[None], mode='bilinear', padding_mode="zeros", align_corners=False)
+    # assert out.item() == 1

{batchgeneratorsv2-0.1.1 → batchgeneratorsv2-0.2}/batchgeneratorsv2/transforms/utils/random.py

@@ -16,4 +16,8 @@ class RandomTransform(BasicTransform):
         if params['apply_transform']:
             return self.transform(**data_dict)
         else:
-            return data_dict
+            return data_dict
+    
+    def __repr__(self):
+        ret_str = f"{type(self).__name__}(p={self.apply_probability}, transform={self.transform})"
+        return ret_str

batchgeneratorsv2-0.2/batchgeneratorsv2/transforms/utils/seg_to_regions.py

@@ -0,0 +1,24 @@
+from typing import Union, List, Tuple
+import torch
+
+from batchgeneratorsv2.transforms.base.basic_transform import SegOnlyTransform
+
+
+class ConvertSegmentationToRegionsTransform(SegOnlyTransform):
+    def __init__(self, regions: Union[List, Tuple], channel_in_seg: int = 0):
+        super().__init__()
+        self.regions = [torch.Tensor(i) if not isinstance(i, int) else torch.Tensor([i]) for i in regions]
+        self.channel_in_seg = channel_in_seg
+
+    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
+        num_regions = len(self.regions)
+        region_output = torch.zeros((num_regions, *segmentation.shape[1:]), dtype=torch.bool, device=segmentation.device)
+        for region_id, region_labels in enumerate(self.regions):
+            if len(region_labels) == 1:
+                region_output[region_id] = segmentation[self.channel_in_seg] == region_labels
+            else:
+                region_output[region_id] = torch.isin(segmentation[self.channel_in_seg], region_labels)
+        # we return bool here and leave it to the loss function to cast it to whatever it needs. Transferring bool to
+        # device followed by cast on device should be faster than having fp32 here and transferring that
+        return region_output
+

{batchgeneratorsv2-0.1.1 → batchgeneratorsv2-0.2}/batchgeneratorsv2.egg-info/PKG-INFO

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

{batchgeneratorsv2-0.1.1 → batchgeneratorsv2-0.2}/pyproject.toml

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

batchgeneratorsv2-0.1.1/batchgeneratorsv2/transforms/utils/seg_to_regions.py

@@ -1,23 +0,0 @@
-from typing import Union, List, Tuple
-
-import torch
-
-from batchgeneratorsv2.transforms.base.basic_transform import SegOnlyTransform
-
-
-class ConvertSegmentationToRegionsTransform(SegOnlyTransform):
-    def __init__(self, regions: Union[List, Tuple], channel_in_seg: int = 0):
-        super().__init__()
-        self.regions = regions
-        self.channel_in_seg = channel_in_seg
-
-    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
-        num_regions = len(self.regions)
-        region_output = torch.zeros((num_regions, *segmentation.shape[1:]), dtype=torch.bool, device=segmentation.device)
-        if isinstance(region_labels, int) or len(region_labels) == 1:
-            if not isinstance(region_labels, int):
-                region_labels = region_labels[0]
-            region_output[:, region_id] = seg[:, self.seg_channel] == region_labels
-        else:
-            region_output[:, region_id] |= np.isin(seg[:, self.seg_channel], region_labels)
-        return region_output.to(segmentation.dtype)