batchgeneratorsv2 0.3.2__py3-none-any.whl

batchgeneratorsv2/transforms/nnunet/seg_to_onehot.py

@@ -0,0 +1,32 @@
+from typing import Union, List, Tuple
+
+import torch
+
+from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
+
+
+class MoveSegAsOneHotToDataTransform(BasicTransform):
+    def __init__(self, source_channel_idx: int, all_labels: Union[Tuple[int, ...], List[int]],
+                 remove_channel_from_source: bool = True):
+        """
+        Used in nnU-Net to append segmentations from the previous stage to the image as additional input
+        Args:
+            source_channel_idx:
+            all_labels:
+            remove_channel_from_source:
+        """
+        super().__init__()
+        self.source_channel_idx = source_channel_idx
+        self.all_labels = all_labels
+        self.remove_channel_from_source = remove_channel_from_source
+
+    def apply(self, data_dict, **params):
+        seg = data_dict['segmentation'][self.source_channel_idx]
+        seg_onehot = torch.zeros((len(self.all_labels), *seg.shape), dtype=data_dict['image'].dtype)
+        for i, l in enumerate(self.all_labels):
+            seg_onehot[i][seg == l] = 1
+        data_dict['image'] = torch.cat((data_dict['image'], seg_onehot))
+        if self.remove_channel_from_source:
+            remaining_channels = [i for i in range(data_dict['segmentation'].shape[0]) if i != self.source_channel_idx]
+            data_dict['segmentation'] = data_dict['segmentation'][remaining_channels]
+        return data_dict

batchgeneratorsv2/transforms/noise/blank_rectangle.py

@@ -0,0 +1,150 @@
+import numpy as np
+import torch
+from typing import Union, Tuple, List, Callable
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+class ColorFunctionExtractor:
+    def __init__(self, rectangle_value: Union[int, float, Tuple[float, float], Callable]):
+        self.rectangle_value = rectangle_value
+
+    def __call__(self, x: torch.Tensor) -> float:
+        if np.isscalar(self.rectangle_value):
+            return float(self.rectangle_value)
+        elif callable(self.rectangle_value):
+            return float(self.rectangle_value(x))
+        elif isinstance(self.rectangle_value, (tuple, list)):
+            return float(np.random.uniform(*self.rectangle_value))
+        else:
+            raise RuntimeError("Unrecognized format for rectangle_value")
+
+
+class BlankRectangleTransform(ImageOnlyTransform):
+    """
+    Overwrites random rectangles in the image with a constant or sampled value.
+
+    Supports 2D/3D data and various configurations of rectangle size/value.
+    """
+
+    def __init__(self,
+                 rectangle_size: Union[int,
+                                       Tuple[int, ...],
+                                       Tuple[Tuple[int, int], ...]],
+                 rectangle_value: Union[int, float, Tuple[float, float], Callable],
+                 num_rectangles: Union[int, Tuple[int, int]],
+                 force_square: bool = False,
+                 p_per_channel: float = 1.0):
+        super().__init__()
+        self.rectangle_size = rectangle_size
+        self.num_rectangles = num_rectangles
+        self.force_square = force_square
+        self.p_per_channel = p_per_channel
+        self.color_fn = ColorFunctionExtractor(rectangle_value)
+
+    def get_parameters(self, image: torch.Tensor, **kwargs) -> dict:
+        C = image.shape[0]
+        spatial_shape = image.shape[1:]
+        D = len(spatial_shape)
+
+        apply_channel = [np.random.rand() < self.p_per_channel for _ in range(C)]
+
+        # Number of rectangles
+        if isinstance(self.num_rectangles, int):
+            n_rects = [self.num_rectangles for _ in range(C)]
+        else:
+            n_rects = [np.random.randint(self.num_rectangles[0], self.num_rectangles[1]) for _ in range(C)]
+
+        # Precompute all rectangles for all channels
+        rectangles = [[] for _ in range(C)]
+        for c in range(C):
+            if not apply_channel[c]:
+                continue
+            for _ in range(n_rects[c]):
+                size = self._sample_rectangle_size(D)
+                lb = [np.random.randint(0, spatial_shape[d] - size[d] + 1) for d in range(D)]
+                ub = [lb[d] + size[d] for d in range(D)]
+                rectangles[c].append((lb, ub))
+
+        return {
+            'apply_channel': apply_channel,
+            'rectangles': rectangles
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        out = img
+        for c, (apply, rects) in enumerate(zip(params['apply_channel'], params['rectangles'])):
+            if not apply:
+                continue
+            for lb, ub in rects:
+                slices = tuple([slice(l, u) for l, u in zip(lb, ub)])
+                intensity = self.color_fn(out[c][slices])
+                out[c][slices] = intensity
+        return out
+
+    def _sample_rectangle_size(self, ndim: int) -> List[int]:
+        if isinstance(self.rectangle_size, int):
+            return [self.rectangle_size] * ndim
+
+        elif isinstance(self.rectangle_size, (tuple, list)) and all(isinstance(x, int) for x in self.rectangle_size):
+            return list(self.rectangle_size)
+
+        elif isinstance(self.rectangle_size, (tuple, list)) and all(isinstance(x, (tuple, list)) for x in self.rectangle_size):
+            if self.force_square:
+                val = np.random.randint(*self.rectangle_size[0])
+                return [val] * ndim
+            else:
+                return [np.random.randint(*self.rectangle_size[d]) for d in range(ndim)]
+
+        raise RuntimeError("Unrecognized format for rectangle_size")
+
+
+if __name__ == '__main__':
+    # import matplotlib.pyplot as plt
+    # from skimage.data import camera
+    # from skimage.util import img_as_float32
+    #
+    # img = torch.from_numpy(img_as_float32(camera())).unsqueeze(0)  # (C, H, W)
+    #
+    # transform = BlankRectangleTransform(
+    #     rectangle_size=((10, 30), (20, 40)),
+    #     rectangle_value=(0.0, 1.0),
+    #     num_rectangles=(2, 5),
+    #     force_square=False,
+    #     p_per_channel=1.0
+    # )
+    #
+    # params = transform.get_parameters(image=img)
+    # img_aug = transform._apply_to_image(img, **params)
+    #
+    # plt.subplot(1, 2, 1)
+    # plt.imshow(img.squeeze().numpy(), cmap='gray')
+    # plt.title("Original")
+    #
+    # plt.subplot(1, 2, 2)
+    # plt.imshow(img_aug.squeeze().numpy(), cmap='gray')
+    # plt.title("With Blank Rectangles")
+    #
+    # plt.tight_layout()
+    # plt.show()
+    import torch
+    import numpy as np
+    import matplotlib.pyplot as plt
+
+    # Create a random 3D image (C, D, H, W)
+    image = torch.rand(1, 32, 64, 64)  # Single-channel 3D volume
+
+    # Instantiate the transform
+    transform = BlankRectangleTransform(
+        rectangle_size=((4, 10), (10, 20), (10, 20)),  # (Z, Y, X) size ranges
+        rectangle_value=(0.0, 1.0),  # Random intensity per rectangle
+        num_rectangles=(3, 7),  # 3 to 6 rectangles per channel
+        force_square=False,
+        p_per_channel=1.0  # Always apply to the channel
+    )
+
+    # Sample transform parameters and apply
+    params = transform.get_parameters(image=image)
+    image_aug = transform._apply_to_image(image, **params)
+
+    from batchviewer import view_batch
+    view_batch(image, image_aug)

batchgeneratorsv2/transforms/noise/gaussian_blur.py

@@ -0,0 +1,260 @@
+from copy import deepcopy
+
+import numpy as np
+from time import time
+import torch
+from skimage.data import camera
+from torch.nn.functional import pad, conv3d, conv1d, conv2d
+
+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 blur_dimension(img: torch.Tensor, sigma: float, dim_to_blur: int, force_use_fft: bool = None, truncate: float = 6):
+    """
+    Smoothes an input image with a 1D Gaussian kernel along the specified dimension.
+    The function supports 1D, 2D, and 3D images.
+
+    :param img: Input image tensor with shape (C, X), (C, X, Y), or (C, X, Y, Z),
+                where C is the channel dimension and X, Y, Z are spatial dimensions.
+    :param sigma: The standard deviation of the Gaussian kernel.
+    :param dim_to_blur: The dimension along which to apply the Gaussian blur (0 for X, 1 for Y, 2 for Z).
+    :return: The blurred image tensor.
+    """
+    assert img.ndim - 1 > dim_to_blur, "dim_to_blur must be a valid spatial dimension of the input image."
+    # Adjustments for kernel based on image dimensions
+    spatial_dims = img.ndim - 1  # Number of spatial dimensions in the input image
+    kernel = _build_kernel(sigma, truncate=truncate)
+
+    ksize = kernel.shape[0]
+
+    # Dynamically set up padding, convolution operation, and kernel shape based on the number of spatial dimensions
+    conv_ops = {1: conv1d, 2: conv2d, 3: conv3d}
+    if force_use_fft is not None:
+        conv_op = conv_ops[spatial_dims] if not force_use_fft else fft_conv
+    else:
+        conv_op = conv_ops[spatial_dims]
+
+    # Adjust kernel and padding for the specified blur dimension and input dimensions
+    if spatial_dims == 1:
+        kernel = kernel[None, None, :]
+        padding = [ksize // 2, ksize // 2]
+    elif spatial_dims == 2:
+        if dim_to_blur == 0:
+            kernel = kernel[None, None, :, None]
+            padding = [0, 0, ksize // 2, ksize // 2]
+        else:  # dim_to_blur == 1
+            kernel = kernel[None, None, None, :]
+            padding = [ksize // 2, ksize // 2, 0, 0]
+    else:  # spatial_dims == 3
+        # Expand kernel and adjust padding based on the blur dimension
+        if dim_to_blur == 0:
+            kernel = kernel[None, None, :, None, None]
+            padding = [0, 0, 0, 0, ksize // 2, ksize // 2]
+        elif dim_to_blur == 1:
+            kernel = kernel[None, None, None, :, None]
+            padding = [0, 0, ksize // 2, ksize // 2, 0, 0]
+        else:  # dim_to_blur == 2
+            kernel = kernel[None, None, None, None, :]
+            padding = [ksize // 2, ksize // 2, 0, 0, 0, 0]
+
+    # Apply padding
+    img_padded = pad(img, padding, mode="reflect")
+
+    # Apply convolution
+    # remember that weights are [c_out, c_in, ...]
+    img_blurred = conv_op(img_padded[None], kernel.expand(img_padded.shape[0], *[-1] * (kernel.ndim - 1)), groups=img_padded.shape[0])[0]
+    return img_blurred
+
+
+class GaussianBlurTransform(ImageOnlyTransform):
+    def __init__(self,
+                 blur_sigma: RandomScalar = (1, 5),
+                 synchronize_channels: bool = False,  # todo make this p_synchronize_channels
+                 synchronize_axes: bool = False,  # todo make this p_synchronize_axes
+                 p_per_channel: float = 1,
+                 benchmark: bool = False
+                 ):
+        """
+        uses separable gaussian filters for all the speed
+
+        blur_sigma, if callable, will be called as blur_sigma(image, shape, dim) where shape is (c, x(, y, z) and dim i
+        s 1, 2 or 3 for x, y and z, respectively)
+        :param blur_sigma:
+        :param synchronize_channels:
+        :param synchronize_axes:
+        :param p_per_channel:
+        """
+        super().__init__()
+        self.blur_sigma = blur_sigma
+        self.benchmark = benchmark
+        self.synchronize_channels = synchronize_channels
+        self.synchronize_axes = synchronize_axes
+        self.p_per_channel = p_per_channel
+        self.benchmark_use_fft = {}  # shape -> kernel size -> use fft yes or no
+        self.benchmark_num_runs = 9
+
+    def get_parameters(self, **data_dict) -> dict:
+        shape = data_dict['image'].shape
+        dims = len(shape) - 1
+        dct = {}
+        dct['apply_to_channel'] = torch.rand(shape[0]) < self.p_per_channel
+        if self.synchronize_axes:
+            dct['sigmas'] = \
+                [[sample_scalar(self.blur_sigma, shape, dim=None)] * dims
+                 for _ in range(sum(dct['apply_to_channel']))] \
+                    if not self.synchronize_channels else \
+                    [sample_scalar(self.blur_sigma, shape, dim=None)] * dims
+        else:
+            dct['sigmas'] = \
+                [[sample_scalar(self.blur_sigma, shape, dim=i + 1) for i in range(dims)]
+                 for _ in range(sum(dct['apply_to_channel']))] \
+                    if not self.synchronize_channels else \
+                    [sample_scalar(self.blur_sigma, shape[i + 1]) for i in range(dims)]
+        return dct
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        if len(params['apply_to_channel']) == 0:
+            return img
+        dim = len(img.shape[1:])
+
+        # print(params['sigmas'])
+        if self.synchronize_channels:
+            # we can compute that in one go as the conv implementation supports arbitrary input channels (with expanded kernel)
+            for d in range(dim):
+                # print(d, params['sigmas'][d])
+                if not self.benchmark:
+                    img[params['apply_to_channel']] = blur_dimension(img[params['apply_to_channel']], params['sigmas'][d], d)
+                else:
+                    img[params['apply_to_channel']] = self._benchmark_wrapper(img[params['apply_to_channel']], params['sigmas'][d], d)
+        else:
+            # we have to go through all the channels, build the kernel for each channel etc
+            idx = np.where(params['apply_to_channel'])[0]
+            for j, i in enumerate(idx):
+                for d in range(dim):
+                    # print(i, d, params['sigmas'][i][d])
+                    if not self.benchmark:
+                        img[i:i+1] = blur_dimension(img[i:i+1], params['sigmas'][j][d], d)
+                    else:
+                        img[i:i+1] = self._benchmark_wrapper(img[i:i+1], params['sigmas'][j][d], d)
+        return img
+
+    def _benchmark_wrapper(self, img: torch.Tensor, sigma: float, dim_to_blur: int):
+        kernel_size = _compute_kernel_size(sigma)
+        shp = img.shape[dim_to_blur + 1]
+        # check if we already benchmarked this
+        if shp in self.benchmark_use_fft.keys() and kernel_size in self.benchmark_use_fft[shp].keys():
+            return blur_dimension(img, sigma, dim_to_blur, force_use_fft=self.benchmark_use_fft[shp][kernel_size])
+        else:
+            # let's not mess up the original image!
+            if shp not in self.benchmark_use_fft.keys():
+                self.benchmark_use_fft[shp] = {}
+            dummy_img = deepcopy(img)
+            times_nonfft = []
+            for _ in range(self.benchmark_num_runs):
+                st = time()
+                blur_dimension(dummy_img, sigma, dim_to_blur, force_use_fft=False)
+                times_nonfft.append(time() - st)
+            times_fft = []
+            for _ in range(self.benchmark_num_runs):
+                st = time()
+                blur_dimension(dummy_img, sigma, dim_to_blur, force_use_fft=True)
+                times_fft.append(time() - st)
+            # print(shp, kernel_size, np.median(times_fft), np.median(times_nonfft), np.median(times_fft) < np.median(times_nonfft))
+            self.benchmark_use_fft[shp][kernel_size] = np.median(times_fft) < np.median(times_nonfft)
+            # convenience stuff
+            self.benchmark_use_fft[shp] = dict(sorted(self.benchmark_use_fft[shp].items()))
+            # now create the real return value
+            return blur_dimension(img, sigma, dim_to_blur, force_use_fft=self.benchmark_use_fft[shp][kernel_size])
+
+
+def _build_kernel(sigma: float, truncate: float = 4) -> torch.Tensor:
+    kernel_size = _compute_kernel_size(sigma, truncate=truncate)
+    ksize_half = (kernel_size - 1) * 0.5
+
+    x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size)
+    pdf = torch.exp(-0.5 * (x / sigma).pow(2))
+    kernel1d = pdf / pdf.sum()
+
+    return kernel1d
+
+
+def _round_to_nearest_odd(n):
+    rounded = round(n)
+    # If the rounded number is odd, return it
+    if rounded % 2 == 1:
+        return rounded
+    # If the rounded number is even, adjust to the nearest odd number
+    return rounded + 1 if n - rounded >= 0 else rounded - 1
+
+
+def _compute_kernel_size(sigma, truncate: float = 4):
+    ksize = _round_to_nearest_odd(sigma * truncate + 0.5)
+    return ksize
+
+
+if __name__ == "__main__":
+    # this is the fastest for larger kernels but it doesn't fit well into the curent benchmark scheme
+
+    # tmp = np.fft.fftn(offsets[d].numpy())
+    # tmp = fourier_gaussian(tmp, sigmas)
+    # offsets[d] = torch.from_numpy(np.fft.ifftn(tmp).real)
+
+    import os
+    from batchgenerators.transforms.noise_transforms import GaussianBlurTransform as GBTBG
+    from batchviewer import view_batch
+
+    os.environ['OMP_NUM_THREADS'] = '1'
+    torch.set_num_threads(1)
+
+    data = camera()
+    data_dict = {'image': torch.from_numpy(camera()[None]).float()}
+
+    gnt2 = GaussianBlurTransform(2, False, False, 1, benchmark=False)
+    out = gnt2(**data_dict)
+    # view_batch(out['image'], torch.from_numpy(camera()[None]).float())
+
+
+    shape = (128, 164, 64)
+    num_warmup_for_benchmark = 1
+    num_repeats = 10
+    for sigma_range in (0.1, 1, 10, 20):
+        print(shape, sigma_range)
+        gnt2 = GaussianBlurTransform(sigma_range, False, False, 1, benchmark=False)
+        times = []
+        for _ in range(num_repeats):
+            data_dict = {'image': torch.ones((2, *shape))}
+            data_dict['image'][tuple([slice(data_dict['image'].shape[0])] + [slice(0, i // 2) for i in shape])] = 200
+            st = time()
+            out = gnt2(**data_dict)
+            times.append(time() - st)
+        print('w /o benchmark', np.median(times))
+
+        gnt = GaussianBlurTransform(sigma_range, False, False, 1, benchmark=True)
+        # warmup
+        for _ in range(num_warmup_for_benchmark):
+            data_dict = {'image': torch.ones((2, *shape))}
+            data_dict['image'][tuple([slice(data_dict['image'].shape[0])] + [slice(0, i // 2) for i in shape])] = 200
+            out = gnt(**data_dict)
+        times = []
+        for _ in range(num_repeats):
+            data_dict = {'image': torch.ones((2, *shape))}
+            data_dict['image'][tuple([slice(data_dict['image'].shape[0])] + [slice(0, i // 2) for i in shape])] = 200
+            st = time()
+            out = gnt(**data_dict)
+            times.append(time() - st)
+        print('with benchmark', np.median(times))
+
+        gnt3 = GBTBG(sigma_range, True, True, 0, 1, 1)
+        times = []
+        for _ in range(num_repeats):
+            data_dict = {'data': np.ones((1, 2, *shape))}
+            data_dict['data'][tuple([slice(data_dict['data'].shape[0])] + [slice(0, i // 2) for i in shape])] = 200
+            st = time()
+            out = gnt3(**data_dict)
+            times.append(time() - st)
+        print('batchgenerator', np.median(times))
+        print()
+        #
+        # print(gnt.benchmark_use_fft)

batchgeneratorsv2/transforms/noise/median_filter.py

@@ -0,0 +1,52 @@
+import numpy as np
+import torch
+from typing import Union, Tuple
+from scipy.ndimage import median_filter
+
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+class MedianFilterTransform(ImageOnlyTransform):
+    """
+    Applies a median filter to selected image channels.
+
+    Attributes:
+        filter_size (int or Tuple[int, int]): Either fixed filter size or range for random sampling.
+        p_same_for_each_channel (float): Probability that all channels share the same filter size.
+        p_per_channel (float): Probability of applying the filter to a given channel.
+    """
+
+    def __init__(self,
+                 filter_size: Union[int, Tuple[int, int]],
+                 p_same_for_each_channel: float = 0.0,
+                 p_per_channel: float = 1.0):
+        super().__init__()
+        self.filter_size = filter_size
+        self.p_same_for_each_channel = p_same_for_each_channel
+        self.p_per_channel = p_per_channel
+
+    def get_parameters(self, image: torch.Tensor, **kwargs) -> dict:
+        C = image.shape[0]
+        use_same = np.random.rand() < self.p_same_for_each_channel
+
+        if isinstance(self.filter_size, int):
+            sizes = [self.filter_size] * C
+        elif use_same:
+            sampled_size = int(np.random.randint(*self.filter_size))
+            sizes = [sampled_size] * C
+        else:
+            sizes = [int(np.random.randint(*self.filter_size)) for _ in range(C)]
+
+        apply_channel = [np.random.rand() < self.p_per_channel for _ in range(C)]
+
+        return {
+            'filter_sizes': sizes,
+            'apply_channel': apply_channel
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        img_np = img.cpu().numpy()
+        for c, (apply, size) in enumerate(zip(params['apply_channel'], params['filter_sizes'])):
+            if apply:
+                img_np[c] = median_filter(img_np[c], size=size)
+        return torch.from_numpy(img_np).to(img.device)

batchgeneratorsv2/transforms/noise/rician.py

@@ -0,0 +1,61 @@
+import torch
+from typing import Tuple
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+class RicianNoiseTransform(ImageOnlyTransform):
+    """
+    Adds Rician noise to simulate MRI characteristics.
+
+    Args:
+        noise_variance (Tuple[float, float]): Range to sample Gaussian noise variance used in Rician computation.
+    """
+
+    def __init__(self, noise_variance: Tuple[float, float] = (0.0, 0.1)):
+        super().__init__()
+        self.noise_variance = noise_variance
+
+    def get_parameters(self, image: torch.Tensor, **kwargs) -> dict:
+        variance = float(np.random.uniform(*self.noise_variance))
+        return {'variance': variance}
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        var = params['variance']
+        noise_real = torch.empty_like(img).normal_(mean=0.0, std=var)
+        noise_imag = torch.empty_like(img).normal_(mean=0.0, std=var)
+
+        min_val = img.min()
+        shifted = img - min_val
+
+        rician = torch.sqrt((shifted + noise_real).pow_(2).add_(noise_imag.pow_(2)))
+        rician = rician + min_val
+
+        # Normalize to match original mean and std
+        input_mean, input_std = img.mean(), img.std()
+        rician_mean, rician_std = rician.mean(), rician.std()
+
+        if rician_std > 0:
+            rician = (rician - rician_mean) / rician_std * input_std + input_mean
+        else:
+            rician = rician * 0 + input_mean  # fallback if std is zero (flat image)
+
+        return rician
+
+
+if __name__ == '__main__':
+    import torch
+    import numpy as np
+
+    # Create a synthetic normalized 3D image: (C, D, H, W)
+    image = torch.ones(1, 32, 64, 64) * 0.5  # z-score normalized MRI-like noise
+    image[0,1,1,1] = 2
+
+    # Instantiate the transform
+    transform = RicianNoiseTransform(noise_variance=(0.05, 0.1))
+
+    # Sample parameters and apply transform
+    params = transform.get_parameters(image=image)
+    image_noisy = transform._apply_to_image(image, **params)
+
+    from batchviewer import view_batch
+    view_batch(image, image_noisy)

batchgeneratorsv2/transforms/noise/sharpen.py

@@ -0,0 +1,128 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from typing import Union, Tuple, List
+
+from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
+
+
+class SharpeningTransform(ImageOnlyTransform):
+    """
+    Applies sharpening to 2D or 3D images using Laplacian-based contrast enhancement.
+    Preserves global intensity by explicitly adding scaled Laplacian to the original image.
+
+    Attributes:
+        strength (float or Tuple[float, float]): Sharpening strength.
+        p_same_for_each_channel (float): Probability of using the same strength for all channels.
+        p_per_channel (float): Probability of applying sharpening to a given channel.
+    """
+
+    def __init__(self,
+                 strength: Union[float, Tuple[float, float]] = 0.2,
+                 p_same_for_each_channel: float = 0.0,
+                 p_per_channel: float = 1.0,
+                 p_clamp_intensities: float = 0):
+        super().__init__()
+        self.strength = strength
+        self.p_same_for_each_channel = p_same_for_each_channel
+        self.p_per_channel = p_per_channel
+        self.p_clamp_intensities: float = p_clamp_intensities
+
+    def get_parameters(self, image: torch.Tensor, **kwargs) -> dict:
+        C = image.shape[0]
+        use_same = np.random.rand() < self.p_same_for_each_channel
+
+        if use_same:
+            strength = self._sample_strength()
+            strengths = [strength] * C
+            clamp = np.random.uniform() < self.p_clamp_intensities
+            clamp_intensities = [clamp] * C
+        else:
+            strengths = [self._sample_strength() for _ in range(C)]
+            clamp_intensities = [np.random.uniform() < self.p_clamp_intensities for _ in range(C)]
+
+        apply_channel = [np.random.rand() < self.p_per_channel for _ in range(C)]
+
+        return {
+            'strengths': strengths,
+            'clamp_intensities': clamp_intensities,
+            'apply_channel': apply_channel
+        }
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        out = img
+        spatial_dims = img.dim() - 1  # 2 for (C, H, W), 3 for (C, D, H, W)
+
+        if spatial_dims == 2:
+            kernel = torch.tensor([[0, -1, 0],
+                                   [-1, 4, -1],
+                                   [0, -1, 0]], dtype=torch.float32, device=img.device)
+            kernel = kernel.unsqueeze(0).unsqueeze(0)  # (1, 1, 3, 3)
+            pad = (1, 1, 1, 1)  # left, right, top, bottom
+
+        elif spatial_dims == 3:
+            kernel = torch.tensor([[[0, 0, 0],
+                                    [0, -1, 0],
+                                    [0, 0, 0]],
+                                   [[0, -1, 0],
+                                    [-1, 6, -1],
+                                    [0, -1, 0]],
+                                   [[0, 0, 0],
+                                    [0, -1, 0],
+                                    [0, 0, 0]]], dtype=torch.float32, device=img.device)
+            kernel = kernel.unsqueeze(0).unsqueeze(0)  # (1, 1, 3, 3, 3)
+            pad = (1, 1, 1, 1, 1, 1)  # left, right, top, bottom, front, back
+
+        else:
+            raise ValueError(f"Unsupported spatial dimensions: {spatial_dims}. Expected 2 or 3.")
+
+        for c, (apply, strength, clamp) in enumerate(zip(params['apply_channel'], params['strengths'], params['clamp_intensities'])):
+            if not apply:
+                continue
+
+            if clamp:
+                mn, mx = torch.min(img[c]), torch.max(img[c])
+
+            x = img[c].unsqueeze(0).unsqueeze(0)  # (1, 1, H, W) or (1, 1, D, H, W)
+            padded = F.pad(x, pad, mode='replicate')
+
+            if spatial_dims == 2:
+                laplace = F.conv2d(padded, kernel)
+            else:
+                laplace = F.conv3d(padded, kernel)
+
+            sharpened = x + strength * laplace
+            out[c] = sharpened.squeeze()
+
+            if clamp:
+                out[c].clamp_(mn, mx)
+
+        return out
+
+    def _sample_strength(self) -> float:
+        if isinstance(self.strength, float):
+            return self.strength
+        return float(np.random.uniform(*self.strength))
+
+
+if __name__ == '__main__':
+    from skimage.data import camera
+    from skimage.util import img_as_float32
+
+    # Load camera image and prepare it
+    img_np = img_as_float32(camera())  # (H, W), float32, values in [0, 1]
+    img_torch = torch.from_numpy(img_np).unsqueeze(0)  # (1, H, W) = (C, H, W)
+
+    # Instantiate the transform
+    transform = SharpeningTransform(
+        strength=(2, 2.1),  # Sharpening strength range
+        p_same_for_each_channel=1.0,  # Force same strength for all channels (only 1 channel here)
+        p_per_channel=1.0,  # Always apply
+        p_clamp_intensities = 1
+    )
+
+    # Generate parameters and apply
+    params = transform.get_parameters(image=img_torch)
+    sharpened = transform._apply_to_image(img_torch, **params)
+    from batchviewer import view_batch
+    view_batch(img_np, sharpened)