batchgeneratorsv2 0.3.2__py3-none-any.whl

batchgeneratorsv2/transforms/local/local_contrast.py

@@ -0,0 +1,90 @@
+import torch
+import numpy as np
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+from batchgeneratorsv2.transforms.local.local_transform import LocalTransform
+
+
+class LocalContrastTransform(ImageOnlyTransform, LocalTransform):
+    """
+    Applies localized contrast modification using a spatial Gaussian mask.
+
+    A contrast-modified version of the image is blended with the original using a kernel-based interpolation.
+
+    Args:
+        scale (RandomScalar): Gaussian spread for the spatial weighting mask.
+        loc (RandomScalar): Relative center position for the Gaussian (in % of image size).
+        new_contrast (RandomScalar): Multiplicative factor for local contrast. 1.0 = no change.
+        same_for_all_channels (bool): Whether to use one kernel/contrast value for all channels.
+        p_per_channel (float): Probability to apply to each channel.
+    """
+
+    def __init__(self,
+                 scale: RandomScalar,
+                 loc: RandomScalar = (-1, 2),
+                 new_contrast: RandomScalar = (0.5, 1.5),
+                 same_for_all_channels: bool = True,
+                 p_per_channel: float = 1.0):
+        ImageOnlyTransform.__init__(self)
+        LocalTransform.__init__(self, scale, loc)
+
+        self.new_contrast = new_contrast
+        self.same_for_all_channels = same_for_all_channels
+        self.p_per_channel = p_per_channel
+
+    def get_parameters(self, image: torch.Tensor, **kwargs) -> dict:
+        C, *spatial = image.shape
+        apply_channel = [np.random.rand() < self.p_per_channel for _ in range(C)]
+
+        if not any(apply_channel):
+            return {'kernels': [None] * C, 'contrasts': [None] * C}
+
+        if self.same_for_all_channels:
+            kernel = self._generate_kernel(spatial).astype(np.float32)
+            contrast = sample_scalar(self.new_contrast)
+
+            kernels = [kernel if apply else None for apply in apply_channel]
+            contrasts = [contrast if apply else None for apply in apply_channel]
+        else:
+            kernels, contrasts = [], []
+            for apply in apply_channel:
+                if not apply:
+                    kernels.append(None)
+                    contrasts.append(None)
+                    continue
+                kernel = self._generate_kernel(spatial).astype(np.float32)
+                contrast = sample_scalar(self.new_contrast)
+                kernels.append(kernel)
+                contrasts.append(contrast)
+
+        return {'kernels': kernels, 'contrasts': contrasts}
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        img_np = img.cpu().numpy()
+
+        for c, (kernel, contrast) in enumerate(zip(params['kernels'], params['contrasts'])):
+            if kernel is None:
+                continue
+
+            channel = img_np[c]
+            mean = (channel * kernel).sum() / (kernel.sum() + 1e-8)
+            modified = (channel - mean) * contrast + mean
+            img_np[c] = self.run_interpolation(channel, modified, kernel)
+
+        return torch.from_numpy(img_np).to(img.device, dtype=img.dtype)
+
+
+if __name__ == '__main__':
+    from batchviewer import view_batch
+
+    # Single-channel synthetic volume
+    image = torch.rand(1, 32, 64, 64)
+
+    # Or contrast
+    contrast = LocalContrastTransform(scale=(10, 20), new_contrast=(0.3, 2.0), p_per_channel=1.0)
+
+    # Apply either one
+    params = contrast.get_parameters(image=image)
+    image_aug = contrast._apply_to_image(image.clone(), **params)
+
+    view_batch(image, image_aug)

batchgeneratorsv2/transforms/local/local_gamma.py

@@ -0,0 +1,104 @@
+import torch
+import numpy as np
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+from batchgeneratorsv2.transforms.local.local_transform import LocalTransform
+
+
+class LocalGammaTransform(ImageOnlyTransform, LocalTransform):
+    """
+    Applies locally varying gamma correction to an image using a spatial Gaussian weighting mask.
+
+    A Gaussian kernel is randomly placed in the image and used to blend between the original image and
+    a gamma-corrected version. This simulates localized nonlinear intensity shifts, useful for data augmentation
+    in medical imaging or general contrast robustness.
+
+    Parameters:
+        scale (RandomScalar): Controls the width of the Gaussian (std dev). Recommend large values (e.g. 10–30).
+        loc (RandomScalar): Controls Gaussian center as a % of image size. E.g. (-1, 2) allows off-canvas kernels.
+        gamma (RandomScalar): The gamma exponent applied locally. Try wild distributions :)
+        same_for_all_channels (bool): If True, one kernel is reused across all channels. Otherwise sampled per-channel.
+        p_per_channel (float): Probability to apply gamma correction to each channel independently.
+    """
+
+    def __init__(self,
+                 scale: RandomScalar,
+                 loc: RandomScalar = (-1, 2),
+                 gamma: RandomScalar = (0.5, 1),
+                 same_for_all_channels: bool = True,
+                 p_per_channel: float = 1.0):
+        ImageOnlyTransform.__init__(self)
+        LocalTransform.__init__(self, scale, loc)
+
+        self.gamma = gamma
+        self.same_for_all_channels = same_for_all_channels
+        self.p_per_channel = p_per_channel
+
+    def get_parameters(self, image: torch.Tensor, **kwargs) -> dict:
+        C, *spatial = image.shape
+        apply_channel = [np.random.rand() < self.p_per_channel for _ in range(C)]
+
+        if not any(apply_channel):
+            return {'kernels': [None] * C, 'gammas': [None] * C}
+
+        if self.same_for_all_channels:
+            kernel = self._generate_kernel(spatial).astype(np.float32)
+            gamma = sample_scalar(self.gamma)
+
+            kernels = [kernel if apply else None for apply in apply_channel]
+            gammas = [gamma if apply else None for apply in apply_channel]
+        else:
+            kernels, gammas = [], []
+            for apply in apply_channel:
+                if not apply:
+                    kernels.append(None)
+                    gammas.append(None)
+                    continue
+                kernel = self._generate_kernel(spatial).astype(np.float32)
+                gamma = sample_scalar(self.gamma)
+                kernels.append(kernel)
+                gammas.append(gamma)
+
+        return {'kernels': kernels, 'gammas': gammas}
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        img_np = img.cpu().numpy()
+
+        for c, (kernel, gamma) in enumerate(zip(params['kernels'], params['gammas'])):
+            if kernel is None:
+                continue
+
+            channel = img_np[c]
+            min_val, max_val = channel.min(), channel.max()
+            denom = max(max_val - min_val, 1e-8)
+
+            # Normalize to [0, 1]
+            norm = (channel - min_val) / denom
+            gamma_corrected = np.power(norm, gamma)
+
+            # Blend using kernel
+            blended = self.run_interpolation(norm, gamma_corrected, kernel)
+
+            # Rescale to original range
+            img_np[c] = blended * denom + min_val
+
+        return torch.from_numpy(img_np).to(img.device, dtype=img.dtype)
+
+
+if __name__ == '__main__':
+    import torch
+    from batchviewer import view_batch
+
+    image = torch.rand(1, 32, 64, 64)  # (C, D, H, W)
+
+    transform = LocalGammaTransform(
+        scale=(10, 20),
+        gamma=lambda *_: np.random.uniform(0.01, 1) if np.random.rand() < 0.5 else np.random.uniform(1, 3),
+        same_for_all_channels=False,
+        p_per_channel=1.0
+    )
+
+    params = transform.get_parameters(image=image)
+    image_gamma = transform._apply_to_image(image.clone(), **params)
+
+    view_batch(image, image_gamma)

batchgeneratorsv2/transforms/local/local_smoothing.py

@@ -0,0 +1,98 @@
+import torch
+import numpy as np
+from batchgeneratorsv2.transforms.local.local_transform import LocalTransform
+from scipy.ndimage import gaussian_filter
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+
+
+class LocalSmoothingTransform(ImageOnlyTransform, LocalTransform):
+    """
+    Applies localized Gaussian smoothing to parts of the image using a spatial Gaussian mask.
+
+    A blurred copy of the image is interpolated with the original, weighted by a Gaussian kernel.
+    The strength and extent of the blur are both controllable.
+
+    Args:
+        scale (RandomScalar): Gaussian spread for the spatial weighting mask.
+        loc (RandomScalar): Relative center position for the Gaussian (in % of image size).
+        smoothing_strength (RandomScalar): Max weight of the smoothed image in the interpolation [0, 1].
+        kernel_size (RandomScalar): Sigma for the actual Gaussian smoothing of the image.
+        same_for_all_channels (bool): Whether to apply the same kernel to all channels.
+        p_per_channel (float): Probability of applying transform per channel.
+    """
+
+    def __init__(self,
+                 scale: RandomScalar,
+                 loc: RandomScalar = (-1, 2),
+                 smoothing_strength: RandomScalar = (0.5, 1.0),
+                 kernel_size: RandomScalar = (0.5, 1.5),
+                 same_for_all_channels: bool = True,
+                 p_per_channel: float = 1.0):
+        ImageOnlyTransform.__init__(self)
+        LocalTransform.__init__(self, scale, loc)
+
+        self.smoothing_strength = smoothing_strength
+        self.kernel_size = kernel_size
+        self.same_for_all_channels = same_for_all_channels
+        self.p_per_channel = p_per_channel
+
+    def get_parameters(self, image: torch.Tensor, **kwargs) -> dict:
+        C, *spatial = image.shape
+        apply_channel = [np.random.rand() < self.p_per_channel for _ in range(C)]
+
+        if not any(apply_channel):
+            return {'kernels': [None] * C, 'sigma': None, 'strengths': [None] * C}
+
+        sigma = sample_scalar(self.kernel_size)
+
+        if self.same_for_all_channels:
+            kernel = self._generate_kernel(spatial).astype(np.float32)
+            strength = sample_scalar(self.smoothing_strength)
+
+            kernels = [kernel if apply else None for apply in apply_channel]
+            strengths = [strength if apply else None for apply in apply_channel]
+        else:
+            kernels, strengths = [], []
+            for apply in apply_channel:
+                if not apply:
+                    kernels.append(None)
+                    strengths.append(None)
+                    continue
+                kernel = self._generate_kernel(spatial).astype(np.float32)
+                strength = sample_scalar(self.smoothing_strength)
+                kernels.append(kernel)
+                strengths.append(strength)
+
+        return {'kernels': kernels, 'sigma': sigma, 'strengths': strengths}
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        img_np = img.cpu().numpy()
+        sigma = params['sigma']
+
+        for c, (kernel, strength) in enumerate(zip(params['kernels'], params['strengths'])):
+            if kernel is None:
+                continue
+
+            kernel = kernel * strength  # scale kernel by smoothing strength
+            smoothed = gaussian_filter(img_np[c], sigma=sigma)
+            img_np[c] = self.run_interpolation(img_np[c], smoothed, kernel)
+
+        return torch.from_numpy(img_np).to(img.device, dtype=img.dtype)
+
+
+if __name__ == '__main__':
+    from batchviewer import view_batch
+
+    # Single-channel synthetic volume
+    image = torch.rand(1, 32, 64, 64)
+
+    # Or contrast
+    smoother = LocalSmoothingTransform(loc=(0, 1), scale=(10, 20), kernel_size=(3, 10), p_per_channel=1.0)
+
+    # Apply either one
+    params = smoother.get_parameters(image=image)
+    image_aug = smoother._apply_to_image(image.clone(), **params)
+
+    view_batch(image, image_aug)
+

batchgeneratorsv2/transforms/local/local_transform.py

@@ -0,0 +1,86 @@
+import numpy as np
+import scipy.stats as st
+from abc import ABC
+from typing import Tuple
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+
+
+class LocalTransform(ABC):
+    def __init__(self, scale: RandomScalar, loc: RandomScalar = (-1, 2)):
+        self.loc = loc
+        self.scale = scale
+
+    def _generate_kernel(self, img_shp: Tuple[int, ...]) -> np.ndarray:
+        ndim = len(img_shp)
+        x_grids = [np.arange(-0.5, s + 0.5, dtype=np.float32) for s in img_shp]
+        kernels = []
+
+        for d in range(ndim):
+            loc_val = sample_scalar(self.loc, img_shp, d)
+            scale_val = sample_scalar(self.scale, img_shp, d)
+            loc_rescaled = loc_val * img_shp[d]
+            cdf = st.norm.cdf(x_grids[d], loc=loc_rescaled, scale=scale_val)
+            kernels.append(np.diff(cdf).astype(np.float32))
+
+        kernel = kernels[0][:, None] @ kernels[1][None]
+        if ndim == 3:
+            kernel = kernel[:, :, None] @ kernels[2][None]
+
+        kernel -= kernel.min()
+        kernel_max = kernel.max()
+        if kernel_max > 0:
+            kernel /= kernel_max
+        return kernel
+
+    def _generate_multiple_kernel_image(self, img_shp: Tuple[int, ...], num_kernels: int) -> np.ndarray:
+        """
+        Places multiple additive Gaussians in the image and normalizes the sum to [0, 1].
+
+        Parameters:
+            img_shp (Tuple[int, ...]): Spatial shape (e.g., (X, Y[, Z]))
+            num_kernels (int): Number of kernels to generate and sum
+
+        Returns:
+            np.ndarray: Combined kernel image with values in [0, 1]
+        """
+        kernel_image = np.zeros(img_shp, dtype=np.float32)
+        for _ in range(num_kernels):
+            kernel_image += self._generate_kernel(img_shp)
+
+        kernel_image -= kernel_image.min()
+        kernel_max = kernel_image.max()
+        if kernel_max > 0:
+            kernel_image /= kernel_max
+        return kernel_image
+
+    @staticmethod
+    def invert_kernel(kernel_image: np.ndarray) -> np.ndarray:
+        """
+        Inverts a normalized kernel: 1 - kernel
+
+        Assumes input is in [0, 1].
+
+        Parameters:
+            kernel_image (np.ndarray): Input kernel in [0, 1]
+
+        Returns:
+            np.ndarray: Inverted kernel in [0, 1]
+        """
+        return 1.0 - kernel_image
+
+    @staticmethod
+    def run_interpolation(original_image: np.ndarray,
+                          modified_image: np.ndarray,
+                          kernel_image: np.ndarray) -> np.ndarray:
+        """
+        Blends original and modified images using the given kernel as a per-pixel weight map.
+
+        Parameters:
+            original_image (np.ndarray): Unmodified input image
+            modified_image (np.ndarray): Modified version (e.g., gamma-corrected)
+            kernel_image (np.ndarray): Kernel in [0, 1], where 0 = keep original, 1 = keep modified
+
+        Returns:
+            np.ndarray: Blended result
+        """
+        return original_image * (1.0 - kernel_image) + modified_image * kernel_image

batchgeneratorsv2/transforms/nnunet/random_binary_operator.py

@@ -0,0 +1,190 @@
+from time import time
+from typing import Union, List, Tuple, Callable
+
+import numpy as np
+import torch
+from skimage.morphology import ball, disk
+
+from batchgeneratorsv2.helpers.fft_conv import fft_conv
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+def binary_dilation_torch(input_tensor, structure_element):
+    # Convert the boolean tensor to float
+    input_tensor = input_tensor.float()
+
+    # Get the number of dimensions of the input tensor
+    num_dims = input_tensor.dim()
+
+    # Prepare the structure element for convolution
+    # Adding extra dimensions to match the input shape for convolution
+    if num_dims == 2:  # For 2D inputs
+        structure_element = structure_element.unsqueeze(0).unsqueeze(0).float()
+    elif num_dims == 3:  # For 3D inputs, adding batch dimension
+        structure_element = structure_element.unsqueeze(0).unsqueeze(0).float()
+    else:
+        raise ValueError("Input tensor must be 2D (X, Y) or 3D (X, Y, Z).")
+
+    # Perform the convolution
+    # if num_dims == 2:  # 2D convolution
+    #     output = F.conv2d(input_tensor.unsqueeze(0).unsqueeze(0), structure_element, padding='same')
+    # elif num_dims == 3:  # 3D convolution
+    #     output = F.conv3d(input_tensor.unsqueeze(0).unsqueeze(0), structure_element, padding='same')
+    output = torch.round(fft_conv(input_tensor.unsqueeze(0).unsqueeze(0), structure_element, padding='same'), decimals=0)
+
+    # Threshold to get binary output
+    output = output > 0
+
+    # Squeeze the batch dimension out and convert to bool
+    return output.squeeze(0).squeeze(0).bool()
+
+
+def binary_erosion_torch(input_tensor, structure_element):
+    return ~binary_dilation_torch(~input_tensor, structure_element)
+
+
+def binary_opening_torch(input_tensor, structure_element):
+    return binary_dilation_torch(binary_erosion_torch(input_tensor, structure_element), structure_element)
+
+
+def binary_closing_torch(input_tensor, structure_element):
+    return binary_erosion_torch(binary_dilation_torch(input_tensor, structure_element), structure_element)
+
+
+class ApplyRandomBinaryOperatorTransform(ImageOnlyTransform):
+    def __init__(self,
+                 channel_idx: Union[int, List[int], Tuple[int, ...]],
+                 any_of_these: Tuple[Callable, ...] = (binary_dilation_torch, binary_erosion_torch, binary_closing_torch, binary_opening_torch),
+                 strel_size: RandomScalar = (1, 10),
+                 p_per_label: float = 1):
+        """
+        We use fft conv. Slower for small kernels but boi does it perform on larger kernels
+
+        Args:
+            channel_idx:
+            any_of_these:
+            strel_size:
+            p_per_label:
+        """
+        super().__init__()
+        if not isinstance(channel_idx, (list, tuple)):
+            channel_idx = [channel_idx]
+        if isinstance(channel_idx, tuple):
+            channel_idx = list(channel_idx)
+        self.channel_idx = channel_idx
+        self.any_of_these = any_of_these
+        self.strel_size = strel_size
+        self.p_per_label = p_per_label
+
+    def get_parameters(self, **data_dict) -> dict:
+        # this needs to be applied in random order to the channels
+        np.random.shuffle(self.channel_idx)
+        apply_to_channels = [self.channel_idx[i] for i, j in enumerate(torch.rand(len(self.channel_idx)) < self.p_per_label) if j]
+        operators = [np.random.choice(self.any_of_these) for _ in apply_to_channels]
+        strel_size = [sample_scalar(self.strel_size, image=data_dict['image'], channel=a) for a in apply_to_channels]
+        return {
+            'apply_to_channels': apply_to_channels,
+            'operators': operators,
+            'strel_size': strel_size,
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        for a, o, s in zip(params['apply_to_channels'], params['operators'], params['strel_size']):
+            # this is a binary map so bool is fine
+            workon = img[a]#.numpy()
+            orig_dtype = workon.dtype
+            workon = workon.to(bool)
+            if workon.ndim == 2:
+                strel = disk(s, dtype=bool)
+            else:
+                strel = ball(s, dtype=bool)
+            result = o(workon, torch.from_numpy(strel))
+            other_ch = [i for i in self.channel_idx if i != a]
+            if len(other_ch) > 0:
+                was_added_mask = result & (~workon)
+                for oc in other_ch:
+                    img[oc][was_added_mask] = 0
+            img[a] = result.to(orig_dtype)#torch.from_numpy(result)
+        return img
+
+
+# class ApplyRandomBinaryOperatorTransformNpy(ImageOnlyTransform):
+#     def __init__(self,
+#                  channel_idx: Union[int, List[int], Tuple[int, ...]],
+#                  any_of_these: Tuple[Callable, ...] = (binary_dilation, binary_erosion, binary_closing, binary_opening),
+#                  strel_size: ScalarType = (1, 10),
+#                  p_per_label: float = 1):
+#         """
+#         We stick to the nnunet implementation and dont optimize. This is a TODO for the future
+#
+#         Args:
+#             channel_idx:
+#             any_of_these:
+#             strel_size:
+#             p_per_label:
+#         """
+#         super().__init__()
+#         if not isinstance(channel_idx, (list, tuple)):
+#             channel_idx = [channel_idx]
+#         if isinstance(channel_idx, tuple):
+#             channel_idx = list(channel_idx)
+#         self.channel_idx = channel_idx
+#         self.any_of_these = any_of_these
+#         self.strel_size = strel_size
+#         self.p_per_label = p_per_label
+#
+#     def get_parameters(self, **data_dict) -> dict:
+#         # this needs to be applied in random order to the channels
+#         np.random.shuffle(self.channel_idx)
+#         apply_to_channels = [self.channel_idx[i] for i, j in enumerate(torch.rand(len(self.channel_idx)) < self.p_per_label) if j]
+#         operators = [np.random.choice(self.any_of_these) for _ in apply_to_channels]
+#         strel_size = [sample_scalar(self.strel_size, image=data_dict['image'], channel=a) for a in apply_to_channels]
+#         return {
+#             'apply_to_channels': apply_to_channels,
+#             'operators': operators,
+#             'strel_size': strel_size,
+#         }
+#
+#     def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+#         for a, o, s in zip(params['apply_to_channels'], params['operators'], params['strel_size']):
+#             # this is a binary map so bool is fine
+#             workon = img[a].numpy()
+#             orig_dtype = workon.dtype
+#             workon = workon.astype(bool)
+#             if workon.ndim == 2:
+#                 strel = disk(s, dtype=bool)
+#             else:
+#                 strel = ball(s, dtype=bool)
+#             result = o(workon, strel)
+#             other_ch = [i for i in self.channel_idx if i != a]
+#             if len(other_ch) > 0:
+#                 was_added_mask = result & (~workon)
+#                 for oc in other_ch:
+#                     img[oc][was_added_mask] = 0
+#             img[a] = torch.from_numpy(result)
+#         return img
+
+
+if __name__ == '__main__':
+    torch.set_num_threads(1)
+
+    tr = ApplyRandomBinaryOperatorTransform(
+        channel_idx=(0, 1),
+        strel_size=1,
+    )
+
+    times_torch = []
+    for _ in range(10):
+        # img = (torch.rand((1, 128, 128, 128)) < 0.5).to(torch.int16)
+        # img = torch.cat((img, 1 - img))
+        img = torch.zeros((1, 50, 50, 50))
+        img[0, :10, :10, :10] = 1
+        data_dict = {'image': torch.cat((img, 1 - img))}
+        st = time()
+        out = tr(**data_dict)
+        times_torch.append(time() - st)
+    print('torch', np.mean(times_torch))
+
+    # from batchviewer import view_batch
+    # view_batch(torch.cat((img, 1 - img)), out['image'], torch.cat((img, 1 - img)) - out['image'])

batchgeneratorsv2/transforms/nnunet/remove_connected_components.py

@@ -0,0 +1,86 @@
+from time import time
+from typing import Union, List, Tuple
+import numpy as np
+import torch
+from acvl_utils.morphology.morphology_helper import label_with_component_sizes
+
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+class RemoveRandomConnectedComponentFromOneHotEncodingTransform(ImageOnlyTransform):
+    def __init__(self,
+                 channel_idx: Union[int, List[int], Tuple[int, ...]],
+                 fill_with_other_class_p: float = 0.25,
+                 dont_do_if_covers_more_than_x_percent: float = 0.25,
+                 p_per_label: float = 1
+                 ):
+        super().__init__()
+        if not isinstance(channel_idx, (list, tuple)):
+            channel_idx = [channel_idx]
+        if isinstance(channel_idx, tuple):
+            channel_idx = list(channel_idx)
+        self.channel_idx = channel_idx
+        self.fill_with_other_class_p = fill_with_other_class_p
+        self.dont_do_if_covers_more_than_x_percent = dont_do_if_covers_more_than_x_percent
+        self.p_per_label = p_per_label
+
+    def get_parameters(self, **data_dict) -> dict:
+        # this needs to be applied in random order to the channels
+        np.random.shuffle(self.channel_idx)
+        apply_to_channels = [self.channel_idx[i] for i, j in enumerate(torch.rand(len(self.channel_idx)) < self.p_per_label) if j]
+
+        # self.fill_with_other_class_p cannot be resolved here because we don't know how many components there are
+        return {
+            'apply_to_channels': apply_to_channels,
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        for a in params['apply_to_channels']:
+            workon = img[a].to(bool).numpy()
+            if not np.any(workon):
+                continue
+            num_voxels = np.prod(workon.shape, dtype=np.uint64)
+            lab, component_sizes = label_with_component_sizes(workon.astype(bool))
+            if len(component_sizes) > 0:
+                valid_component_ids = [i for i, j in component_sizes.items() if j <
+                                       num_voxels * self.dont_do_if_covers_more_than_x_percent]
+                # print('RemoveRandomConnectedComponentFromOneHotEncodingTransform', c,
+                # np.unique(data[b, c]), len(component_sizes), valid_component_ids,
+                # len(valid_component_ids))
+                if len(valid_component_ids) > 0:
+                    random_component = np.random.choice(valid_component_ids)
+                    img[a][lab == random_component] = 0
+                    if np.random.uniform() < self.fill_with_other_class_p:
+                        other_ch = [i for i in self.channel_idx if i != a]
+                        if len(other_ch) > 0:
+                            other_class = np.random.choice(other_ch)
+                            img[other_class][lab == random_component] = 1
+        return img
+
+
+if __name__ == '__main__':
+    torch.set_num_threads(1)
+
+    tr = RemoveRandomConnectedComponentFromOneHotEncodingTransform(
+        channel_idx=(0, 1),
+        fill_with_other_class_p=1,
+        dont_do_if_covers_more_than_x_percent=0.25,
+        p_per_label=1
+    )
+
+    times_torch = []
+    for _ in range(10):
+        # img = (torch.rand((1, 128, 128, 128)) < 0.5).to(torch.int16)
+        # img = torch.cat((img, 1 - img))
+        img = torch.zeros((1, 64, 64, 64))
+        img[0, :10, :10, :10] = 1
+        img[0, -10:, -10:, -10:] = 1
+        img[0, -10:, :10, -10:] = 1
+        data_dict = {'image': torch.cat((img, 1 - img))}
+        st = time()
+        out = tr(**data_dict)
+        times_torch.append(time() - st)
+    print('torch', np.mean(times_torch))
+
+    from batchviewer import view_batch
+    view_batch(torch.cat((img, 1 - img)), out['image'], torch.cat((img, 1 - img)) - out['image'])