From aa0f239fe49c848a3690516fa6b2b50aa4a8118a Mon Sep 17 00:00:00 2001 From: PIXDev Date: Tue, 13 Oct 2026 02:21:37 -0700 Subject: [PATCH] Return the input image unmodified if sigma is zero in gaussian blur. PiperOrigin-RevId: 998485835 --- dm_pix/_src/augment.py | 39 ++++--------------------------------- dm_pix/_src/augment_test.py | 26 ------------------------- 2 files changed, 4 insertions(+), 61 deletions(-) diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py index d462506..56e0483 100644 --- a/dm_pix/_src/augment.py +++ b/dm_pix/_src/augment.py @@ -480,9 +480,7 @@ 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. 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). + sigma: the standard deviation (in pixels) of the gaussian kernel. 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. @@ -494,40 +492,13 @@ 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) - - return jax.lax.cond( - jnp.all(jnp.isclose(sigma, 0.0)), - lambda x: x, - 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.0 * sigma**2)) + blur_filter = jnp.exp(-x**2 / (2. * 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]) @@ -543,16 +514,14 @@ def _gaussian_blur( strides=(1, 1), padding=padding, channel_axis=channel_axis, - dimension_numbers=dimension_numbers, - ) + dimension_numbers=dimension_numbers) blurred = _depthwise_conv2d( blurred, kernel=blur_v, strides=(1, 1), padding=padding, channel_axis=channel_axis, - dimension_numbers=dimension_numbers, - ) + dimension_numbers=dimension_numbers) if expand_batch_dim: blurred = jnp.squeeze(blurred, axis=0) return blurred diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py index 74edf09..df6f3d9 100644 --- a/dm_pix/_src/augment_test.py +++ b/dm_pix/_src/augment_test.py @@ -605,32 +605,6 @@ 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) - - 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"