batchgeneratorsv2 0.3.2__py3-none-any.whl

batchgeneratorsv2/transforms/spatial/spatial.py

@@ -0,0 +1,601 @@
+import math
+from copy import deepcopy
+from typing import Tuple, List, Union
+
+import numpy as np
+import pandas as pd
+import torch
+from scipy.ndimage import fourier_gaussian
+from torch.nn.functional import grid_sample
+
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
+from batchgeneratorsv2.transforms.utils.cropping import crop_tensor
+
+
+class SpatialTransform(BasicTransform):
+    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),
+                 elastic_deform_magnitude: RandomScalar = (0, 0.2),
+                 p_synchronize_def_scale_across_axes: float = 0,
+                 p_rotation: float = 0,
+                 rotation: RandomScalar = (0, 2 * np.pi),
+                 p_rot_per_axis: float = 1,
+                 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',
+                 border_mode_seg: str = "zeros",
+                 center_deformation: bool = True,
+                 mode_image: str = 'bilinear',
+                 padding_mode_image: str = "zeros",
+                 padding_value_seg: float = 0,
+                 padding_value_image: float = 0,
+                 align_corners: bool = False
+                 ):
+        """
+        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", "constant", "reflection", "border"
+
+        border_mode_seg: can be "zeros", "constant", "reflection", "border". padding values are only considered for
+        the corresponding "constant" modes.
+        """
+        super().__init__()
+        self.patch_size = patch_size
+        if not isinstance(patch_center_dist_from_border, (tuple, list)):
+            patch_center_dist_from_border = [patch_center_dist_from_border] * len(patch_size)
+        self.patch_center_dist_from_border = patch_center_dist_from_border
+        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 pixels!!
+        self.p_rotation = p_rotation
+        self.rotation = rotation
+        self.p_rot_per_axis = p_rot_per_axis
+        self.p_scaling = p_scaling
+        self.scaling = scaling  # larger numbers = smaller objects!
+        self.p_synchronize_scaling_across_axes = p_synchronize_scaling_across_axes
+        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.mode_image = mode_image
+        self.padding_mode_image = padding_mode_image
+        self.padding_value_seg = padding_value_seg
+        self.padding_value_image = padding_value_image
+        self.align_corners = align_corners
+        self._grid_cache = {}  # key: (patch_size, dtype) -> base grid tensor
+
+    def _get_base_grid_clone(self) -> torch.Tensor:
+        key = tuple(self.patch_size)
+        g = self._grid_cache.get(key)
+        if g is None:
+            g = _create_centered_identity_grid2(self.patch_size).float().contiguous()
+            self._grid_cache[key] = g
+        return g.clone()
+
+    @staticmethod
+    def _get_crop_pad_settings(padding_mode: str, padding_value: float):
+        if padding_mode == 'reflection':
+            return 'reflect', {}
+        if padding_mode == 'border':
+            return 'replicate', {}
+        if padding_mode == 'zeros':
+            return 'constant', {'value': 0}
+        if padding_mode == 'constant':
+            return 'constant', {'value': padding_value}
+        raise RuntimeError(f'Unknown pad mode: {padding_mode}')
+
+    @staticmethod
+    def _get_grid_sample_padding_mode(padding_mode: str) -> str:
+        if padding_mode in ('zeros', 'constant'):
+            return 'zeros'
+        if padding_mode in ('border', 'reflection'):
+            return padding_mode
+        raise RuntimeError(f'Unknown pad mode: {padding_mode}')
+
+    @staticmethod
+    def _requires_constant_padding_fixup(padding_mode: str, padding_value: float) -> bool:
+        return padding_mode == 'constant' and padding_value != 0
+
+    def _compute_out_of_bounds_mask(self, grid: torch.Tensor, spatial_shape: Tuple[int, ...]) -> torch.Tensor:
+        if self.align_corners:
+            lo = grid.new_tensor(-1.)
+            hi = grid.new_tensor(1.)
+        else:
+            size = grid.new_tensor(spatial_shape)
+            lo = -1 + 1 / size
+            hi = 1 - 1 / size
+        return ((grid < lo) | (grid > hi)).any(dim=-1)
+
+    def get_parameters(self, **data_dict) -> dict:
+        dim = data_dict['image'].ndim - 1
+
+        do_rotation = np.random.uniform() < self.p_rotation
+        do_scale = np.random.uniform() < self.p_scaling
+        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(0, dim)]
+            if self.p_rot_per_axis < 1:
+                for i in range(dim):
+                    if np.random.uniform() > self.p_rot_per_axis:
+                        angles[i] = 0
+        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(0, dim)]
+        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[-1], scales)
+            else:
+                raise RuntimeError(f'Unsupported dimension: {dim}')
+        else:
+            affine = None  # this will allow us to detect that we can skip computations
+
+        # elastic deformation. We need to create the displacement field here
+        # we use the method from augment_spatial_2 in batchgenerators
+        if do_deform:
+            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)
+                ]
+
+            # 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)]
+
+            magnitude = [
+                sample_scalar(self.elastic_deform_magnitude, image=data_dict['image'], patch_size=self.patch_size,
+                              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))
+
+            # all the additional time elastic deform takes is spent here
+            for d in range(dim):
+                # fft torch, slower
+                # for i in range(offsets.ndim - 1):
+                #     offsets[d] = blur_dimension(offsets[d][None], sigmas[d], i, force_use_fft=True, truncate=6)[0]
+
+                # fft numpy, this is faster o.O
+                tmp = np.fft.fftn(offsets[d].numpy())
+                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))
+            spatial_dims = tuple(list(range(1, dim + 1)))
+            offsets = torch.permute(offsets, (*spatial_dims, 0))
+        else:
+            offsets = None
+
+        shape = data_dict['image'].shape[1:]
+        if not self.random_crop:
+            center_location_in_pixels = [i / 2 for i in shape]
+        else:
+            center_location_in_pixels = []
+            for d in range(0, dim):
+                mn = self.patch_center_dist_from_border[d]
+                mx = shape[d] - self.patch_center_dist_from_border[d]
+                if mx < mn:
+                    center_location_in_pixels.append(shape[d] / 2)
+                else:
+                    center_location_in_pixels.append(np.random.uniform(mn, mx))
+        # Precompute the deformed grid once (shared by image, segmentation, regression target)
+        if affine is not None or offsets is not None:
+            grid = self._get_base_grid_clone()
+
+            # we deform first, then rotate
+            if offsets is not None:
+                grid += offsets
+            if affine is not None:
+                grid = torch.matmul(grid, torch.from_numpy(affine).float())
+
+            # 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 self.center_deformation and offsets is not None:
+                mn = grid.mean(dim=list(range(len(shape))))
+            else:
+                mn = 0
+
+            new_center = torch.Tensor([c - s / 2 for c, s in zip(center_location_in_pixels, shape)])
+            grid += (new_center - mn)
+            grid = _convert_my_grid_to_grid_sample_grid(grid, shape)
+        else:
+            grid = None
+
+        return {
+            'center_location_in_pixels': center_location_in_pixels,
+            'grid': grid,
+            # we don't need them but we keep them so that we can debug better
+            'affine': affine,
+            'elastic_offsets': offsets,
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        if params['grid'] is None:
+            # 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
+            pad_mode, pad_kwargs = self._get_crop_pad_settings(self.padding_mode_image, self.padding_value_image)
+            return 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,
+            )
+
+        grid = params['grid']
+        result = grid_sample(
+            img[None],
+            grid[None],
+            mode=self.mode_image,
+            padding_mode=self._get_grid_sample_padding_mode(self.padding_mode_image),
+            align_corners=self.align_corners,
+        )[0]
+        if self._requires_constant_padding_fixup(self.padding_mode_image, self.padding_value_image):
+            out_of_bounds_mask = self._compute_out_of_bounds_mask(grid, img.shape[1:])
+            result.masked_fill_(out_of_bounds_mask.unsqueeze(0), self.padding_value_image)
+        return result
+
+    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
+        segmentation = segmentation.contiguous()
+        if params['grid'] is None:
+            # 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
+            pad_mode, pad_kwargs = self._get_crop_pad_settings(self.border_mode_seg, self.padding_value_seg)
+            return crop_tensor(
+                segmentation,
+                [math.floor(i) for i in params['center_location_in_pixels']],
+                self.patch_size,
+                pad_mode=pad_mode,
+                pad_kwargs=pad_kwargs,
+            )
+
+        grid = params['grid']
+        grid_sample_padding_mode = self._get_grid_sample_padding_mode(self.border_mode_seg)
+
+        if self.mode_seg == 'nearest':
+            result_seg = grid_sample(
+                segmentation[None].float(),
+                grid[None],
+                mode=self.mode_seg,
+                padding_mode=grid_sample_padding_mode,
+                align_corners=self.align_corners
+            )[0].to(segmentation.dtype)
+        else:
+            result_seg = torch.zeros((segmentation.shape[0], *self.patch_size), dtype=segmentation.dtype)
+            if self.bg_style_seg_sampling:
+                for c in range(segmentation.shape[0]):
+                    labels = torch.from_numpy(np.sort(pd.unique(segmentation[c].numpy().ravel())))
+                    # if we only have 2 labels then we can save compute time
+                    if len(labels) == 2:
+                        out = grid_sample(
+                            ((segmentation[c] == labels[1]).float())[None, None],
+                            grid[None],
+                            mode=self.mode_seg,
+                            padding_mode=grid_sample_padding_mode,
+                            align_corners=self.align_corners
+                        )[0][0] >= 0.5
+                        result_seg[c][out] = labels[1]
+                        result_seg[c][~out] = labels[0]
+                    else:
+                        for i, u in enumerate(labels):
+                            result_seg[c][
+                                grid_sample(
+                                    ((segmentation[c] == u).float())[None, None],
+                                    grid[None],
+                                    mode=self.mode_seg,
+                                    padding_mode=grid_sample_padding_mode,
+                                    align_corners=self.align_corners
+                                )[0][0] >= 0.5] = u
+            else:
+                for c in range(segmentation.shape[0]):
+                    labels = torch.from_numpy(np.sort(pd.unique(segmentation[c].numpy().ravel())))
+                    # torch.where(torch.bincount(segmentation.ravel()) > 0)[0].to(segmentation.dtype)
+                    tmp = torch.zeros((len(labels), *self.patch_size), dtype=torch.float16)
+                    scale_factor = 1000
+                    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=grid_sample_padding_mode,
+                            align_corners=self.align_corners
+                        )[0][0]
+                        mask = tmp[i] > (0.7 * scale_factor)
+                        result_seg[c][mask] = u
+                        done_mask = done_mask | mask
+                    if not torch.all(done_mask):
+                        result_seg[c][~done_mask] = labels[tmp[:, ~done_mask].argmax(0)]
+                    del tmp
+
+        if self._requires_constant_padding_fixup(self.border_mode_seg, self.padding_value_seg):
+            out_of_bounds_mask = self._compute_out_of_bounds_mask(grid, segmentation.shape[1:])
+            result_seg.masked_fill_(out_of_bounds_mask.unsqueeze(0), self.padding_value_seg)
+        del grid
+        return result_seg.contiguous()
+
+    def _apply_to_regr_target(self, regression_target, **params) -> torch.Tensor:
+        return self._apply_to_image(regression_target, **params)
+
+    def _apply_to_keypoints(self, keypoints, **params):
+        raise NotImplementedError
+
+    def _apply_to_bbox(self, bbox, **params):
+        raise NotImplementedError
+    
+
+def create_affine_matrix_3d(rotation_angles, scaling_factors):
+    # Rotation matrices for each axis
+    Rx = np.array([[1, 0, 0],
+                   [0, np.cos(rotation_angles[0]), -np.sin(rotation_angles[0])],
+                   [0, np.sin(rotation_angles[0]), np.cos(rotation_angles[0])]])
+
+    Ry = np.array([[np.cos(rotation_angles[1]), 0, np.sin(rotation_angles[1])],
+                   [0, 1, 0],
+                   [-np.sin(rotation_angles[1]), 0, np.cos(rotation_angles[1])]])
+
+    Rz = np.array([[np.cos(rotation_angles[2]), -np.sin(rotation_angles[2]), 0],
+                   [np.sin(rotation_angles[2]), np.cos(rotation_angles[2]), 0],
+                   [0, 0, 1]])
+
+    # Scaling matrix
+    S = np.diag(scaling_factors)
+
+    # Combine rotation and scaling
+    RS = Rz @ Ry @ Rx @ S
+    return RS
+
+
+def create_affine_matrix_2d(rotation_angle, scaling_factors):
+    # Rotation matrix
+    R = np.array([[np.cos(rotation_angle), -np.sin(rotation_angle)],
+                  [np.sin(rotation_angle), np.cos(rotation_angle)]])
+
+    # Scaling matrix
+    S = np.diag(scaling_factors)
+
+    # Combine rotation and scaling
+    RS = R @ S
+    return RS
+
+
+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)
+    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
+
+
+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
+
+
+    import SimpleITK as sitk
+    # img = camera()
+    # seg = None
+    img = sitk.GetArrayFromImage(sitk.ReadImage('/media/isensee/raw_data/nnUNet_raw/Dataset226_BraTS2024-BraTS-GLI/imagesTr/BraTS-GLI-00005-100_0001.nii.gz'))
+    seg = sitk.GetArrayFromImage(sitk.ReadImage('/media/isensee/raw_data/nnUNet_raw/Dataset226_BraTS2024-BraTS-GLI/labelsTr/BraTS-GLI-00005-100.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 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=(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 = [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')
+    # 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/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))})