batchgeneratorsv2 0.1__tar.gz → 0.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: batchgeneratorsv2
-Version: 0.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.2/batchgeneratorsv2/benchmarks/bg_comparison/nnUNet_pipeline_bg.py

@@ -0,0 +1,90 @@
+from time import time
+
+import numpy as np
+import torch
+from batchgenerators.transforms.color_transforms import BrightnessMultiplicativeTransform, \
+    ContrastAugmentationTransform, GammaTransform
+from batchgenerators.transforms.noise_transforms import GaussianNoiseTransform, GaussianBlurTransform
+from batchgenerators.transforms.resample_transforms import SimulateLowResolutionTransform
+from batchgenerators.transforms.spatial_transforms import SpatialTransform, MirrorTransform
+from batchgenerators.transforms.utility_transforms import RemoveLabelTransform, NumpyToTensor
+from nnunetv2.training.data_augmentation.custom_transforms.deep_supervision_donwsampling import \
+    DownsampleSegForDSTransform2
+from nnunetv2.training.data_augmentation.custom_transforms.masking import MaskTransform
+from nnunetv2.training.data_augmentation.custom_transforms.region_based_training import \
+    ConvertSegmentationToRegionsTransform
+from nnunetv2.training.data_augmentation.custom_transforms.transforms_for_dummy_2d import Convert3DTo2DTransform, \
+    Convert2DTo3DTransform
+
+if __name__ == '__main__':
+    regions = ((1, 2, 3), (2, 3), (3, ))
+    do_dummy_2d_data_aug = False
+    patch_size = (128, 128, 128)
+    rotation_for_DA = (0, 2*np.pi)
+    deep_supervision_scales = ((1, 1, 1), (0.5, 0.5, 0.5), (0.25, 0.25, 0.25))
+
+    tr_transforms = []
+    if do_dummy_2d_data_aug:
+        ignore_axes = (0,)
+        tr_transforms.append(Convert3DTo2DTransform())
+        patch_size_spatial = patch_size[1:]
+    else:
+        patch_size_spatial = patch_size
+        ignore_axes = None
+
+    tr_transforms.append(SpatialTransform(
+        patch_size_spatial, patch_center_dist_from_border=None,
+        do_elastic_deform=False, alpha=(0, 0), sigma=(0, 0),
+        do_rotation=True, angle_x=rotation_for_DA, angle_y=rotation_for_DA, angle_z=rotation_for_DA,
+        p_rot_per_axis=1,  # todo experiment with this
+        do_scale=True, scale=(0.7, 1.4),
+        border_mode_data="constant", border_cval_data=0, order_data=3,
+        border_mode_seg="constant", border_cval_seg=-1, order_seg=1,
+        random_crop=False,  # random cropping is part of our dataloaders
+        p_el_per_sample=0, p_scale_per_sample=1, p_rot_per_sample=1,
+        independent_scale_for_each_axis=False  # todo experiment with this
+    ))
+
+    if do_dummy_2d_data_aug:
+        tr_transforms.append(Convert2DTo3DTransform())
+
+    tr_transforms.append(GaussianNoiseTransform(p_per_sample=1, p_per_channel=1))
+    tr_transforms.append(GaussianBlurTransform((0.5, 1.), different_sigma_per_channel=True, p_per_sample=1,
+                                               p_per_channel=1))
+    tr_transforms.append(BrightnessMultiplicativeTransform(multiplier_range=(0.75, 1.25), p_per_sample=1))
+    tr_transforms.append(ContrastAugmentationTransform(p_per_sample=1, p_per_channel=1))
+    tr_transforms.append(SimulateLowResolutionTransform(zoom_range=(0.5, 1), per_channel=True,
+                                                        p_per_channel=1,
+                                                        order_downsample=0, order_upsample=3, p_per_sample=1,
+                                                        ignore_axes=ignore_axes))
+    tr_transforms.append(GammaTransform((0.7, 1.5), True, True, retain_stats=True, p_per_sample=1))
+    tr_transforms.append(GammaTransform((0.7, 1.5), False, True, retain_stats=True, p_per_sample=1))
+
+    tr_transforms.append(MirrorTransform((0, 1, 2)))
+
+    tr_transforms.append(MaskTransform([0, 1, 2, 3],
+                                       mask_idx_in_seg=0, set_outside_to=0))
+
+    tr_transforms.append(RemoveLabelTransform(-1, 0))
+
+    tr_transforms.append(ConvertSegmentationToRegionsTransform(regions, 'seg', 'seg'))
+
+    tr_transforms.append(DownsampleSegForDSTransform2(deep_supervision_scales, 0, input_key='seg',
+                                                      output_key='seg'))
+
+    tr_transforms.append(NumpyToTensor(['data', 'seg'], 'float'))
+
+    compute_times = [[] for i in range(len(tr_transforms))]
+
+    torch.set_num_threads(1)
+    for iter in range(50):
+        print(iter)
+        data_dict = {'data': np.random.uniform(size=(1, 4, 128, 128, 128)),
+                     'seg': np.round(4.5 * np.random.uniform(size=(1, 1, 128, 128, 128)) - 1, decimals=0).astype(np.int8)}
+        for i, t in enumerate(tr_transforms):
+            st = time()
+            data_dict = t(**data_dict)
+            compute_times[i].append(time() - st)
+
+    for t, ct in zip(tr_transforms, compute_times):
+        print(t.__class__.__name__, np.mean(ct))

batchgeneratorsv2-0.2/batchgeneratorsv2/benchmarks/bg_comparison/nnUNet_pipeline_here.py

@@ -0,0 +1,138 @@
+from time import time
+
+import numpy as np
+import torch
+
+from batchgeneratorsv2.transforms.intensity.brightness import MultiplicativeBrightnessTransform
+from batchgeneratorsv2.transforms.intensity.contrast import BGContrast, ContrastTransform
+from batchgeneratorsv2.transforms.intensity.gamma import GammaTransform
+from batchgeneratorsv2.transforms.intensity.gaussian_noise import GaussianNoiseTransform
+from batchgeneratorsv2.transforms.noise.gaussian_blur import GaussianBlurTransform
+from batchgeneratorsv2.transforms.spatial.low_resolution import SimulateLowResolutionTransform
+from batchgeneratorsv2.transforms.spatial.mirroring import MirrorTransform
+from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform
+from batchgeneratorsv2.transforms.utils.compose import ComposeTransforms
+from batchgeneratorsv2.transforms.utils.deep_supervision_downsampling import DownsampleSegForDSTransform
+from batchgeneratorsv2.transforms.utils.nnunet_masking import MaskImageTransform
+from batchgeneratorsv2.transforms.utils.pseudo2d import Convert2DTo3DTransform, Convert3DTo2DTransform
+from batchgeneratorsv2.transforms.utils.random import RandomTransform
+from batchgeneratorsv2.transforms.utils.remove_label import RemoveLabelTansform
+from batchgeneratorsv2.transforms.utils.seg_to_regions import ConvertSegmentationToRegionsTransform
+
+if __name__ == '__main__':
+    regions = ((1, 2, 3), (2, 3), (3, ))
+    do_dummy_2d_data_aug = False
+    patch_size = (128, 128, 128)
+    rotation_for_DA = (0, 2*np.pi)
+    deep_supervision_scales = ((1, 1, 1), (0.5, 0.5, 0.5), (0.25, 0.25, 0.25))
+
+    transforms = []
+    if do_dummy_2d_data_aug:
+        ignore_axes = (0,)
+        transforms.append(Convert3DTo2DTransform())
+        patch_size_spatial = patch_size[1:]
+    else:
+        patch_size_spatial = patch_size
+        ignore_axes = None
+    transforms.append(
+        SpatialTransform(
+            patch_size_spatial, patch_center_dist_from_border=0, random_crop=False, p_elastic_deform=0,
+            p_rotation=1,
+            rotation=rotation_for_DA, p_scaling=1, scaling=(0.7, 1.4), p_synchronize_scaling_across_axes=1
+        )
+    )
+    if do_dummy_2d_data_aug:
+        transforms.append(Convert2DTo3DTransform())
+
+    transforms.append(
+        GaussianNoiseTransform(
+            noise_variance=(0, 0.1),
+            p_per_channel=1,
+            synchronize_channels=True
+        )
+    )
+    transforms.append(
+        GaussianBlurTransform(
+            blur_sigma=(0.5, 1.),
+            synchronize_channels=False,
+            synchronize_axes=False,
+            p_per_channel=1, benchmark=True
+        ))
+    transforms.append(
+        MultiplicativeBrightnessTransform(
+            multiplier_range=BGContrast((0.75, 1.25)),
+            synchronize_channels=False,
+            p_per_channel=1
+        ))
+    transforms.append(
+        ContrastTransform(
+            contrast_range=BGContrast((0.75, 1.25)),
+            preserve_range=True,
+            synchronize_channels=False,
+            p_per_channel=1
+        ))
+    transforms.append(
+        SimulateLowResolutionTransform(
+            scale=(0.5, 1),
+            synchronize_channels=False,
+            synchronize_axes=True,
+            ignore_axes=ignore_axes,
+            allowed_channels=None,
+            p_per_channel=1
+        ))
+    transforms.append(
+        GammaTransform(
+            gamma=BGContrast((0.7, 1.5)),
+            p_invert_image=1,
+            synchronize_channels=False,
+            p_per_channel=1,
+            p_retain_stats=1
+        ))
+    transforms.append(
+        GammaTransform(
+            gamma=BGContrast((0.7, 1.5)),
+            p_invert_image=0,
+            synchronize_channels=False,
+            p_per_channel=1,
+            p_retain_stats=1
+        ))
+    transforms.append(
+        MirrorTransform(
+            allowed_axes=(0, 1, 2)
+        )
+    )
+
+    transforms.append(MaskImageTransform(
+        apply_to_channels=[0, 1, 2, 3],
+        channel_idx_in_seg=0,
+        set_outside_to=0,
+    ))
+
+    transforms.append(
+        RemoveLabelTansform(-1, 0)
+    )
+
+    transforms.append(
+        ConvertSegmentationToRegionsTransform(
+            regions=regions,
+            channel_in_seg=0
+        )
+    )
+
+    transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales))
+
+    compute_times = [[] for i in range(len(transforms))]
+
+    with torch.no_grad():
+        torch.set_num_threads(1)
+        for iter in range(50):
+            print(iter)
+            data_dict = {'image': torch.rand((4, 128, 128, 128)),
+                         'segmentation': torch.round(4.5 * torch.rand((1, 128, 128, 128)) - 1, decimals=0).to(torch.int8)}
+            for i, t in enumerate(transforms):
+                st = time()
+                data_dict = t(**data_dict)
+                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.mean(ct))

batchgeneratorsv2-0.2/batchgeneratorsv2/benchmarks/unique_values.py

@@ -0,0 +1,55 @@
+import torch
+import numpy as np
+from time import time
+import pandas as pd
+
+def unique_torch(tensor):
+    return torch.unique(tensor)
+
+def unique_npy(tensor):
+    return np.unique(tensor.numpy())
+
+def unique_pandas(tensor):
+    np.sort(pd.unique(tensor.numpy().ravel()))
+
+def unique_bincount(tensor):
+    return torch.where(torch.bincount(tensor.ravel()) > 0)[0]
+
+
+if __name__ == '__main__':
+    torch.set_num_threads(1)
+    shape = (64, 64, 64)
+    labels = 200
+
+    times = []
+    for _ in range(10):
+        seg = torch.round(torch.rand(shape) * 20, decimals=0).to(torch.uint8)
+        st = time()
+        unique = unique_torch(seg)
+        times.append(time() - st)
+    print('unique_torch', np.median(times))
+
+    times = []
+    for _ in range(10):
+        seg = torch.round(torch.rand(shape) * 20, decimals=0).to(torch.uint8)
+        st = time()
+        unique = unique_npy(seg)
+        times.append(time() - st)
+    print('unique_npy', np.median(times))
+
+    times = []
+    for _ in range(10):
+        seg = torch.round(torch.rand(shape) * 20, decimals=0).to(torch.uint8)
+        st = time()
+        unique = unique_pandas(seg)
+        times.append(time() - st)
+    print('unique_pandas', np.median(times))
+
+    times = []
+    for _ in range(10):
+        seg = torch.round(torch.rand(shape) * 20, decimals=0).to(torch.uint8)
+        st = time()
+        unique = unique_bincount(seg)
+        times.append(time() - st)
+    print('unique_bincount', np.median(times))
+

batchgeneratorsv2-0.2/batchgeneratorsv2/helpers/scalar_type.py

@@ -0,0 +1,28 @@
+from typing import Union, Tuple, Callable
+import numpy as np
+
+
+RandomScalar = Union[int, float, Tuple[float, float], Callable[..., Union[int, float]]]
+
+
+def sample_scalar(scalar_type: RandomScalar, *args, **kwargs):
+    if isinstance(scalar_type, (int, float)):
+        return scalar_type
+    elif isinstance(scalar_type, (list, tuple)):
+        assert len(scalar_type) == 2, 'if list is provided, its length must be 2'
+        assert scalar_type[0] <= scalar_type[1], 'if list is provided, first entry must be smaller or equal than second entry, ' \
+                                                'otherwise we cannot sample using np.random.uniform'
+        if scalar_type[0] == scalar_type[1]:
+            return scalar_type[0]
+        return np.random.uniform(*scalar_type)
+    elif callable(scalar_type):
+        return scalar_type(*args, **kwargs)
+    else:
+        raise RuntimeError('Unknown type: %s. Expected: int, float, list, tuple, callable', type(scalar_type))
+
+
+if __name__ == '__main__':
+    sample_scalar(0.5)
+    sample_scalar((0, 1))
+    sample_scalar(lambda: np.random.uniform(-1, 2))
+    sample_scalar(lambda x, y: np.random.uniform(x, y), 0.5, 2)

batchgeneratorsv2-0.2/batchgeneratorsv2/transforms/base/basic_transform.py

@@ -0,0 +1,77 @@
+import abc
+import torch
+
+
+class BasicTransform(abc.ABC):
+    """
+    Transforms are applied to each sample individually. The dataloader is responsible for collating, or we might consider a CollateTransform
+
+    We expect (C, X, Y) or (C, X, Y, Z) shaped inputs for image and seg (yes seg can have more color channels)
+
+    No idea what keypoint and bbox will look like, this is Michaels turf
+    """
+    def __init__(self):
+        pass
+
+    def __call__(self, **data_dict) -> dict:
+        params = self.get_parameters(**data_dict)
+        return self.apply(data_dict, **params)
+
+    def apply(self, data_dict, **params):
+        if data_dict.get('image') is not None:
+            data_dict['image'] = self._apply_to_image(data_dict['image'], **params)
+
+        if data_dict.get('regression_target') is not None:
+            data_dict['regression_target'] = self._apply_to_segmentation(data_dict['regression_target'], **params)
+
+        if data_dict.get('segmentation') is not None:
+            data_dict['segmentation'] = self._apply_to_segmentation(data_dict['segmentation'], **params)
+
+        if data_dict.get('keypoints') is not None:
+            data_dict['keypoints'] = self._apply_to_keypoints(data_dict['keypoints'], **params)
+
+        if data_dict.get('bbox') is not None:
+            data_dict['bbox'] = self._apply_to_bbox(data_dict['bbox'], **params)
+
+        return data_dict
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        pass
+
+    def _apply_to_regr_target(self, regression_target, **params) -> torch.Tensor:
+        pass
+
+    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
+        pass
+
+    def _apply_to_keypoints(self, keypoints, **params):
+        pass
+
+    def _apply_to_bbox(self, bbox, **params):
+        pass
+
+    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:
+        if data_dict.get('image') is not None:
+            data_dict['image'] = self._apply_to_image(data_dict['image'], **params)
+        return data_dict
+
+
+class SegOnlyTransform(BasicTransform):
+    def apply(self, data_dict: dict, **params) -> dict:
+        if data_dict.get('segmentation') is not None:
+            data_dict['segmentation'] = self._apply_to_segmentation(data_dict['segmentation'], **params)
+        return data_dict
+
+
+if __name__ == '__main__':
+    pass

batchgeneratorsv2-0.2/batchgeneratorsv2/transforms/intensity/brightness.py

@@ -0,0 +1,63 @@
+import torch
+
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+class MultiplicativeBrightnessTransform(ImageOnlyTransform):
+    def __init__(self, multiplier_range: RandomScalar, synchronize_channels: bool, p_per_channel: float = 1):
+        super().__init__()
+        self.multiplier_range = multiplier_range
+        self.synchronize_channels = synchronize_channels
+        self.p_per_channel = p_per_channel
+
+    def get_parameters(self, **data_dict) -> dict:
+        shape = data_dict['image'].shape
+        apply_to_channel = torch.where(torch.rand(shape[0]) < self.p_per_channel)[0]
+        if self.synchronize_channels:
+            multipliers = torch.Tensor([sample_scalar(self.multiplier_range, image=data_dict['image'], channel=None)] * len(apply_to_channel))
+        else:
+            multipliers = torch.Tensor([sample_scalar(self.multiplier_range, image=data_dict['image'], channel=c) for c in apply_to_channel])
+        return {
+            'apply_to_channel': apply_to_channel,
+            'multipliers': multipliers
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        if len(params['apply_to_channel']) == 0:
+            return img
+        # even though this is array notation it's a lot slower. Shame shame
+        # img[params['apply_to_channel']] *= params['multipliers'].view(-1, *[1]*(img.ndim - 1))
+        for c, m in zip(params['apply_to_channel'], params['multipliers']):
+            img[c] *= m
+        return img
+
+
+if __name__ == '__main__':
+    from time import time
+    import numpy as np
+    import os
+
+    os.environ['OMP_NUM_THREADS'] = '1'
+    torch.set_num_threads(1)
+
+    mbt = MultiplicativeBrightnessTransform((0.5, 2.), False, 1)
+
+    times_torch = []
+    for _ in range(1000):
+        data_dict = {'image': torch.ones((2, 128, 192, 64))}
+        st = time()
+        out = mbt(**data_dict)
+        times_torch.append(time() - st)
+    print('torch', np.mean(times_torch))
+
+    from batchgenerators.transforms.color_transforms import BrightnessMultiplicativeTransform
+
+    gnt_bg = BrightnessMultiplicativeTransform((0.5, 2), True, p_per_sample=1)
+    times_bg = []
+    for _ in range(1000):
+        data_dict = {'data': np.ones((1, 2, 128, 192, 64))}
+        st = time()
+        out = gnt_bg(**data_dict)
+        times_bg.append(time() - st)
+    print('bg', np.mean(times_bg))

batchgeneratorsv2-0.2/batchgeneratorsv2/transforms/intensity/contrast.py

@@ -0,0 +1,96 @@
+import torch
+
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+import numpy as np
+
+
+class BGContrast():
+    def __init__(self, contrast_range):
+        self.contrast_range = contrast_range
+
+    def sample_contrast(self, *args, **kwargs):
+        if callable(self.contrast_range):
+            factor = self.contrast_range()
+        else:
+            if np.random.random() < 0.5 and self.contrast_range[0] < 1:
+                factor = np.random.uniform(self.contrast_range[0], 1)
+            else:
+                factor = np.random.uniform(max(self.contrast_range[0], 1), self.contrast_range[1])
+        return factor
+
+    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__()
+        self.contrast_range = contrast_range
+        self.preserve_range = preserve_range
+        self.synchronize_channels = synchronize_channels
+        self.p_per_channel = p_per_channel
+
+    def get_parameters(self, **data_dict) -> dict:
+        shape = data_dict['image'].shape
+        apply_to_channel = torch.where(torch.rand(shape[0]) < self.p_per_channel)[0]
+        if self.synchronize_channels:
+            multipliers = torch.Tensor([sample_scalar(self.contrast_range, image=data_dict['image'], channel=None)] * len(apply_to_channel))
+        else:
+            multipliers = torch.Tensor([sample_scalar(self.contrast_range, image=data_dict['image'], channel=c) for c in apply_to_channel])
+        return {
+            'apply_to_channel': apply_to_channel,
+            'multipliers': multipliers
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        if len(params['apply_to_channel']) == 0:
+            return img
+        # array notation is not faster, let's leave it like this
+        for i in range(len(params['apply_to_channel'])):
+            c = params['apply_to_channel'][i]
+            mean = img[c].mean()
+            if self.preserve_range:
+                minm = img[c].min()
+                maxm = img[c].max()
+
+            # this is faster than having it in one line because this circumvents reallocating memory
+            img[c] -= mean
+            img[c] *= params['multipliers'][i]
+            img[c] += mean
+
+            if self.preserve_range:
+                img[c].clamp_(minm, maxm)
+
+        return img
+
+
+if __name__ == '__main__':
+    from time import time
+    import os
+
+    os.environ['OMP_NUM_THREADS'] = '1'
+    torch.set_num_threads(1)
+
+    mbt = ContrastTransform(BGContrast((0.75, 1.25)).sample_contrast, True, False, p_per_channel=1)
+
+    times_torch = []
+    for _ in range(100):
+        data_dict = {'image': torch.ones((2, 128, 192, 64))}
+        st = time()
+        out = mbt(**data_dict)
+        times_torch.append(time() - st)
+    print('torch', np.mean(times_torch))
+
+    from batchgenerators.transforms.color_transforms import ContrastAugmentationTransform
+
+    gnt_bg = ContrastAugmentationTransform((0.75, 1.25), preserve_range=True, per_channel=True, p_per_channel=1)
+    times_bg = []
+    for _ in range(100):
+        data_dict = {'data': np.ones((1, 2, 128, 192, 64))}
+        st = time()
+        out = gnt_bg(**data_dict)
+        times_bg.append(time() - st)
+    print('bg', np.mean(times_bg))

batchgeneratorsv2-0.2/batchgeneratorsv2/transforms/intensity/gamma.py

@@ -0,0 +1,88 @@
+from typing import Callable, Union
+
+import torch
+
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+class GammaTransform(ImageOnlyTransform):
+    def __init__(self, gamma: RandomScalar, p_invert_image: float, synchronize_channels: bool, p_per_channel: float,
+                 p_retain_stats: float):
+        super().__init__()
+        self.gamma = gamma
+        self.p_invert_image = p_invert_image
+        self.synchronize_channels = synchronize_channels
+        self.p_per_channel = p_per_channel
+        self.p_retain_stats = p_retain_stats
+
+    def get_parameters(self, **data_dict) -> dict:
+        shape = data_dict['image'].shape
+        apply_to_channel = torch.where(torch.rand(shape[0]) < self.p_per_channel)[0]
+        retain_stats = torch.rand(len(apply_to_channel)) < self.p_retain_stats
+        invert_image = torch.rand(len(apply_to_channel)) < self.p_invert_image
+
+        if self.synchronize_channels:
+            gamma = torch.Tensor([sample_scalar(self.gamma, image=data_dict['image'], channel=None)] * len(apply_to_channel))
+        else:
+            gamma = torch.Tensor([sample_scalar(self.gamma, image=data_dict['image'], channel=c) for c in apply_to_channel])
+        return {
+            'apply_to_channel': apply_to_channel,
+            'retain_stats': retain_stats,
+            'invert_image': invert_image,
+            'gamma': gamma
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        for c, r, i, g in zip(params['apply_to_channel'], params['retain_stats'], params['invert_image'], params['gamma']):
+            if i:
+                img[c] *= -1
+            if r:
+                # std_mean is for whatever reason slower than doing the computations separately!?
+                # std, mean = torch.std_mean(img[c])
+                mean = torch.mean(img[c])
+                std = torch.std(img[c])
+            minm = torch.min(img[c])
+            rnge = torch.max(img[c]) - minm
+            img[c] = torch.pow(((img[c] - minm) / torch.clamp(rnge, min=1e-7)), g) * rnge + minm
+            if r:
+                # std_here, mn_here = torch.std_mean(img[c])
+                mn_here = torch.mean(img[c])
+                std_here = torch.std(img[c])
+                img[c] -= mn_here
+                img[c] *= (std / torch.clamp(std_here, min=1e-7))
+                img[c] += mean
+
+            if i:
+                img[c] *= -1
+        return img
+
+
+if __name__ == '__main__':
+    from time import time
+    import numpy as np
+    import os
+
+    os.environ['OMP_NUM_THREADS'] = '1'
+    torch.set_num_threads(1)
+
+    mbt = GammaTransform((0.7, 1.5), 0, False, 1, 1)
+
+    times_torch = []
+    for _ in range(100):
+        data_dict = {'image': torch.ones((2, 128, 192, 64))}
+        st = time()
+        out = mbt(**data_dict)
+        times_torch.append(time() - st)
+    print('torch', np.mean(times_torch))
+
+    from batchgenerators.transforms.color_transforms import GammaTransform as BGGamma
+
+    gnt_bg = BGGamma((0.7, 1.5), False, True, retain_stats=True, p_per_sample=1)
+    times_bg = []
+    for _ in range(100):
+        data_dict = {'data': np.ones((1, 2, 128, 192, 64))}
+        st = time()
+        out = gnt_bg(**data_dict)
+        times_bg.append(time() - st)
+    print('bg', np.mean(times_bg))