Return the input image unmodified if sigma is zero in gaussian blur.

dm_pix/_src/augment.py

@@ -480,7 +480,9 @@ def gaussian_blur(
   Args:
     image: the input image, as a [0-1] float tensor. Should have 3 or 4
       dimensions with two spatial dimensions.
-    sigma: the standard deviation (in pixels) of the gaussian kernel.
+    sigma: the standard deviation (in pixels) of the gaussian kernel. If it is
+      0, the image is returned without blurring. If it is a negative scalar, an
+      error is raised (behavior is undefined for negative traced/array values).
     kernel_size: the size (in pixels) of the square gaussian kernel. Will be
       "rounded" to the next odd integer.
     padding: either "SAME" or "VALID", passed to the underlying convolution.
@@ -492,13 +494,49 @@ def gaussian_blur(
   # DO NOT REMOVE - Logging usage.
 
   chex.assert_rank(image, {3, 4})
+  chex.assert_size(sigma, 1)
+  if isinstance(sigma, chex.Scalar):
+    chex.assert_scalar_non_negative(sigma)
+
+  radius = int(kernel_size / 2)
+  if padding.upper() == "VALID" and radius > 0:
+    if _channels_last(image, channel_axis):
+      true_fun = lambda x: x[..., radius:-radius, radius:-radius, :]
+    else:
+      true_fun = lambda x: x[..., radius:-radius, radius:-radius]
+  else:
+    true_fun = lambda x: x
+
+  return jax.lax.cond(
+      jnp.all(jnp.isclose(sigma, 0.0)),
+      true_fun,
+      lambda x: _gaussian_blur(
+          x,
+          sigma=sigma,
+          kernel_size=kernel_size,
+          padding=padding,
+          channel_axis=channel_axis,
+      ),
+      image,
+  )
+
+
+def _gaussian_blur(
+    image: chex.Array,
+    sigma: chex.Numeric,
+    kernel_size: float,
+    *,
+    padding: str = "SAME",
+    channel_axis: int = -1,
+) -> chex.Array:
+  """Internal helper for `gaussian_blur` when `sigma > 0`."""
   data_format = "NHWC" if _channels_last(image, channel_axis) else "NCHW"
   dimension_numbers = (data_format, "HWIO", data_format)
   num_channels = image.shape[channel_axis]
   radius = int(kernel_size / 2)
   kernel_size_ = 2 * radius + 1
   x = jnp.arange(-radius, radius + 1).astype(jnp.float32)
-  blur_filter = jnp.exp(-x**2 / (2. * sigma**2))
+  blur_filter = jnp.exp(-(x**2) / (2.0 * sigma**2))
   blur_filter = blur_filter / jnp.sum(blur_filter)
   blur_v = jnp.reshape(blur_filter, [kernel_size_, 1, 1, 1])
   blur_h = jnp.reshape(blur_filter, [1, kernel_size_, 1, 1])

dm_pix/_src/augment_test.py

@@ -605,6 +605,95 @@ def test_pad_to_size_batch_chw_shape(self):
     )
     self.assertEqual(result.shape, (2, 3, 6, 6))
 
+  @parameterized.named_parameters(
+      ("scalar", 0.0),
+      ("0d_array", lambda: jnp.array(0.0)),
+      ("1d_array", lambda: jnp.array([0.0])),
+  )
+  def test_zero_sigma_for_gaussian_blur(self, sigma):
+    sigma = sigma() if callable(sigma) else sigma
+    image = jax.random.uniform(jax.random.PRNGKey(0), shape=(4, 4, 3))
+    result = augment.gaussian_blur(image, sigma=sigma, kernel_size=3)
+    np.testing.assert_array_equal(result, image)
+    result_jit = jax.jit(augment.gaussian_blur, static_argnames="kernel_size")(
+        image, sigma=sigma, kernel_size=3
+    )
+    np.testing.assert_array_equal(result_jit, image)
+
+  @parameterized.named_parameters(
+      (
+          "channels_last_3d",
+          (6, 6, 3),
+          -1,
+          (slice(1, -1), slice(1, -1), slice(None)),
+      ),
+      (
+          "channels_last_4d",
+          (2, 6, 6, 3),
+          -1,
+          (slice(None), slice(1, -1), slice(1, -1), slice(None)),
+      ),
+      (
+          "channels_first_3d",
+          (3, 6, 6),
+          -3,
+          (slice(None), slice(1, -1), slice(1, -1)),
+      ),
+      (
+          "channels_first_4d",
+          (2, 3, 6, 6),
+          -3,
+          (slice(None), slice(None), slice(1, -1), slice(1, -1)),
+      ),
+  )
+  def test_zero_sigma_valid_padding_for_gaussian_blur(
+      self, shape, channel_axis, crop_slice
+  ):
+    image = jax.random.uniform(jax.random.PRNGKey(0), shape=shape)
+    expected = image[crop_slice]
+    result = augment.gaussian_blur(
+        image,
+        sigma=0.0,
+        kernel_size=3,
+        padding="VALID",
+        channel_axis=channel_axis,
+    )
+    np.testing.assert_array_equal(result, expected)
+    result_jit = jax.jit(
+        augment.gaussian_blur,
+        static_argnames=("kernel_size", "padding", "channel_axis"),
+    )(
+        image,
+        sigma=0.0,
+        kernel_size=3,
+        padding="VALID",
+        channel_axis=channel_axis,
+    )
+    np.testing.assert_array_equal(result_jit, expected)
+
+  def test_gaussian_blur_valid_padding_nonzero_sigma(self):
+    image = jax.random.uniform(jax.random.PRNGKey(0), shape=(6, 6, 3))
+    result = augment.gaussian_blur(
+        image, sigma=1.0, kernel_size=3, padding="VALID"
+    )
+    self.assertEqual(result.shape, (4, 4, 3))
+    result_jit = jax.jit(
+        augment.gaussian_blur, static_argnames=("kernel_size", "padding")
+    )(image, sigma=1.0, kernel_size=3, padding="VALID")
+    np.testing.assert_array_equal(result, result_jit)
+
+  def test_gaussian_blur_sigma_types(self):
+    image = jax.random.uniform(jax.random.PRNGKey(0), shape=(4, 4, 3))
+    expected = augment.gaussian_blur(image, sigma=0.5, kernel_size=3)
+    for sigma in (jnp.array(0.5), jnp.array([0.5])):
+      result = augment.gaussian_blur(image, sigma=sigma, kernel_size=3)
+      np.testing.assert_allclose(result, expected, rtol=1e-6, atol=1e-6)
+
+  def test_negative_scalar_sigma_for_gaussian_blur(self):
+    image = jnp.zeros((4, 4, 3))
+    with self.assertRaises(AssertionError):
+      augment.gaussian_blur(image, sigma=-1.0, kernel_size=3)
+
 
 if __name__ == "__main__":
   os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"