batchgeneratorsv2 0.2.3__tar.gz → 0.3.2__tar.gz

{batchgeneratorsv2-0.2.3 → batchgeneratorsv2-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: batchgeneratorsv2
-Version: 0.2.3
+Version: 0.3.2
 Summary: Batchgenerators but better
 Author: Helmholtz Imaging Applied Computer Vision Lab
 Author-email: Fabian Isensee <f.isensee@dkfz-heidelberg.de>
@@ -220,8 +220,8 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0.0
 Requires-Dist: numpy
-Requires-Dist: fft-conv-pytorch
 Requires-Dist: batchgenerators>=0.25
+Dynamic: license-file
 
 # batchgeneratorsv2
 This repository is work in progress. If builds upon the [batchgenerators](https://github.com/MIC-DKFZ/batchgenerators) 

batchgeneratorsv2-0.3.2/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-0.2.3 → batchgeneratorsv2-0.3.2}/batchgeneratorsv2/transforms/base/basic_transform.py

@@ -22,7 +22,7 @@ class BasicTransform(abc.ABC):
             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)
+            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)

batchgeneratorsv2-0.3.2/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))

{batchgeneratorsv2-0.2.3 → batchgeneratorsv2-0.3.2}/batchgeneratorsv2/transforms/intensity/contrast.py

@@ -25,44 +25,71 @@ class BGContrast():
     def __repr__(self):
         return self.__class__.__name__ + f"(contrast_range={self.contrast_range})"
 
+
+import torch
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
 class ContrastTransform(ImageOnlyTransform):
-    def __init__(self, contrast_range: RandomScalar, preserve_range: bool, synchronize_channels: bool, p_per_channel: float = 1):
+    def __init__(self, contrast_range: RandomScalar, preserve_range: bool, synchronize_channels: bool, p_per_channel: float = 1.0):
         super().__init__()
         self.contrast_range = contrast_range
         self.preserve_range = preserve_range
         self.synchronize_channels = synchronize_channels
-        self.p_per_channel = p_per_channel
+        self.p_per_channel = float(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))
+        img = data_dict["image"]
+        c = img.shape[0]
+
+        # sample on correct device
+        apply_idx = (torch.rand(c, device=img.device) < self.p_per_channel).nonzero(as_tuple=False).flatten()
+        n = apply_idx.numel()
+
+        if n == 0:
+            multipliers = None
+        elif self.synchronize_channels:
+            m = float(sample_scalar(self.contrast_range, image=img, channel=None))
+            multipliers = torch.full((n,), m, device=img.device, dtype=img.dtype)
         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
-        }
+            # Still a Python loop because sample_scalar is scalar-by-scalar
+            # Use .tolist() to avoid iterating tensor scalars in Python
+            ms = [sample_scalar(self.contrast_range, image=img, channel=int(ch)) for ch in apply_idx.tolist()]
+            multipliers = torch.as_tensor(ms, device=img.device, dtype=img.dtype)
+
+        return {"apply_to_channel": apply_idx, "multipliers": multipliers}
 
     def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
-        if len(params['apply_to_channel']) == 0:
+        idx = params["apply_to_channel"]
+        multipliers = params["multipliers"]
+        if multipliers is None or idx.numel() == 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)
+
+        if self.preserve_range:
+            for i in range(idx.numel()):
+                c = int(idx[i])
+                m = multipliers[i]
+
+                x = img[c]
+                mean = x.mean()
+                minm = x.min()
+                maxm = x.max()
+
+                x.sub_(mean)
+                x.mul_(m)
+                x.add_(mean)
+                x.clamp_(minm, maxm)
+        else:
+            for i in range(idx.numel()):
+                c = int(idx[i])
+                m = multipliers[i]
+
+                x = img[c]
+                mean = x.mean()
+                x.sub_(mean)
+                x.mul_(m)
+                x.add_(mean)
 
         return img
 

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

@@ -0,0 +1,135 @@
+from typing import Optional
+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 = float(p_invert_image)
+        self.synchronize_channels = synchronize_channels
+        self.p_per_channel = float(p_per_channel)
+        self.p_retain_stats = float(p_retain_stats)
+
+    def get_parameters(self, **data_dict) -> dict:
+        img: torch.Tensor = data_dict["image"]
+        c = img.shape[0]
+        device = img.device
+        dtype = img.dtype
+
+        apply_idx = (torch.rand(c, device=device) < self.p_per_channel).nonzero(as_tuple=False).flatten()
+        n = apply_idx.numel()
+        if n == 0:
+            return {"apply_to_channel": apply_idx,
+                    "retain_stats": None,
+                    "invert_image": None,
+                    "gamma": None}
+
+        retain_stats = (torch.rand(n, device=device) < self.p_retain_stats)
+        invert_image = (torch.rand(n, device=device) < self.p_invert_image)
+
+        if self.synchronize_channels:
+            g = float(sample_scalar(self.gamma, image=img, channel=None))
+            gamma = torch.full((n,), g, device=device, dtype=dtype)
+        else:
+            # sample_scalar is scalar-based; keep loop but avoid tensor scalar iteration
+            gs = [float(sample_scalar(self.gamma, image=img, channel=int(ch))) for ch in apply_idx.tolist()]
+            gamma = torch.as_tensor(gs, device=device, dtype=dtype)
+
+        return {
+            "apply_to_channel": apply_idx,
+            "retain_stats": retain_stats,
+            "invert_image": invert_image,
+            "gamma": gamma,
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        idx: torch.Tensor = params["apply_to_channel"]
+        if idx.numel() == 0:
+            return img
+
+        retain_stats: torch.Tensor = params["retain_stats"]
+        invert_image: torch.Tensor = params["invert_image"]
+        gamma: torch.Tensor = params["gamma"]
+
+        # constants
+        eps = 1e-7
+
+        # Loop over selected channels (good for small C)
+        for k in range(idx.numel()):
+            c = int(idx[k])
+            r = bool(retain_stats[k])
+            inv = bool(invert_image[k])
+            g = gamma[k]
+
+            x = img[c]
+
+            if inv:
+                x.mul_(-1)
+
+            if r:
+                mean = x.mean()
+                std = x.std()
+
+            minm = x.min()
+            maxm = x.max()
+            rnge = maxm - minm
+            denom = torch.clamp(rnge, min=eps)
+
+            # In-place gamma: x = (((x - min) / denom) ** g) * rnge + min
+            x.sub_(minm)
+            x.div_(denom)
+            x.pow_(g)
+            x.mul_(rnge)
+            x.add_(minm)
+
+            if r:
+                mn_here = x.mean()
+                std_here = x.std()
+                x.sub_(mn_here)
+                x.mul_(std / torch.clamp(std_here, min=eps))
+                x.add_(mean)
+
+            if inv:
+                x.mul_(-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))

batchgeneratorsv2-0.3.2/batchgeneratorsv2/transforms/intensity/gaussian_noise.py

@@ -0,0 +1,104 @@
+from typing import Tuple
+import torch
+
+from batchgeneratorsv2.helpers.scalar_type import RandomScalar, sample_scalar
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+class GaussianNoiseTransform(ImageOnlyTransform):
+    def __init__(self,
+                 noise_variance: RandomScalar = (0, 0.1),
+                 p_per_channel: float = 1.,
+                 synchronize_channels: bool = False):
+        super().__init__()
+        self.noise_variance = noise_variance
+        self.p_per_channel = p_per_channel
+        self.synchronize_channels = synchronize_channels
+
+    def get_parameters(self, **data_dict) -> dict:
+        img = data_dict["image"]
+        c = img.shape[0]
+
+        # bool mask on same device as image
+        apply = torch.rand(c, device=img.device) < self.p_per_channel
+
+        # store also count / indices to avoid recomputing later
+        idx = apply.nonzero(as_tuple=False).flatten()
+        n = idx.numel()
+
+        if n == 0:
+            sigmas = None
+        elif self.synchronize_channels:
+            sigmas = sample_scalar(self.noise_variance, img)
+        else:
+            # still uses sample_scalar, but avoids list->cat in _apply
+            # if sample_scalar is cheap, this is fine; otherwise see note below
+            sigmas = [sample_scalar(self.noise_variance, img) for _ in range(n)]
+
+        return {"apply_mask": apply, "apply_idx": idx, "num_apply": n, "sigmas": sigmas}
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        n = params["num_apply"]
+        if n == 0:
+            return img
+
+        idx = params["apply_idx"]
+        spatial = img.shape[1:]
+        device = img.device
+        dtype = img.dtype
+
+        sigmas = params["sigmas"]
+
+        if sigmas is None:
+            return img
+
+        # Create noise only for selected channels
+        if not self.synchronize_channels:
+            # vectorize per-channel sigma by creating a tensor of shape (n, 1, 1, ...)
+            # list->tensor is small (n floats), then broadcast
+            sigma_t = torch.as_tensor(sigmas, device=device, dtype=dtype)
+            view_shape = (n,) + (1,) * len(spatial)
+            sigma_t = sigma_t.view(view_shape)
+
+            noise = torch.empty((n, *spatial), device=device, dtype=dtype).normal_()
+            noise.mul_(sigma_t)
+        else:
+            sigma = sigmas
+            noise = torch.empty((n, *spatial), device=device, dtype=dtype).normal_(mean=0.0, std=float(sigma))
+
+        # Advanced indexing (img[idx]) returns a copy, so use indexed assignment
+        # to make sure modifications are written back to img.
+        img[idx] += noise
+        return img
+    
+
+if __name__ == "__main__":
+    from time import time
+    import numpy as np
+
+    os.environ['OMP_NUM_THREADS'] = '1'
+    torch.set_num_threads(1)
+
+    gnt = GaussianNoiseTransform((0, 0.1), 1, False)
+
+    times = []
+    for _ in range(1000):
+        data_dict = {'image': torch.ones((2, 32, 32, 32))}
+        st = time()
+        out = gnt(**data_dict)
+        times.append(time() - st)
+    print('torch', np.mean(times))
+
+    from batchgenerators.transforms.noise_transforms import GaussianNoiseTransform
+
+    gnt_bg = GaussianNoiseTransform((0, 0.1), 1, 1, True)
+
+    times = []
+    for _ in range(1000):
+        data_dict = {'data': np.ones((1, 2, 32, 32, 32))}
+        st = time()
+        out = gnt_bg(**data_dict)
+        times.append(time() - st)
+
+    print('bg', np.mean(times))
+    # torch is 2.5x faster