batchgeneratorsv2 0.3.2__py3-none-any.whl

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/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/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/helpers/fft_conv.py

@@ -0,0 +1,149 @@
+from typing import Iterable, Tuple, Union
+
+import torch
+import torch.nn.functional as f
+from torch import Tensor, nn
+from torch.fft import irfftn, rfftn
+from math import ceil, floor
+
+
+# Taken from here: https://github.com/vcasellesb/fft-conv-pytorch/tree/non-tup-slice-fix. THANK YOU!
+# Original codebase is here: https://github.com/fkodom/fft-conv-pytorch -> unfortunately was not updated for pytorch 2.9.
+
+
+def complex_matmul(a: Tensor, b: Tensor, groups: int = 1) -> Tensor:
+    """Multiplies two complex-valued tensors."""
+    # Scalar matrix multiplication of two tensors, over only the first channel
+    # dimensions. Dimensions 3 and higher will have the same shape after multiplication.
+    # We also allow for "grouped" multiplications, where multiple sections of channels
+    # are multiplied independently of one another (required for group convolutions).
+    a = a.view(a.size(0), groups, -1, *a.shape[2:])
+    b = b.view(groups, -1, *b.shape[1:])
+
+    a = torch.movedim(a, 2, a.dim() - 1).unsqueeze(-2)
+    b = torch.movedim(b, (1, 2), (b.dim() - 1, b.dim() - 2))
+
+    # complex value matrix multiplication
+    real = a.real @ b.real - a.imag @ b.imag
+    imag = a.imag @ b.real + a.real @ b.imag
+    real = torch.movedim(real, real.dim() - 1, 2).squeeze(-1)
+    imag = torch.movedim(imag, imag.dim() - 1, 2).squeeze(-1)
+    c = torch.zeros(real.shape, dtype=torch.complex64, device=a.device)
+    c.real, c.imag = real, imag
+
+    return c.view(c.size(0), -1, *c.shape[3:])
+
+
+def to_ntuple(val: Union[int, Iterable[int]], n: int) -> Tuple[int, ...]:
+    """Casts to a tuple with length 'n'.  Useful for automatically computing the
+    padding and stride for convolutions, where users may only provide an integer.
+
+    Args:
+        val: (Union[int, Iterable[int]]) Value to cast into a tuple.
+        n: (int) Desired length of the tuple
+
+    Returns:
+        (Tuple[int, ...]) Tuple of length 'n'
+    """
+    if isinstance(val, Iterable):
+        out = tuple(val)
+        if len(out) == n:
+            return out
+        else:
+            raise ValueError(f"Cannot cast tuple of length {len(out)} to length {n}.")
+    else:
+        return n * (val,)
+
+
+def fft_conv(
+    signal: Tensor,
+    kernel: Tensor,
+    bias: Tensor = None,
+    padding: Union[int, Iterable[int], str] = 0,
+    padding_mode: str = "constant",
+    stride: Union[int, Iterable[int]] = 1,
+    dilation: Union[int, Iterable[int]] = 1,
+    groups: int = 1,
+) -> Tensor:
+    """Performs N-d convolution of Tensors using a fast fourier transform, which
+    is very fast for large kernel sizes. Also, optionally adds a bias Tensor after
+    the convolution (in order ot mimic the PyTorch direct convolution).
+
+    Args:
+        signal: (Tensor) Input tensor to be convolved with the kernel.
+        kernel: (Tensor) Convolution kernel.
+        bias: (Tensor) Bias tensor to add to the output.
+        padding: (Union[int, Iterable[int], str) If int, Number of zero samples to pad then
+            input on the last dimension. If str, "same" supported to pad input for size preservation.
+        padding_mode: (str) Padding mode to use from {constant, reflection, replication}.
+                      reflection not available for 3d.
+        stride: (Union[int, Iterable[int]) Stride size for computing output values.
+        dilation: (Union[int, Iterable[int]) Dilation rate for the kernel.
+        groups: (int) Number of groups for the convolution.
+
+    Returns:
+        (Tensor) Convolved tensor
+    """
+
+    # Cast padding, stride & dilation to tuples.
+    n = signal.ndim - 2
+    stride_ = to_ntuple(stride, n=n)
+    dilation_ = to_ntuple(dilation, n=n)
+    if isinstance(padding, str):
+        if padding == "same":
+            if stride != 1 or dilation != 1:
+                raise ValueError("stride must be 1 for padding='same'.")
+            padding_ = [(k - 1) / 2 for k in kernel.shape[2:]]
+        # else:
+        #     raise ValueError(f"Padding mode {padding} not supported.")
+    else:
+        padding_ = to_ntuple(padding, n=n)
+
+    # internal dilation offsets
+    offset = torch.zeros(1, 1, *dilation_, device=signal.device, dtype=signal.dtype)
+    offset[(slice(None), slice(None), *((0,) * n))] = 1.0
+
+    # correct the kernel by cutting off unwanted dilation trailing zeros
+    cutoff = tuple(slice(None, -d + 1 if d != 1 else None) for d in dilation_)
+
+    # pad the kernel internally according to the dilation parameters
+    kernel = torch.kron(kernel, offset)[(slice(None), slice(None)) + cutoff]
+
+    # Pad the input signal & kernel tensors (round to support even sized convolutions)
+    signal_padding = [r(p) for p in padding_[::-1] for r in (floor, ceil)]
+    signal = f.pad(signal, signal_padding, mode=padding_mode)
+
+    # Because PyTorch computes a *one-sided* FFT, we need the final dimension to
+    # have *even* length.  Just pad with one more zero if the final dimension is odd.
+    signal_size = signal.size()  # original signal size without padding to even
+    if signal.size(-1) % 2 != 0:
+        signal = f.pad(signal, [0, 1])
+
+    kernel_padding = [
+        pad
+        for i in reversed(range(2, signal.ndim))
+        for pad in [0, signal.size(i) - kernel.size(i)]
+    ]
+    padded_kernel = f.pad(kernel, kernel_padding)
+
+    # Perform fourier convolution -- FFT, matrix multiply, then IFFT
+    signal_fr = rfftn(signal.float(), dim=tuple(range(2, signal.ndim)))
+    kernel_fr = rfftn(padded_kernel.float(), dim=tuple(range(2, signal.ndim)))
+
+    kernel_fr.imag *= -1
+    output_fr = complex_matmul(signal_fr, kernel_fr, groups=groups)
+    output = irfftn(output_fr, dim=tuple(range(2, signal.ndim)))
+
+    # Remove extra padded values
+    crop_slices = (slice(None), slice(None)) + tuple(
+        slice(0, (signal_size[i] - kernel.size(i) + 1), stride_[i - 2])
+        for i in range(2, signal.ndim)
+    )
+    output = output[crop_slices].contiguous()
+
+    # Optionally, add a bias term before returning.
+    if bias is not None:
+        bias_shape = tuple([1, -1] + (signal.ndim - 2) * [1])
+        output += bias.view(bias_shape)
+
+    return output

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/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_regr_target(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/transforms/intensity/brightness.py

@@ -0,0 +1,123 @@
+import numpy as np
+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
+
+
+class BrightnessAdditiveTransform(ImageOnlyTransform):
+    """
+    Adds random additive brightness noise sampled from a Gaussian distribution (mu, sigma).
+
+    Supports either synchronized brightness shift across all channels or per-channel brightness shift.
+
+    Args:
+        mu (float): Mean of the Gaussian used to sample brightness shifts.
+        sigma (float): Standard deviation of the Gaussian.
+        synchronize_channels (bool): If True, brightness shifts are shared across all channels.
+        p_per_channel (float): Probability to apply the brightness shift to each channel.
+    """
+
+    def __init__(self,
+                 mu: float,
+                 sigma: float,
+                 synchronize_channels: bool = True,  # Changed to synchronize_channels
+                 p_per_channel: float = 1.0):
+        super().__init__()
+        self.mu = mu
+        self.sigma = sigma
+        self.synchronize_channels = synchronize_channels  # Now it's being used
+        self.p_per_channel = p_per_channel
+
+    def get_parameters(self, **data_dict) -> dict:
+        img = data_dict["image"]
+        c = img.shape[0]
+        apply_to_channel = (torch.rand(c, device=img.device) < self.p_per_channel).nonzero(as_tuple=False).flatten()
+
+        if len(apply_to_channel) == 0:
+            return {"apply_to_channel": apply_to_channel, "shift": None}
+
+        # Apply either synchronized or per-channel brightness shift
+        if self.synchronize_channels:
+            shift_value = float(sample_scalar((self.mu, self.sigma), image=img, channel=None))
+            shift = torch.full((c,), shift_value, device=img.device)
+        else:
+            shift = torch.empty(c, device=img.device).normal_(float(self.mu), float(self.sigma))
+
+        return {"apply_to_channel": apply_to_channel, "shift": shift}
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        if params["shift"] is None:
+            return img
+
+        apply_idx = params["apply_to_channel"]
+        if apply_idx.numel() == 0:
+            return img
+
+        shift = params["shift"]
+        # Build full per-channel shift vector; non-selected channels get shift 0
+        shift_full = torch.zeros((img.shape[0],), device=img.device, dtype=img.dtype)
+        shift_full[apply_idx] = shift[apply_idx]
+
+        view_shape = (img.shape[0],) + (1,) * (img.ndim - 1)
+        img.add_(shift_full.view(view_shape))
+        return img
+
+
+if __name__ == '__main__':
+    from time import time
+    import os
+
+    os.environ['OMP_NUM_THREADS'] = '1'
+    torch.set_num_threads(1)
+
+    # mbt = BrightnessAdditiveTransform(0, 0.5,True, 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))