diff --git a/dm_pix/_src/metrics.py b/dm_pix/_src/metrics.py index 6dfd6be..767e85f 100644 --- a/dm_pix/_src/metrics.py +++ b/dm_pix/_src/metrics.py @@ -278,9 +278,8 @@ def filter_fn_x(z): epsilon = jnp.finfo(jnp.float32).eps ** 2 sigma00 = jnp.maximum(epsilon, sigma00) sigma11 = jnp.maximum(epsilon, sigma11) - sigma01 = jnp.sign(sigma01) * jnp.minimum( - jnp.sqrt(sigma00 * sigma11), jnp.abs(sigma01) - ) + covariance_bound = jnp.sqrt(sigma00 * sigma11) + sigma01 = jnp.clip(sigma01, -covariance_bound, covariance_bound) c1 = (k1 * max_val) ** 2 c2 = (k2 * max_val) ** 2 diff --git a/dm_pix/_src/metrics_test.py b/dm_pix/_src/metrics_test.py index 63daab4..0bba63e 100644 --- a/dm_pix/_src/metrics_test.py +++ b/dm_pix/_src/metrics_test.py @@ -15,9 +15,11 @@ import os from absl.testing import absltest +from absl.testing import parameterized import chex from dm_pix._src import metrics import jax +import jax.numpy as jnp import numpy as np import tensorflow as tf @@ -206,6 +208,112 @@ def test_ssim_ignore_nans(self): assert not np.any(np.isnan(ssim)) +def _mean_filter(x): + return jnp.mean(x, axis=(-3, -2), keepdims=True) + + +def _smooth_ssim(a, b): + ma, mb = _mean_filter(a), _mean_filter(b) + va, vb = _mean_filter((a - ma) ** 2), _mean_filter((b - mb) ** 2) + covariance = _mean_filter((a - ma) * (b - mb)) + # These fixtures have valid covariance, so no clipping is required. + result = ( + (2 * ma * mb + 0.01**2) + * (2 * covariance + 0.03**2) + / ((ma**2 + mb**2 + 0.01**2) * (va + vb + 0.03**2)) + ) + return result.mean(axis=(-3, -2, -1)) + + +def _ssim_gradient_images(batch=False, channels=1): + a = jnp.array([[0.75, 0.75], [0.25, 0.25]])[..., None] + b = jnp.array([[0.75, 0.25], [0.75, 0.25]])[..., None] + a, b = jnp.repeat(a, channels, -1), jnp.repeat(b, channels, -1) + if batch: + a, b = jnp.stack([a, a]), jnp.stack([b, b]) + return a, b + + +class SsimCovarianceGradientTest(parameterized.TestCase): + + @parameterized.product( + batch=[False, True], channels=[1, 3], custom_filter=[False, True] + ) + def test_zero_covariance_gradient_matches_independent_formula( + self, batch, channels, custom_filter + ): + a, b = _ssim_gradient_images(batch, channels) + kwargs = ( + {"filter_fn": _mean_filter} + if custom_filter + else {"filter_size": 2, "filter_sigma": 1.0} + ) + objective = lambda x: metrics.ssim(a, x, **kwargs).sum() + reference = lambda x: _smooth_ssim(a, x).sum() + expected = jax.value_and_grad(reference)(b) + for fn in ( + jax.value_and_grad(objective), + jax.jit(jax.value_and_grad(objective)), + ): + actual = fn(b) + for value, target in zip(actual, expected): + np.testing.assert_allclose(value, target, rtol=3e-6, atol=1e-7) + _, tangent = jax.jvp(objective, (b,), (a - 0.5,)) + step = 0.01 + finite_diff = ( + objective(b + step * (a - 0.5)) - objective(b - step * (a - 0.5)) + ) / (2 * step) + np.testing.assert_allclose(tangent, finite_diff, rtol=5e-4, atol=1e-6) + + def test_default_gaussian_passes_gradients_at_a_zero_second_image(self): + a, b = _ssim_gradient_images() + b = jnp.zeros_like(b) + actual = jax.jit( + jax.value_and_grad(lambda x: metrics.ssim(a, x, filter_size=2)) + )(b) + expected = jax.value_and_grad(lambda x: _smooth_ssim(a, x))(b) + for value, target in zip(actual, expected): + np.testing.assert_allclose(value, target, rtol=3e-6, atol=1e-8) + + @parameterized.parameters(-0.5, 0.5) + def test_nonzero_covariance_values_and_gradients_are_preserved(self, weight): + a, b = _ssim_gradient_images() + b = b + weight * (a - 0.5) + actual = jax.value_and_grad( + lambda x: metrics.ssim(a, x, filter_fn=_mean_filter) + )(b) + expected = jax.value_and_grad(lambda x: _smooth_ssim(a, x))(b) + for value, target in zip(actual, expected): + np.testing.assert_allclose(value, target, rtol=3e-6, atol=1e-7) + + def test_correlation_parameter_can_learn_from_zero(self): + a, b = _ssim_gradient_images() + score = lambda parameter: metrics.ssim( + a, b + parameter * (a - 0.5), filter_size=2, filter_sigma=1.0 + ) + parameter = jnp.array(0.0) + derivative = jax.jit(jax.grad(score))(parameter) + self.assertGreater(float(derivative), 0.9) + self.assertGreater( + float(score(parameter + 0.1 * derivative)), float(score(parameter)) + ) + + @parameterized.parameters(False, True) + def test_map_and_identity_scores_keep_existing_values(self, negative): + a, _ = _ssim_gradient_images(batch=True) + b = 1 - a if negative else a + actual = metrics.ssim(a, b, filter_size=2, return_map=True) + expected = _smooth_ssim(a, b) + np.testing.assert_allclose( + actual.mean(axis=(-3, -2, -1)), expected, rtol=3e-6, atol=1e-7 + ) + self.assertEqual(actual.shape, (2, 1, 1, 1)) + + +if __name__ == "__main__": + absltest.main() + + if __name__ == "__main__": os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" absltest.main()