batchgeneratorsv2 0.2__tar.gz → 0.2.2__tar.gz

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

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

batchgeneratorsv2-0.2.2/batchgeneratorsv2/transforms/intensity/inversion.py

@@ -0,0 +1,51 @@
+import numpy as np
+import torch
+
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+class InvertImageTransform(ImageOnlyTransform):
+    def __init__(self, p_invert_image: float, p_synchronize_channels: float = 1, p_per_channel: float = 1):
+        super().__init__()
+        self.p_invert_image = p_invert_image
+        self.p_synchronize_channels = p_synchronize_channels
+        self.p_per_channel = p_per_channel
+
+    def get_parameters(self, **data_dict) -> dict:
+        shape = data_dict['image'].shape
+        apply = np.random.uniform() < self.p_invert_image
+        if apply:
+            if np.random.uniform() < self.p_synchronize_channels:
+                apply_to_channel = torch.arange(0, shape[0])
+            else:
+                apply_to_channel = torch.where(torch.rand(shape[0]) < self.p_per_channel)[0]
+        else:
+            apply_to_channel = []
+        return {
+            'apply_to_channel': apply_to_channel,
+            'apply': apply,
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        if not params['apply']:
+            return img
+        else:
+            for ch in params['apply_to_channel']:
+                mn = img[ch].mean()
+                img[ch] -= mn
+                img[ch] *= -1
+                img[ch] += mn
+        return img
+
+
+if __name__ == '__main__':
+    mbt = InvertImageTransform(0.5, 0.5, 0.5)
+    from batchviewer import view_batch
+
+    for _ in range(100):
+        data_dict = {'image': torch.ones((2, 20, 192, 64))}
+        data_dict['image'][0, :10] = -1
+        data_dict['image'][1, :5] = -1
+        ret = mbt(**data_dict)
+        print(ret['image'][0, 0, 0, 0], ret['image'][1, 0, 0, 0])
+    view_batch(mbt(**data_dict)['image'])

{batchgeneratorsv2-0.2 → batchgeneratorsv2-0.2.2}/batchgeneratorsv2/transforms/spatial/low_resolution.py

@@ -8,9 +8,13 @@ from torch.nn.functional import interpolate
 
 
 class SimulateLowResolutionTransform(ImageOnlyTransform):
-    def __init__(self, scale: RandomScalar, synchronize_channels: bool, synchronize_axes: bool,
+    def __init__(self,
+                 scale: RandomScalar,
+                 synchronize_channels: bool,
+                 synchronize_axes: bool,
                  ignore_axes: Tuple[int, ...],
-                 allowed_channels: Tuple[int, ...] = None, p_per_channel: float = 1):
+                 allowed_channels: Tuple[int, ...] = None,
+                 p_per_channel: float = 1):
         super().__init__()
         self.scale = scale
         self.synchronize_channels = synchronize_channels

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

@@ -7,7 +7,7 @@ import SimpleITK
 import numpy as np
 import pandas as pd
 import torch
-from scipy.ndimage import fourier_gaussian
+from scipy.ndimage import fourier_gaussian, gaussian_filter
 from torch import Tensor
 from torch.nn.functional import grid_sample
 
@@ -31,8 +31,18 @@ class SpatialTransform(BasicTransform):
                  scaling: RandomScalar = (0.7, 1.3),
                  p_synchronize_scaling_across_axes: float = 0,
                  bg_style_seg_sampling: bool = True,
-                 mode_seg: str = 'bilinear'
+                 mode_seg: str = 'bilinear',
+                 border_mode_seg: str = "zeros",
+                 center_deformation: bool = True,
+                 padding_mode_image: str = "zeros"
                  ):
+        """
+        magnitude must be given in pixels!
+        deformation scale is given as a paercentage of the edge length
+        
+        padding_mode_image: see torch grid_sample documentation. This currently applies to image and regression target 
+        because both call self._apply_to_image. Can be "zeros", "reflection", "border"
+        """
         super().__init__()
         self.patch_size = patch_size
         if not isinstance(patch_center_dist_from_border, (tuple, list)):
@@ -41,7 +51,7 @@ class SpatialTransform(BasicTransform):
         self.random_crop = random_crop
         self.p_elastic_deform = p_elastic_deform
         self.elastic_deform_scale = elastic_deform_scale  # sigma for blurring offsets, in % of patch size. Larger values mean coarser deformation
-        self.elastic_deform_magnitude = elastic_deform_magnitude  # determines the maximum displacement, measured in % of patch size
+        self.elastic_deform_magnitude = elastic_deform_magnitude  # determines the maximum displacement, measured in pixels!!
         self.p_rotation = p_rotation
         self.rotation = rotation
         self.p_scaling = p_scaling
@@ -50,9 +60,11 @@ class SpatialTransform(BasicTransform):
         self.p_synchronize_def_scale_across_axes = p_synchronize_def_scale_across_axes
         self.bg_style_seg_sampling = bg_style_seg_sampling
         self.mode_seg = mode_seg
+        self.border_mode_seg = border_mode_seg
+        self.center_deformation = center_deformation
+        self.padding_mode_image = padding_mode_image
 
     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
@@ -60,14 +72,14 @@ 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 - 1, -1, -1)]
+            angles = [sample_scalar(self.rotation, image=data_dict['image'], dim=i) for i in range(0, 3)]
         else:
             angles = [0] * dim
         if do_scale:
             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)]
+                scales = [sample_scalar(self.scaling, image=data_dict['image'], dim=i) for i in range(0, 3)]
         else:
             scales = [1] * dim
 
@@ -85,7 +97,6 @@ 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)][::-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)
@@ -93,16 +104,16 @@ class SpatialTransform(BasicTransform):
             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)
+                    for i in range(dim)
                     ]
 
             # 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)][::-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.
+            sigmas = [i * j for i, j in zip(deformation_scales, self.patch_size)]
+
             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 - 1, -1, -1)]
+                              dim=i, deformation_scale=deformation_scales[i])
+                for i in range(dim)]
             # doing it like this for better memory layout for blurring
             offsets = torch.normal(mean=0, std=1, size=(dim, *self.patch_size))
 
@@ -117,18 +128,24 @@ class SpatialTransform(BasicTransform):
                 tmp = fourier_gaussian(tmp, sigmas[d])
                 offsets[d] = torch.from_numpy(np.fft.ifftn(tmp).real)
 
+                # tmp = offsets[d].numpy().astype(np.float64)
+                # gaussian_filter(tmp, sigmas[d], 0, output=tmp)
+                # offsets[d] = torch.from_numpy(tmp).to(offsets.dtype)
+                # print(offsets.dtype)
+
                 mx = torch.max(torch.abs(offsets[d]))
                 offsets[d] /= (mx / np.clip(magnitude[d], a_min=1e-8, a_max=np.inf))
-            offsets = torch.permute(offsets, (1, 2, 3, 0))
+            spatial_dims = tuple(list(range(1, dim + 1)))
+            offsets = torch.permute(offsets, (*spatial_dims, 0))
         else:
             offsets = None
-        # 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][::-1]
+            center_location_in_pixels = [i / 2 for i in shape]
         else:
             center_location_in_pixels = []
-            for d in range(dim - 1, -1, -1):
+            for d in range(0, 3):
                 mn = self.patch_center_dist_from_border[d]
                 mx = shape[d] - self.patch_center_dist_from_border[d]
                 if mx < mn:
@@ -146,15 +163,25 @@ 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, [math.floor(i) for i in params['center_location_in_pixels']][::-1], self.patch_size, pad_mode='constant',
-                              pad_kwargs={'value': 0})
+
+            # torch is inconsistent. AAAAaaah
+            if self.padding_mode_image == 'reflection':
+                pad_mode = 'reflect'
+                pad_kwargs = {}
+            elif self.padding_mode_image == 'zeros':
+                pad_mode = 'constant'
+                {'value': 0}
+            elif self.padding_mode_image == 'border':
+                pad_mode = 'replicate'
+                pad_kwargs = {}
+            else:
+                raise RuntimeError('Unknown pad mode')
+
+            img = crop_tensor(img, [math.floor(i) for i in params['center_location_in_pixels']], self.patch_size, pad_mode=pad_mode,
+                              pad_kwargs=pad_kwargs)
             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)][::-1])
-            grid /= grid_scale
+            grid = _create_centered_identity_grid2(self.patch_size)
 
             # we deform first, then rotate
             if params['elastic_offsets'] is not None:
@@ -164,14 +191,16 @@ class SpatialTransform(BasicTransform):
 
             # we center the grid around the center_location_in_pixels. We should center the mean of the grid, not the center position
             # only do this if we elastic deform
-            if params['elastic_offsets'] is not None:
+            if self.center_deformation and 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:][::-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]
+
+            new_center = torch.Tensor([c - s / 2 for c, s in zip(params['center_location_in_pixels'], img.shape[1:])])
+            grid += (new_center - mn)
+            # print(f'grid sample with pad mode {self.padding_mode_image}')
+            return grid_sample(img[None], _convert_my_grid_to_grid_sample_grid(grid, img.shape[1:])[None],
+                               mode='bilinear', padding_mode=self.padding_mode_image, align_corners=False)[0]
 
     def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
         segmentation = segmentation.contiguous()
@@ -180,17 +209,13 @@ class SpatialTransform(BasicTransform):
             # This saves compute.
             # cropping requires the center to be given as integer coordinates
             segmentation = crop_tensor(segmentation,
-                                       [math.floor(i) for i in params['center_location_in_pixels']][::-1],
+                                       [math.floor(i) for i in params['center_location_in_pixels']],
                                        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)][::-1])
-            grid /= grid_scale
+            grid = _create_centered_identity_grid2(self.patch_size)
 
             # we deform first, then rotate
             if params['elastic_offsets'] is not None:
@@ -199,20 +224,22 @@ 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
-            if params['elastic_offsets'] is not None:
+            if self.center_deformation and 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:][::-1], params['center_location_in_pixels'])])
-            grid += - mn + new_center
+
+            new_center = torch.Tensor([c - s / 2 for c, s in zip(params['center_location_in_pixels'], segmentation.shape[1:])])
+
+            grid += (new_center - mn)
+            grid = _convert_my_grid_to_grid_sample_grid(grid, segmentation.shape[1:])
 
             if self.mode_seg == 'nearest':
                 result_seg = grid_sample(
                                 segmentation[None].float(),
                                 grid[None],
                                 mode=self.mode_seg,
-                                padding_mode="zeros",
+                                padding_mode=self.border_mode_seg,
                                 align_corners=False
                             )[0].to(segmentation.dtype)
             else:
@@ -226,7 +253,7 @@ class SpatialTransform(BasicTransform):
                                     ((segmentation[c] == labels[1]).float())[None, None],
                                     grid[None],
                                     mode=self.mode_seg,
-                                    padding_mode="zeros",
+                                    padding_mode=self.border_mode_seg,
                                     align_corners=False
                                 )[0][0] >= 0.5
                             result_seg[c][out] = labels[1]
@@ -238,7 +265,7 @@ class SpatialTransform(BasicTransform):
                                         ((segmentation[c] == u).float())[None, None],
                                         grid[None],
                                         mode=self.mode_seg,
-                                        padding_mode="zeros",
+                                        padding_mode=self.border_mode_seg,
                                         align_corners=False
                                     )[0][0] >= 0.5] = u
                 else:
@@ -250,7 +277,7 @@ class SpatialTransform(BasicTransform):
                         done_mask = torch.zeros(*self.patch_size, dtype=torch.bool)
                         for i, u in enumerate(labels):
                             tmp[i] = grid_sample(((segmentation[c] == u).float() * scale_factor)[None, None], grid[None],
-                                                 mode=self.mode_seg, padding_mode="zeros", align_corners=False)[0][0]
+                                                 mode=self.mode_seg, padding_mode=self.border_mode_seg, align_corners=False)[0][0]
                             mask = tmp[i] > (0.7 * scale_factor)
                             result_seg[c][mask] = u
                             done_mask = done_mask | mask
@@ -305,15 +332,38 @@ def create_affine_matrix_2d(rotation_angle, scaling_factors):
     return RS
 
 
-def _create_identity_grid(size: List[int]) -> Tensor:
-    space = [torch.linspace((-s + 1) / s, (s - 1) / s, s) for s in size[::-1]]
+# def _create_identity_grid(size: List[int]) -> Tensor:
+#     space = [torch.linspace((-s + 1) / s, (s - 1) / s, s) for s in size[::-1]]
+#     grid = torch.meshgrid(space, indexing="ij")
+#     grid = torch.stack(grid, -1)
+#     spatial_dims = list(range(len(size)))
+#     grid = grid.permute((*spatial_dims[::-1], len(size)))
+#     return grid
+
+
+def _create_centered_identity_grid2(size: Union[Tuple[int, ...], List[int]]) -> torch.Tensor:
+    space = [torch.linspace((1 - s) / 2, (s - 1) / 2, s) for s in size]
     grid = torch.meshgrid(space, indexing="ij")
     grid = torch.stack(grid, -1)
-    spatial_dims = list(range(len(size)))
-    grid = grid.permute((*spatial_dims[::-1], len(size)))
     return grid
 
 
+def _convert_my_grid_to_grid_sample_grid(my_grid: torch.Tensor, original_shape: Union[Tuple[int, ...], List[int]]):
+    # rescale
+    for d in range(len(original_shape)):
+        s = original_shape[d]
+        my_grid[..., d] /= (s / 2)
+    my_grid = torch.flip(my_grid, (len(my_grid.shape) - 1, ))
+    # my_grid = my_grid.flip((len(my_grid.shape) - 1,))
+    return my_grid
+
+
+# size = (4, 5, 6)
+# grid_old = _create_identity_grid(size)
+# grid_new = _create_centered_identity_grid2(size)
+# grid_new_converted = _convert_my_grid_to_grid_sample_grid(grid_new, size)
+# torch.all(torch.isclose(grid_new_converted, grid_old))
+
 # An alternative way of generating the displacement fieldQ
 # def displacement_field(data: torch.Tensor):
 #     downscaling_global = np.random.uniform() ** 2 * 4 + 2
@@ -395,30 +445,86 @@ if __name__ == '__main__':
     # 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
-    #
+    def eldef_scale(image, dim, patch_size):
+        return 0.1
+
+    def eldef_magnitude(image, dim, patch_size, deformation_scale):
+        return 10 if dim == 2 else 0
+
+    def rot(image, dim):
+        return 45/360 * 2 * np.pi if dim == 0 else 0
+
+    def scaling(image, dim):
+        return 0.5 if dim == 0 else 1
+
+    # lines
+    patch = torch.zeros((1, 64, 60, 68))
+    patch[:, :, 10, 30] = 1
+    patch[:, 50, :, 30] = 1
+    patch[:, 40, 20, :] = 1
+
+    # patch_block
+    patch_block = torch.zeros((1, 64, 60, 68))
+    patch_block[:, 22:42, 20:40, 24:44] = 1
+
+    patch_line = torch.zeros((1, 64, 60, 128))
+    patch_line[:, 22:24, 30:32, 10:-10] = 1
+    use = patch_line
+
+    sp = SpatialTransform(
+        patch_size=patch.shape[1:],
+        patch_center_dist_from_border=0,
+        random_crop=False,
+        p_elastic_deform=0,
+        p_rotation=1,
+        p_scaling=0,
+        elastic_deform_scale=eldef_scale,
+        elastic_deform_magnitude=eldef_magnitude,
+        p_synchronize_def_scale_across_axes=0,
+        rotation=rot,
+        scaling=scaling,
+        p_synchronize_scaling_across_axes=0,
+        bg_style_seg_sampling=False,
+        mode_seg='bilinear'
+    )
+
+
+    SimpleITK.WriteImage(SimpleITK.GetImageFromArray(use[0].numpy()), 'orig.nii.gz')
+
+    params = sp.get_parameters(image=use)
+    transformed = sp._apply_to_image(use, **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')
+
+    #################
+    # with this part I verify that the crop through spatialtransforms grid sample yields the same result as crop_tensor
+    #################
+
     # sp = SpatialTransform(
-    #     patch_size=(64, 60, 68),
+    #     patch_size=(48, 52, 54),
     #     patch_center_dist_from_border=0,
-    #     random_crop=False,
+    #     random_crop=True,
     #     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'
+    #     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,
     # )
     #
     # patch = torch.zeros((1, 64, 60, 68))
@@ -427,62 +533,20 @@ if __name__ == '__main__':
     # patch[:, 40, 20, :] = 1
     # SimpleITK.WriteImage(SimpleITK.GetImageFromArray(patch[0].numpy()), 'orig.nii.gz')
     #
+    # center_coords = [50, 10, 16]
     # 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, **params2)
     #
     # 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')
-
-    #################
-    # with this part I verify that the crop through spatialtransforms grid sample yields the same result as crop_tensor
-    #################
-
-    sp = SpatialTransform(
-        patch_size=(48, 52, 54),
-        patch_center_dist_from_border=0,
-        random_crop=True,
-        p_elastic_deform=0,
-        p_rotation=1,
-        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,
-    )
-
-    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')
+    # 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)
     ####################

batchgeneratorsv2-0.2.2/batchgeneratorsv2/transforms/spatial/transpose.py

@@ -0,0 +1,67 @@
+from typing import Set
+import numpy as np
+import torch
+
+from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
+
+
+class TransposeAxesTransform(BasicTransform):
+    """
+    A transformation class to permute specified spatial axes of an image and related data.
+
+    Attributes:
+        allowed_axes (Set[int]): Set of spatial axes allowed for permutation (e.g., {1, 2} for y and z axes in an
+        image of shape (c, x, y, z)).
+    """
+
+    def __init__(self, allowed_axes: Set[int]):
+        """
+        Initialize the transform with allowed spatial axes for permutation.
+
+        Args:
+            allowed_axes (Set[int]): Set of spatial axis indices for permutation.
+        """
+        super().__init__()
+        self.allowed_axes = allowed_axes
+
+    def get_parameters(self, **data_dict) -> dict:
+        """
+        Generate a random axis permutation order.
+
+        Args:
+            data_dict (dict): Dictionary containing `image` tensor data.
+
+        Returns:
+            dict: Permutation order of axes as 'axis_order'.
+        """
+        shape_of_allowed = [data_dict['image'].shape[1 + i] for i in self.allowed_axes]
+        if len(shape_of_allowed) < 2:
+            return {'axis_order': list(range(len(data_dict['image'].shape)))}
+        if not all(i == shape_of_allowed[0] for i in shape_of_allowed[1:]):
+            raise ValueError(f"Axis shapes are not identical: {shape_of_allowed}. Cannot permute.\n"
+                             f"Image shape: {data_dict['image'].shape}. Allowed axes: {self.allowed_axes}")
+
+        axes = [i + 1 for i in self.allowed_axes]
+        np.random.shuffle(axes)
+        axis_order = np.arange(len(data_dict['image'].shape))
+        axis_order[np.isin(axis_order, axes)] = axes
+        return {'axis_order': [int(i) for i in axis_order]}
+
+    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
+        return segmentation.permute(params['axis_order']).contiguous()
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        return img.permute(params['axis_order']).contiguous()
+
+    def _apply_to_regr_target(self, regression_target, **params) -> torch.Tensor:
+        return regression_target.permute(params['axis_order']).contiguous()
+
+    def _apply_to_bbox(self, bbox, **params):
+        raise NotImplementedError
+
+    def _apply_to_keypoints(self, keypoints, **params):
+        raise NotImplementedError
+
+if __name__ == '__main__':
+    t = TransposeAxesTransform((1, 2))
+    ret = t(**{'image': torch.rand((2, 31, 32, 32)), 'segmentation': torch.ones((1, 31, 32, 32))})

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

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

{batchgeneratorsv2-0.2 → batchgeneratorsv2-0.2.2}/batchgeneratorsv2.egg-info/SOURCES.txt

@@ -24,6 +24,7 @@ batchgeneratorsv2/transforms/intensity/brightness.py
 batchgeneratorsv2/transforms/intensity/contrast.py
 batchgeneratorsv2/transforms/intensity/gamma.py
 batchgeneratorsv2/transforms/intensity/gaussian_noise.py
+batchgeneratorsv2/transforms/intensity/inversion.py
 batchgeneratorsv2/transforms/nnunet/__init__.py
 batchgeneratorsv2/transforms/nnunet/random_binary_operator.py
 batchgeneratorsv2/transforms/nnunet/remove_connected_components.py
@@ -34,6 +35,7 @@ batchgeneratorsv2/transforms/spatial/__init__.py
 batchgeneratorsv2/transforms/spatial/low_resolution.py
 batchgeneratorsv2/transforms/spatial/mirroring.py
 batchgeneratorsv2/transforms/spatial/spatial.py
+batchgeneratorsv2/transforms/spatial/transpose.py
 batchgeneratorsv2/transforms/utils/__init__.py
 batchgeneratorsv2/transforms/utils/compose.py
 batchgeneratorsv2/transforms/utils/cropping.py

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

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