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))