refactor: one kernels module instead of five copies

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

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.

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:

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(