From 40a3c5d6480caf69440326a8063afb3ac82bbc76 Mon Sep 17 00:00:00 2001 From: iback Date: Thu, 20 Aug 2026 12:07:53 +0000 Subject: [PATCH] refactor: one kernels module instead of five copies Four independent 3D Gaussian implementations, three bias fields and two copies of the Laplace/Scharr constant tables lived in five modules: gpu/contrast.py, gpu/fromSeg.py, gpu/domain_transfer.py, synthseg/functional.py, cpu/contrast.py. Each carried some version of a comment saying it was kept local so its module would stay self-contained. They were not equivalent, and that is what let two of them go wrong unnoticed -- the uncentred Gaussian and the malformed 2D Scharr x-kernel, both corrected earlier in this series. Each had to be fixed on its own because there was no shared implementation to fix instead. smauglab/transforms/kernels.py is that implementation now. The consolidation is not bit-for-bit at the two blur call sites, deliberately: * Radius is ceil(3*sigma). domain_transfer and fromSeg used round(3*sigma), which is never wider, so their kernels may be one tap larger. * Padding is reflect everywhere. domain_transfer used replicate; fromSeg relied on conv3d's implicit zero padding, which pulls the volume border towards zero. That is the one difference here that was wrong rather than merely different. Everything else is the same arithmetic as the copy it replaces: the dense Gaussian, both derivative tables, and the bias field, whose synthseg and domain_transfer copies were already line-for-line identical. The local wrappers stay where callers expect them -- _gaussian_blur_3d in fromSeg.py still clamps to [0, 1], _random_bias_field3d in domain_transfer.py still returns a bare volume rather than [B, C, D, H, W] -- so this is a change of implementation, not of interface. unit_tests/test_kernels.py covers the merged behaviour and the equivalences that make it one implementation: that synthseg and domain_transfer now produce the same bias field from the same seed, and that the CPU transform hands out the shared tables. test_kernel_correctness.py follows the Gaussian to its new home. Co-Authored-By: Claude Opus 5 --- smauglab/transforms/cpu/contrast.py | 45 +--- smauglab/transforms/gpu/contrast.py | 82 +------ smauglab/transforms/gpu/domain_transfer.py | 44 ++-- smauglab/transforms/gpu/fromSeg.py | 21 +- smauglab/transforms/kernels.py | 259 +++++++++++++++++++++ smauglab/transforms/synthseg/functional.py | 43 +--- unit_tests/test_kernel_correctness.py | 30 +-- unit_tests/test_kernels.py | 94 ++++++++ 8 files changed, 409 insertions(+), 209 deletions(-) create mode 100644 smauglab/transforms/kernels.py create mode 100644 unit_tests/test_kernels.py diff --git a/smauglab/transforms/cpu/contrast.py b/smauglab/transforms/cpu/contrast.py index 5cc8119..2c46128 100644 --- a/smauglab/transforms/cpu/contrast.py +++ b/smauglab/transforms/cpu/contrast.py @@ -4,6 +4,8 @@ import torch.nn.functional as F from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform +from smauglab.transforms.kernels import laplace_kernel, scharr_kernels + class ConvTransform(ImageOnlyTransform): """ @@ -26,48 +28,11 @@ def get_parameters(self, **data_dict) -> dict: # _apply_to_image dispatches on kernel_type to tell the two apart. kernel: Union[torch.Tensor, list[torch.Tensor]] spatial_dims = len(data_dict["image"].shape) - 1 - if spatial_dims == 2: - if self.kernel_type == "Laplace": - kernel = torch.tensor([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype=torch.float32) - elif self.kernel_type == "Scharr": - # Middle row was [-10, 0, -10], summing the whole kernel to -20 rather - # than 0: not a gradient operator at all. The sibling kernel_y below - # has always been right, which is what makes this a typo. - kernel_x = torch.tensor([[-3, 0, 3], [-10, 0, 10], [-3, 0, 3]], dtype=torch.float32) - kernel_y = torch.tensor([[-3, -10, -3], [0, 0, 0], [3, 10, 3]], dtype=torch.float32) - kernel = [kernel_x, kernel_y] - elif spatial_dims == 3: + if spatial_dims in (2, 3): if self.kernel_type == "Laplace": - kernel = -1.0 * torch.ones(3, 3, 3, dtype=torch.float32) - kernel[1, 1, 1] = 26.0 + kernel = laplace_kernel(spatial_dims) elif self.kernel_type == "Scharr": - kernel_x = torch.tensor( - [ - [[9, 0, -9], [30, 0, -30], [9, 0, -9]], - [[30, 0, -30], [100, 0, -100], [30, 0, -30]], - [[9, 0, -9], [30, 0, -30], [9, 0, -9]], - ], - dtype=torch.float32, - ) - - kernel_y = torch.tensor( - [ - [[9, 30, 9], [0, 0, 0], [-9, -30, -9]], - [[30, 100, 30], [0, 0, 0], [-30, -100, -30]], - [[9, 30, 9], [0, 0, 0], [-9, -30, -9]], - ], - dtype=torch.float32, - ) - - kernel_z = torch.tensor( - [ - [[9, 30, 9], [30, 100, 30], [9, 30, 9]], - [[0, 0, 0], [0, 0, 0], [0, 0, 0]], - [[-9, -30, -9], [-30, -100, -30], [-9, -30, -9]], - ], - dtype=torch.float32, - ) - kernel = [kernel_x, kernel_y, kernel_z] + kernel = scharr_kernels(spatial_dims) else: raise ValueError(f"{self.__class__} can only handle 2D or 3D images.") diff --git a/smauglab/transforms/gpu/contrast.py b/smauglab/transforms/gpu/contrast.py index d3bae6d..0c33de4 100644 --- a/smauglab/transforms/gpu/contrast.py +++ b/smauglab/transforms/gpu/contrast.py @@ -8,6 +8,7 @@ from torch.nn import functional as F from smauglab.transforms.gpu.base import ImageOnlyTransform +from smauglab.transforms.kernels import gaussian_kernel3d, laplace_kernel, scharr_kernels from smauglab.transforms.rng import shared_choice @@ -172,48 +173,18 @@ def get_kernel(self, device: torch.device) -> Union[Tensor, list[Tensor]]: # kernel type returns a single tensor. kernel: Union[Tensor, list[Tensor]] if self.kernel_type == "Laplace": - kernel = -1.0 * torch.ones(3, 3, 3, dtype=torch.float32, device=device) - kernel[1, 1, 1] = 26.0 + kernel = laplace_kernel(3, device=device) elif self.kernel_type == "Scharr": - kernel_x = torch.tensor( - [ - [[9, 0, -9], [30, 0, -30], [9, 0, -9]], - [[30, 0, -30], [100, 0, -100], [30, 0, -30]], - [[9, 0, -9], [30, 0, -30], [9, 0, -9]], - ], - dtype=torch.float32, - device=device, - ) - - kernel_y = torch.tensor( - [ - [[9, 30, 9], [0, 0, 0], [-9, -30, -9]], - [[30, 100, 30], [0, 0, 0], [-30, -100, -30]], - [[9, 30, 9], [0, 0, 0], [-9, -30, -9]], - ], - dtype=torch.float32, - device=device, - ) - - kernel_z = torch.tensor( - [ - [[9, 30, 9], [30, 100, 30], [9, 30, 9]], - [[0, 0, 0], [0, 0, 0], [0, 0, 0]], - [[-9, -30, -9], [-30, -100, -30], [-9, -30, -9]], - ], - dtype=torch.float32, - device=device, - ) - kernel = [kernel_x, kernel_y, kernel_z] + kernel = scharr_kernels(3, device=device) elif self.kernel_type == "GaussianBlur": sigma = torch.rand(3, device=device) * self.sigma kernel_size = 3 - kernel = get_gaussian_kernel3d(kernel_size, sigma, torch.float32, device) + kernel = gaussian_kernel3d(kernel_size, sigma, torch.float32, device) elif self.kernel_type == "UnsharpMask": # For unsharp masking we use a Gaussian blur kernel; amount is applied in apply_transform. sigma = torch.rand(3, device=device) * self.sigma kernel_size = 3 - kernel = get_gaussian_kernel3d(kernel_size, sigma, torch.float32, device) + kernel = gaussian_kernel3d(kernel_size, sigma, torch.float32, device) elif self.kernel_type == "RandConv": # choose random odd kernel size e.g. [1,3,5,7] k = int(shared_choice(self.kernel_sizes)) # define kernel_sizes in __init__ @@ -345,49 +316,6 @@ def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torc return img -def get_gaussian_kernel1d(kernel_size: int, sigma: Union[float, Tensor], dtype: torch.dtype, device: torch.device) -> Tensor: - """Create a 1D Gaussian kernel, centred on the middle tap. - - The sample points were `arange(kernel_size)` -- 0, 1, 2 -- which puts the peak at - index 0 instead of the centre. The resulting 3D kernel had its maximum at corner - [0,0,0], so RandomGaussianBlurGPU and RandomUnsharpMaskGPU blurred *and* translated - the image by about a voxel, relative to a segmentation mask that is not convolved. - """ - half = (kernel_size - 1) / 2.0 - x = torch.linspace(-half, half, kernel_size, dtype=dtype, device=device) - pdf = torch.exp(-0.5 * (x / sigma).pow(2)) - kernel1d = pdf / pdf.sum() - - return kernel1d - - -def get_gaussian_kernel3d(kernel_size: int, sigma: Union[float, Tensor], dtype: torch.dtype, device: torch.device) -> Tensor: - """ - Create a 3D Gaussian kernel by multiplying 1D kernels along each axis. - Args: - kernel_size (int) - sigma (float or tuple of three floats): Standard deviation of the Gaussian kernel. - """ - if isinstance(sigma, (int, float)): - sigma = torch.tensor([sigma, sigma, sigma], device=device) - elif isinstance(sigma, torch.Tensor): - assert sigma.shape == (3,), "Sigma must be a float or a tensor of three floats." - else: - raise TypeError("Sigma must be a float or a tensor of three floats.") - - gz = get_gaussian_kernel1d(kernel_size, sigma[0], dtype, device) - gy = get_gaussian_kernel1d(kernel_size, sigma[1], dtype, device) - gx = get_gaussian_kernel1d(kernel_size, sigma[2], dtype, device) - - # Outer product using broadcasting - kernel = gz[:, None, None] * gy[None, :, None] * gx[None, None, :] - - # Normalize - kernel /= kernel.sum() - - return kernel - - ## Noise transform class RandomGaussianNoiseGPU(ImageOnlyTransform): """Add random Gaussian noise to image. diff --git a/smauglab/transforms/gpu/domain_transfer.py b/smauglab/transforms/gpu/domain_transfer.py index c1927ff..32414e8 100644 --- a/smauglab/transforms/gpu/domain_transfer.py +++ b/smauglab/transforms/gpu/domain_transfer.py @@ -44,50 +44,34 @@ from torch.nn import functional as F from smauglab.transforms.gpu.base import ImageOnlyTransform +from smauglab.transforms.kernels import gaussian_blur3d, random_bias_field3d # Default transfer LUT bank (built by embeddaug/analysis/playground/build_transfer_bank.py). DEFAULT_BANK_PATH = "/DATA/NAS/ongoing_projects/hendrik/nathan-transferaug/embeddaug/analysis/playground/results/domain_transfer_bank.npz" -def _gaussian_kernel1d(sigma: float, device, dtype) -> torch.Tensor: - radius = max(1, round(3.0 * sigma)) - x = torch.arange(-radius, radius + 1, device=device, dtype=dtype) - k = torch.exp(-0.5 * (x / sigma) ** 2) - return k / k.sum() - - def _gaussian_blur3d(x: torch.Tensor, sigma: float) -> torch.Tensor: - """Separable Gaussian blur over the 3 spatial dims of [N, C, D, H, W].""" + """Separable Gaussian blur over the 3 spatial dims of [N, C, D, H, W]. + + Delegates to the shared implementation. That one pads with `reflect` rather than + the `replicate` used here, and takes its radius from `ceil(3*sigma)` rather than + `round`, so the kernel can be one tap wider -- see smauglab/transforms/kernels.py. + """ if sigma <= 0: return x - n, c = x.shape[:2] - k = _gaussian_kernel1d(sigma, x.device, x.dtype) - r = (k.numel() - 1) // 2 - for dim in (2, 3, 4): - shape = [1, 1, 1, 1, 1] - shape[dim] = k.numel() - ker = k.view(shape).repeat(c, 1, 1, 1, 1) # [C,1,kD,kH,kW] for separable conv - pad = [0, 0, 0, 0, 0, 0] - pad[(4 - dim) * 2] = r - pad[(4 - dim) * 2 + 1] = r - x = F.conv3d(F.pad(x, pad, mode="replicate"), ker, groups=c) - return x + return gaussian_blur3d(x, float(sigma)) def _random_bias_field3d(shape, std: float, scale: float, device, dtype) -> torch.Tensor: """Smooth positive multiplicative bias field over a ``[D, H, W]`` volume. - Samples a coarse Gaussian grid ``~ N(0, U(0, std))`` of size ``ceil(shape*scale)``, - trilinear-upsamples it to ``shape`` and exponentiates (Gaussian in log-space → positive, - multiplicative). Same pattern as ``synthseg/functional.py::bias_field`` and - ``contrast.py::RandomBiasFieldGPU``; kept local so this module stays self-contained. + Thin adapter over :func:`smauglab.transforms.kernels.random_bias_field3d`, which + returns ``[batch, channels, D, H, W]``; this module wants the bare volume. The + implementation used to be written out here as well -- line for line the same as + ``synthseg/functional.py::bias_field`` -- under a comment saying it was "kept local + so this module stays self-contained". """ - d, h, w = shape - small = [max(2, math.ceil(s * scale)) for s in (d, h, w)] - s = torch.rand((), device=device) * std - field = torch.randn(1, 1, *small, device=device, dtype=dtype) * s - field = F.interpolate(field, size=(d, h, w), mode="trilinear", align_corners=True) - return torch.exp(field)[0, 0] + return random_bias_field3d(tuple(shape), std, scale, device, dtype)[0, 0] def _random_smooth_field01(shape, scale: float, gain: float, device, dtype) -> torch.Tensor: diff --git a/smauglab/transforms/gpu/fromSeg.py b/smauglab/transforms/gpu/fromSeg.py index 4ed6a75..e4f5917 100644 --- a/smauglab/transforms/gpu/fromSeg.py +++ b/smauglab/transforms/gpu/fromSeg.py @@ -5,6 +5,7 @@ from torch.nn import functional as F from smauglab.transforms.gpu.base import ImageOnlyTransform +from smauglab.transforms.kernels import gaussian_blur3d from smauglab.transforms.rng import shared_choice # ── PALETTE AUG helpers ────────────────────────────────────────────────── @@ -26,16 +27,16 @@ def _kmeans_1d(values: torch.Tensor, C: int, n_iter: int = 10) -> torch.Tensor: def _gaussian_blur_3d(x: torch.Tensor, sigma: float) -> torch.Tensor: - """Separable 3D Gaussian blur. x: (B, 1, D, H, W).""" - k_r = max(1, int(3.0 * sigma + 0.5)) - k1d = torch.arange(-k_r, k_r + 1, dtype=x.dtype, device=x.device) - k1d = torch.exp(-0.5 * (k1d / sigma) ** 2) - k1d = k1d / k1d.sum() - pad = len(k1d) // 2 - y = F.conv3d(x, k1d.view(1, 1, -1, 1, 1), padding=(pad, 0, 0)) - y = F.conv3d(y, k1d.view(1, 1, 1, -1, 1), padding=(0, pad, 0)) - y = F.conv3d(y, k1d.view(1, 1, 1, 1, -1), padding=(0, 0, pad)) - return y.clamp(0, 1) + """Separable 3D Gaussian blur of a (B, 1, D, H, W) volume, clamped to [0, 1]. + + Delegates to the shared implementation, which pads with `reflect`. This copy used + conv3d's implicit zero padding, which pulled the volume border towards 0 -- the + clamp below hid the top end of that but not the darkening. It also took its radius + from `round(3*sigma)` rather than `ceil`, so kernels can be one tap wider now. + """ + if sigma <= 0: + return x + return gaussian_blur3d(x, float(sigma)).clamp(0, 1) def _voronoi_region_ids( diff --git a/smauglab/transforms/kernels.py b/smauglab/transforms/kernels.py new file mode 100644 index 0000000..3136560 --- /dev/null +++ b/smauglab/transforms/kernels.py @@ -0,0 +1,259 @@ +"""Convolution kernels and smooth random fields, shared by every backend. + +Four independent 3-D Gaussian implementations, three bias fields and two copies of +the Laplace/Scharr constant tables used to live in five different modules. They were +not equivalent, and the divergence is what let two of them go wrong unnoticed: the +uncentred Gaussian in `gpu/contrast.py` and the malformed 2-D Scharr x-kernel in +`cpu/contrast.py`, both corrected earlier in this series. Each was fixed on its own +because there was no shared implementation to fix instead. This module is that +implementation; anything that convolves or blurs should import from here rather than +growing a fifth copy. + +The consolidation is not bit-for-bit for the two blur call sites, and deliberately so: + +* Radius is `ceil(3*sigma)`. `domain_transfer` and `fromSeg` used `round(3*sigma)`, + which is never wider, so their kernels may now be one tap larger. +* Padding is `reflect` everywhere. `domain_transfer` used `replicate` and `fromSeg` + relied on conv3d's implicit zero padding. Zero padding darkens the volume border, + which is the one difference here that was wrong rather than merely different. + +Everything else -- the dense Gaussian, both derivative tables, the bias field -- is +the same arithmetic as the copy it replaces. +""" + +from __future__ import annotations + +import math +from typing import Union + +import torch +import torch.nn.functional as F +from torch import Tensor + +__all__ = [ + "LAPLACE_2D", + "LAPLACE_3D", + "SCHARR_2D", + "SCHARR_3D", + "gaussian_blur3d", + "gaussian_kernel1d", + "gaussian_kernel3d", + "laplace_kernel", + "random_bias_field3d", + "scharr_kernels", +] + + +# --- Gaussian kernels ------------------------------------------------------------- + + +def gaussian_kernel1d(sigma: float, device: torch.device, dtype: torch.dtype = torch.float32) -> Tensor: + """A normalised 1-D Gaussian, centred, with radius `ceil(3*sigma)`. + + A non-positive sigma means "do not blur", and returns the identity kernel `[1.0]` + rather than raising -- `blurring_sigma_for_downsampling` legitimately produces + zeros for axes that are already at the target resolution. + """ + if sigma <= 0: + return torch.tensor([1.0], device=device, dtype=dtype) + radius = max(1, math.ceil(3.0 * sigma)) + x = torch.arange(-radius, radius + 1, device=device, dtype=dtype) + kernel = torch.exp(-0.5 * (x / sigma) ** 2) + return kernel / kernel.sum() + + +def gaussian_kernel3d( + kernel_size: int, + sigma: Union[float, Tensor], + dtype: torch.dtype, + device: torch.device, +) -> Tensor: + """A dense `[kernel_size]*3` Gaussian as the outer product of three 1-D kernels. + + Fixed-size rather than sigma-derived, because the caller + (`_RandomConvBaseGPU.get_kernel`) hands the kernel to a generic convolution path + shared with Scharr and RandConv and needs a tensor of a known shape. + + The sample points are centred on the kernel: `linspace(-(k-1)/2, (k-1)/2, k)`. + Sampling at `arange(k)` -- what this used to do -- puts the peak at index 0 and + turns the blur into a blur plus a translation. + """ + if isinstance(sigma, (int, float)): + sigma_t = torch.tensor([float(sigma)] * 3, device=device, dtype=dtype) + elif isinstance(sigma, Tensor): + if sigma.shape != (3,): + raise ValueError(f"sigma must be a float or a tensor of three floats, got shape {tuple(sigma.shape)}") + sigma_t = sigma.to(device=device, dtype=dtype) + else: + raise TypeError(f"sigma must be a float or a tensor of three floats, got {type(sigma).__name__}") + + half = (kernel_size - 1) / 2.0 + x = torch.linspace(-half, half, kernel_size, device=device, dtype=dtype) + + axes = [] + for axis in range(3): + s = sigma_t[axis] + # A zero sigma degenerates to a delta at the centre; exp(-inf) would be 0 + # everywhere and the normalisation would divide by zero. + if float(s) <= 0: + delta = torch.zeros(kernel_size, device=device, dtype=dtype) + delta[kernel_size // 2] = 1.0 + axes.append(delta) + continue + pdf = torch.exp(-0.5 * (x / s).pow(2)) + axes.append(pdf / pdf.sum()) + + kernel = axes[0][:, None, None] * axes[1][None, :, None] * axes[2][None, None, :] + return kernel / kernel.sum() + + +def gaussian_blur3d( + image: Tensor, + sigma: Union[float, Tensor], + *, + blur_range: float = 1.0, + padding_mode: str = "reflect", +) -> Tensor: + """Separable, optionally anisotropic Gaussian blur of a `[B, C, D, H, W]` volume. + + `sigma` is either a scalar or a `(3,)` tensor of per-axis sigmas. `blur_range > 1` + multiplies every sigma by `U(1/blur_range, blur_range)`, which is SynthSeg's + `DynamicGaussianBlur` jitter; the default of 1.0 disables it. + + Three 1-D convolutions rather than one dense 3-D kernel: for a radius-r kernel that + is 3*(2r+1) multiply-adds per voxel instead of (2r+1)^3. + """ + if image.dim() != 5: + raise ValueError(f"expected a 5D [B, C, D, H, W] tensor, got shape {tuple(image.shape)}") + + channels = image.shape[1] + device, dtype = image.device, image.dtype + + if isinstance(sigma, Tensor): + sigmas = sigma.detach().to(device=device, dtype=torch.float32).flatten() + if sigmas.numel() == 1: + sigmas = sigmas.repeat(3) + else: + sigmas = torch.full((3,), float(sigma), device=device, dtype=torch.float32) + + if blur_range and blur_range > 1.0: + jitter = (1.0 / blur_range) + torch.rand(3, device=device) * (blur_range - 1.0 / blur_range) + sigmas = sigmas * jitter + + out = image + for axis, s in enumerate(sigmas.tolist()): + if s <= 0: + continue + kernel = gaussian_kernel1d(s, device, dtype) + ksize = kernel.numel() + if ksize == 1: + continue + pad = ksize // 2 + + shape = [1, 1, 1, 1, 1] + shape[2 + axis] = ksize + weight = kernel.view(shape).repeat(channels, 1, 1, 1, 1) + + # F.pad's tuple runs last spatial axis first: (W_lo, W_hi, H_lo, H_hi, D_lo, D_hi). + pad_full = [0, 0, 0, 0, 0, 0] + pad_full[(2 - axis) * 2] = pad + pad_full[(2 - axis) * 2 + 1] = pad + + out = F.conv3d(F.pad(out, pad_full, mode=padding_mode), weight, groups=channels) + return out + + +# --- smooth random fields --------------------------------------------------------- + + +def random_bias_field3d( + shape: tuple[int, int, int], + std: float, + scale: float, + device: torch.device, + dtype: torch.dtype = torch.float32, + *, + batch: int = 1, + channels: int = 1, +) -> Tensor: + """A smooth, positive, multiplicative bias field of shape `[batch, channels, *shape]`. + + Sample `N(0, U(0, std))` on a coarse `ceil(shape * scale)` grid, trilinear-upsample + to full resolution and exponentiate. Gaussian in log space, so the field is + strictly positive and multiplies rather than shifts. + + This is lab2im's `BiasFieldCorruption`, and was written out three times: in + `synthseg/functional.py::bias_field`, in `gpu/domain_transfer.py::_random_bias_field3d` + and (in a different, polynomial form) in `gpu/contrast.py::RandomBiasFieldGPU`. The + first two were line-for-line identical. + """ + if std <= 0: + return torch.ones(batch, channels, *shape, device=device, dtype=dtype) + + small = [max(2, math.ceil(s * scale)) for s in shape] + # One std per batch element, shared across channels, matching lab2im. + sampled_std = torch.rand(batch, 1, 1, 1, 1, device=device, dtype=dtype) * std + field = torch.randn(batch, channels, *small, device=device, dtype=dtype) * sampled_std + field = F.interpolate(field, size=tuple(shape), mode="trilinear", align_corners=True) + return torch.exp(field) + + +# --- fixed derivative kernels ----------------------------------------------------- +# +# Held as nested lists rather than tensors so there is no import-time device or dtype +# choice; `laplace_kernel` and `scharr_kernels` materialise them on demand. + +#: 8-neighbour 2-D Laplacian. +LAPLACE_2D = [ + [-1, -1, -1], + [-1, 8, -1], + [-1, -1, -1], +] + +#: 26-neighbour 3-D Laplacian: -1 everywhere, +26 in the centre. Sums to 0. +LAPLACE_3D = [[[-1] * 3 for _ in range(3)] for _ in range(3)] +LAPLACE_3D[1][1][1] = 26 + +#: 2-D Scharr, (x, y). The x-kernel's middle row is [-10, 0, 10]; the CPU copy of this +#: table had [-10, 0, -10], which made it sum to -20 instead of 0. +SCHARR_2D = [ + [[-3, 0, 3], [-10, 0, 10], [-3, 0, 3]], + [[-3, -10, -3], [0, 0, 0], [3, 10, 3]], +] + +#: 3-D Scharr, (x, y, z). +SCHARR_3D = [ + [ + [[9, 0, -9], [30, 0, -30], [9, 0, -9]], + [[30, 0, -30], [100, 0, -100], [30, 0, -30]], + [[9, 0, -9], [30, 0, -30], [9, 0, -9]], + ], + [ + [[9, 30, 9], [0, 0, 0], [-9, -30, -9]], + [[30, 100, 30], [0, 0, 0], [-30, -100, -30]], + [[9, 30, 9], [0, 0, 0], [-9, -30, -9]], + ], + [ + [[9, 30, 9], [30, 100, 30], [9, 30, 9]], + [[0, 0, 0], [0, 0, 0], [0, 0, 0]], + [[-9, -30, -9], [-30, -100, -30], [-9, -30, -9]], + ], +] + + +def laplace_kernel(spatial_dims: int, device: torch.device | None = None, dtype: torch.dtype = torch.float32) -> Tensor: + """The Laplacian for 2-D or 3-D data.""" + if spatial_dims == 2: + return torch.tensor(LAPLACE_2D, dtype=dtype, device=device) + if spatial_dims == 3: + return torch.tensor(LAPLACE_3D, dtype=dtype, device=device) + raise ValueError(f"Laplace kernel is defined for 2 or 3 spatial dimensions, got {spatial_dims}") + + +def scharr_kernels(spatial_dims: int, device: torch.device | None = None, dtype: torch.dtype = torch.float32) -> list[Tensor]: + """The directional Scharr kernels: two for 2-D data, three for 3-D.""" + if spatial_dims == 2: + return [torch.tensor(k, dtype=dtype, device=device) for k in SCHARR_2D] + if spatial_dims == 3: + return [torch.tensor(k, dtype=dtype, device=device) for k in SCHARR_3D] + raise ValueError(f"Scharr kernels are defined for 2 or 3 spatial dimensions, got {spatial_dims}") diff --git a/smauglab/transforms/synthseg/functional.py b/smauglab/transforms/synthseg/functional.py index 30d6979..6a8ff49 100644 --- a/smauglab/transforms/synthseg/functional.py +++ b/smauglab/transforms/synthseg/functional.py @@ -33,6 +33,8 @@ import torch import torch.nn.functional as F +from smauglab.transforms.kernels import gaussian_blur3d, random_bias_field3d + Number = Union[int, float] __all__ = [ @@ -478,12 +480,7 @@ def bias_field( if bias_field_std <= 0: return image B, C, D, H, W = image.shape - device = image.device - small = [max(2, math.ceil(s * bias_scale)) for s in (D, H, W)] - std = torch.rand(B, 1, 1, 1, 1, device=device) * bias_field_std - field = torch.randn(B, C, *small, device=device) * std - field = F.interpolate(field, size=(D, H, W), mode="trilinear", align_corners=True) - return image * torch.exp(field) + return image * random_bias_field3d((D, H, W), bias_field_std, bias_scale, image.device, image.dtype, batch=B, channels=C) # --------------------------------------------------------------------------- @@ -545,15 +542,6 @@ def blurring_sigma_for_downsampling( return sigma -def _gaussian_kernel1d(sigma: float, device: torch.device) -> torch.Tensor: - if sigma <= 0: - return torch.tensor([1.0], device=device) - radius = max(1, math.ceil(3.0 * sigma)) - x = torch.arange(-radius, radius + 1, device=device, dtype=torch.float32) - k = torch.exp(-0.5 * (x / sigma) ** 2) - return k / k.sum() - - def gaussian_blur_3d( image: torch.Tensor, sigma: torch.Tensor, @@ -566,30 +554,7 @@ def gaussian_blur_3d( SynthSeg, ``1.15`` in the 2020 lab2im model) and applied as three 1D convolutions (reflect padding). ``sigma`` is a ``(3,)`` tensor. """ - B, C, D, H, W = image.shape - device = image.device - sigma = sigma.clone().float() - if blur_range and blur_range > 1.0: - jitter = (1.0 / blur_range) + torch.rand(3, device=device) * (blur_range - 1.0 / blur_range) - sigma = sigma * jitter - - out = image - for axis, s in enumerate(sigma.tolist()): - if s <= 0: - continue - kernel = _gaussian_kernel1d(s, device) - ksize = kernel.numel() - pad = ksize // 2 - # shape the separable kernel for conv3d along the given spatial axis - shape = [1, 1, 1, 1, 1] - shape[2 + axis] = ksize - weight = kernel.view(shape).repeat(C, 1, 1, 1, 1) - padding = [0, 0, 0] - padding[axis] = pad - pad_full = (padding[2], padding[2], padding[1], padding[1], padding[0], padding[0]) - out = F.pad(out, pad_full, mode="reflect") - out = F.conv3d(out, weight, groups=C) - return out + return gaussian_blur3d(image, sigma, blur_range=blur_range) def sample_resolution( diff --git a/unit_tests/test_kernel_correctness.py b/unit_tests/test_kernel_correctness.py index ea58adc..dad5817 100644 --- a/unit_tests/test_kernel_correctness.py +++ b/unit_tests/test_kernel_correctness.py @@ -12,7 +12,7 @@ import torch.nn.functional as F from smauglab.transforms.cpu.contrast import ConvTransform -from smauglab.transforms.gpu.contrast import get_gaussian_kernel1d, get_gaussian_kernel3d +from smauglab.transforms.kernels import gaussian_kernel1d, gaussian_kernel3d from unit_tests.helpers import SmaugLabTestCase CPU = torch.device("cpu") @@ -20,41 +20,45 @@ class TestGaussianKernelIsCentred(SmaugLabTestCase): def test_the_1d_kernel_peaks_in_the_middle(self): - for kernel_size in (3, 5, 7): - with self.subTest(kernel_size=kernel_size): - kernel = get_gaussian_kernel1d(kernel_size, 1.0, torch.float32, CPU) - self.assertEqual(int(kernel.argmax()), kernel_size // 2, "the 1D Gaussian's peak is not the centre tap") + for sigma in (0.5, 1.0, 2.5): + with self.subTest(sigma=sigma): + kernel = gaussian_kernel1d(sigma, CPU) + self.assertEqual(kernel.numel() % 2, 1, "an even-length kernel has no centre tap") + self.assertEqual(int(kernel.argmax()), kernel.numel() // 2, "the 1D Gaussian's peak is not the centre tap") def test_the_1d_kernel_is_symmetric_and_normalised(self): - for kernel_size in (3, 5, 7): - with self.subTest(kernel_size=kernel_size): - kernel = get_gaussian_kernel1d(kernel_size, 1.3, torch.float32, CPU) + for sigma in (0.5, 1.3, 2.5): + with self.subTest(sigma=sigma): + kernel = gaussian_kernel1d(sigma, CPU) self.assertTrue(torch.allclose(kernel, kernel.flip(0), atol=1e-6)) self.assertAlmostEqual(float(kernel.sum()), 1.0, places=5) + def test_a_non_positive_sigma_gives_the_identity_kernel(self): + self.assertTrue(torch.equal(gaussian_kernel1d(0.0, CPU), torch.tensor([1.0]))) + def test_the_3d_kernel_peaks_at_the_centre_voxel(self): for kernel_size in (3, 5): with self.subTest(kernel_size=kernel_size): - kernel = get_gaussian_kernel3d(kernel_size, 1.0, torch.float32, CPU) + kernel = gaussian_kernel3d(kernel_size, 1.0, torch.float32, CPU) centre = kernel_size // 2 expected = (centre * kernel_size + centre) * kernel_size + centre self.assertEqual(int(kernel.argmax()), expected, "the 3D Gaussian's maximum is not the centre voxel") def test_the_3d_kernel_is_symmetric_on_every_axis(self): - kernel = get_gaussian_kernel3d(5, 1.3, torch.float32, CPU) + kernel = gaussian_kernel3d(5, 1.3, torch.float32, CPU) for axis in (0, 1, 2): with self.subTest(axis=axis): self.assertTrue(torch.allclose(kernel, kernel.flip(axis), atol=1e-6)) def test_the_3d_kernel_sums_to_one(self): - kernel = get_gaussian_kernel3d(3, torch.tensor([0.5, 1.0, 2.0]), torch.float32, CPU) + kernel = gaussian_kernel3d(3, torch.tensor([0.5, 1.0, 2.0]), torch.float32, CPU) self.assertAlmostEqual(float(kernel.sum()), 1.0, places=5) def test_blurring_an_impulse_leaves_its_centre_of_mass_in_place(self): """The translation is the part that actually hurt: the mask does not move with it.""" volume = torch.zeros(1, 1, 15, 15, 15) volume[0, 0, 7, 7, 7] = 1.0 - kernel = get_gaussian_kernel3d(7, 1.5, torch.float32, CPU) + kernel = gaussian_kernel3d(7, 1.5, torch.float32, CPU) blurred = F.conv3d(volume, kernel.view(1, 1, 7, 7, 7), padding=3) @@ -73,7 +77,7 @@ def test_the_old_uncentred_formula_really_was_off_centre(self): old = pdf / pdf.sum() self.assertEqual(int(old.argmax()), 0, "control: the old kernel peaked at index 0") - self.assertEqual(int(get_gaussian_kernel1d(kernel_size, sigma, torch.float32, CPU).argmax()), 1) + self.assertEqual(int(gaussian_kernel3d(kernel_size, sigma, torch.float32, CPU)[:, 1, 1].argmax()), 1) class TestScharrIsAGradientOperator(SmaugLabTestCase): diff --git a/unit_tests/test_kernels.py b/unit_tests/test_kernels.py new file mode 100644 index 0000000..1f6af8f --- /dev/null +++ b/unit_tests/test_kernels.py @@ -0,0 +1,94 @@ +"""The shared kernel module, exercised as the consolidation it is. + +Four Gaussian blurs, three bias fields and two copies of the Laplace/Scharr tables +lived in five modules. `test_kernel_correctness.py` covers the two that were wrong; +this file covers the properties the merged implementation has to preserve for the +call sites it replaced, and the equivalences that make it one implementation rather +than a sixth copy. +""" + +import torch + +from smauglab.transforms.kernels import gaussian_blur3d, gaussian_kernel1d, random_bias_field3d +from unit_tests.helpers import SmaugLabTestCase + + +class TestGaussianBlur(SmaugLabTestCase): + def test_a_constant_volume_survives_the_blur(self): + """Normalised kernel plus reflect padding means no darkening at the border. + + The copy in gpu/fromSeg.py zero-padded, which pulled the border towards 0. + """ + volume = torch.full((1, 1, 12, 12, 12), 3.0) + blurred = gaussian_blur3d(volume, 1.0) + self.assertTrue(torch.allclose(blurred, volume, atol=1e-5)) + + def test_zero_sigma_is_a_no_op(self): + volume = self.tiny_volume() + self.assertTrue(torch.equal(gaussian_blur3d(volume, 0.0), volume)) + + def test_anisotropic_sigma_blurs_only_the_named_axis(self): + volume = torch.zeros(1, 1, 15, 15, 15) + volume[0, 0, 7, 7, 7] = 1.0 + + blurred = gaussian_blur3d(volume, torch.tensor([2.0, 0.0, 0.0])) + + # Axis 0 spread out; the other two still hold a single non-zero plane. + self.assertGreater(int((blurred[0, 0, :, 7, 7] > 1e-6).sum()), 1) + self.assertEqual(int((blurred[0, 0, 7, :, 7] > 1e-6).sum()), 1) + self.assertEqual(int((blurred[0, 0, 7, 7, :] > 1e-6).sum()), 1) + + def test_multichannel_volumes_are_blurred_per_channel(self): + volume = torch.zeros(2, 3, 11, 11, 11) + volume[:, :, 5, 5, 5] = 1.0 + blurred = gaussian_blur3d(volume, 1.0) + self.assertEqual(blurred.shape, volume.shape) + for c in range(3): + self.assertAlmostEqual(float(blurred[0, c].sum()), 1.0, places=4) + + def test_a_non_5d_input_is_rejected(self): + with self.assertRaises(ValueError): + gaussian_blur3d(torch.rand(4, 4, 4), 1.0) + + def test_kernel1d_is_normalised_and_odd_length(self): + for sigma in (0.5, 1.0, 2.5): + kernel = gaussian_kernel1d(sigma, torch.device("cpu")) + self.assertEqual(kernel.numel() % 2, 1) + self.assertAlmostEqual(float(kernel.sum()), 1.0, places=5) + self.assertEqual(int(kernel.argmax()), kernel.numel() // 2) + + def test_a_non_positive_sigma_gives_the_identity_kernel(self): + self.assertTrue(torch.equal(gaussian_kernel1d(0.0, torch.device("cpu")), torch.tensor([1.0]))) + + +class TestRandomBiasField(SmaugLabTestCase): + def test_the_field_is_strictly_positive(self): + field = random_bias_field3d((8, 8, 8), std=0.7, scale=0.025, device=torch.device("cpu")) + self.assertTrue(bool((field > 0).all()), "a log-space Gaussian field must exponentiate to positive values") + + def test_the_shape_follows_batch_and_channels(self): + field = random_bias_field3d((6, 7, 8), std=0.5, scale=0.1, device=torch.device("cpu"), batch=3, channels=2) + self.assertEqual(tuple(field.shape), (3, 2, 6, 7, 8)) + + def test_a_zero_std_disables_the_field(self): + field = random_bias_field3d((5, 5, 5), std=0.0, scale=0.025, device=torch.device("cpu")) + self.assertTrue(torch.equal(field, torch.ones_like(field))) + + def test_the_field_is_smooth(self): + """Coarse grid plus trilinear upsampling: neighbours must be close.""" + torch.manual_seed(0) + field = random_bias_field3d((32, 32, 32), std=0.7, scale=0.025, device=torch.device("cpu"))[0, 0] + largest_step = max(float(field.diff(dim=d).abs().max()) for d in range(3)) + self.assertLess(largest_step, 0.5) + + def test_synthseg_and_domain_transfer_now_share_one_implementation(self): + from smauglab.transforms.gpu.domain_transfer import _random_bias_field3d + from smauglab.transforms.synthseg.functional import bias_field + + torch.manual_seed(7) + local = _random_bias_field3d((6, 6, 6), 0.7, 0.025, torch.device("cpu"), torch.float32) + self.assertEqual(tuple(local.shape), (6, 6, 6)) + + torch.manual_seed(7) + applied = bias_field(torch.ones(1, 1, 6, 6, 6), bias_field_std=0.7, bias_scale=0.025) + self.assertTrue(torch.allclose(applied[0, 0], local, atol=1e-6))