diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py index 56e0483..a9ead1c 100644 --- a/dm_pix/_src/augment.py +++ b/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]) diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py index df6f3d9..8de2143 100644 --- a/dm_pix/_src/augment_test.py +++ b/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"