fix: centre the Gaussian kernel and correct the CPU Scharr x-kernel

smauglab/transforms/cpu/contrast.py

@@ -30,7 +30,10 @@ def get_parameters(self, **data_dict) -> dict:
             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":
-                kernel_x = torch.tensor([[-3, 0, 3], [-10, 0, -10], [-3, 0, 3]], dtype=torch.float32)
+                # 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:

smauglab/transforms/gpu/contrast.py

@@ -326,9 +326,15 @@ def apply_convolution(img: torch.Tensor, kernel: torch.Tensor, dim: int) -> torc
 
 
 def get_gaussian_kernel1d(kernel_size: int, sigma: Union[float, Tensor], dtype: torch.dtype, device: torch.device) -> Tensor:
-    """Create a 1D Gaussian kernel."""
+    """Create a 1D Gaussian kernel, centred on the middle tap.
 
-    x = torch.arange(kernel_size, dtype=dtype, device=device)
+    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()
 

unit_tests/test_kernel_correctness.py

@@ -0,0 +1,123 @@
+"""Convolution kernels that were not what their name says.
+
+* The 1-D Gaussian was sampled at `arange(k)`, so its peak sat at index 0 and the 3-D
+  kernel built from it had its maximum at corner [0, 0, 0]. A blur through that kernel
+  translates the image by about a voxel as well as blurring it -- while the
+  segmentation mask, which is never convolved, stays put.
+* The 2-D CPU Scharr x-kernel had `[-10, 0, -10]` as its middle row, so it summed to
+  -20 instead of 0 and was not a gradient operator.
+"""
+
+import torch
+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 unit_tests.helpers import SmaugLabTestCase
+
+CPU = torch.device("cpu")
+
+
+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")
+
+    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)
+                self.assertTrue(torch.allclose(kernel, kernel.flip(0), atol=1e-6))
+                self.assertAlmostEqual(float(kernel.sum()), 1.0, places=5)
+
+    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)
+                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)
+        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)
+        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)
+
+        blurred = F.conv3d(volume, kernel.view(1, 1, 7, 7, 7), padding=3)
+
+        grid = torch.arange(15, dtype=torch.float32)
+        for axis in (2, 3, 4):
+            with self.subTest(axis=axis):
+                marginal = blurred.sum(dim=[d for d in (2, 3, 4) if d != axis]).flatten()
+                centre_of_mass = float((marginal * grid).sum() / marginal.sum())
+                self.assertAlmostEqual(centre_of_mass, 7.0, places=3, msg=f"blur shifted the impulse along axis {axis}")
+
+    def test_the_old_uncentred_formula_really_was_off_centre(self):
+        """Pin down what was wrong, so nobody reintroduces it as a 'simplification'."""
+        kernel_size, sigma = 3, 1.0
+        x = torch.arange(kernel_size, dtype=torch.float32)  # the old sample points
+        pdf = torch.exp(-0.5 * (x / sigma).pow(2))
+        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)
+
+
+class TestScharrIsAGradientOperator(SmaugLabTestCase):
+    def _kernels(self, spatial_dims: int):
+        image = torch.rand(1, *([8] * spatial_dims))
+        return ConvTransform(kernel_type="Scharr").get_parameters(image=image)["kernel"]
+
+    def test_every_2d_scharr_kernel_sums_to_zero(self):
+        for axis, kernel in enumerate(self._kernels(2)):
+            with self.subTest(axis=axis):
+                self.assertAlmostEqual(float(kernel.sum()), 0.0, places=5)
+
+    def test_every_3d_scharr_kernel_sums_to_zero(self):
+        for axis, kernel in enumerate(self._kernels(3)):
+            with self.subTest(axis=axis):
+                self.assertAlmostEqual(float(kernel.sum()), 0.0, places=5)
+
+    def test_the_2d_kernels_are_antisymmetric_about_their_axis(self):
+        """A gradient operator negates when its differencing axis is flipped."""
+        kernel_x, kernel_y = self._kernels(2)
+        self.assertTrue(torch.allclose(kernel_x, -kernel_x.flip(1), atol=1e-6))
+        self.assertTrue(torch.allclose(kernel_y, -kernel_y.flip(0), atol=1e-6))
+
+    def test_a_constant_image_gives_no_gradient(self):
+        constant = torch.full((1, 1, 7, 7), 4.0)
+        for axis, kernel in enumerate(self._kernels(2)):
+            with self.subTest(axis=axis):
+                response = F.conv2d(constant, kernel.view(1, 1, 3, 3))
+                self.assertTrue(torch.allclose(response, torch.zeros_like(response), atol=1e-4))
+
+    def test_the_2d_x_kernel_responds_to_a_horizontal_edge(self):
+        """Sanity: it must still detect the thing it is for."""
+        image = torch.zeros(1, 1, 7, 7)
+        image[0, 0, :, 4:] = 1.0
+        kernel_x = self._kernels(2)[0]
+
+        response = F.conv2d(image, kernel_x.view(1, 1, 3, 3))
+
+        self.assertGreater(float(response.abs().max()), 1.0)
+
+    def test_the_laplace_kernels_still_sum_to_zero(self):
+        """Untouched by this change, but the neighbouring branch of the same method."""
+        for spatial_dims in (2, 3):
+            with self.subTest(spatial_dims=spatial_dims):
+                image = torch.rand(1, *([8] * spatial_dims))
+                kernel = ConvTransform(kernel_type="Laplace").get_parameters(image=image)["kernel"]
+                self.assertAlmostEqual(float(kernel.sum()), 0.0, places=5)