diff --git a/dm_pix/_src/color_conversion.py b/dm_pix/_src/color_conversion.py index 0c47cf2..694fa90 100644 --- a/dm_pix/_src/color_conversion.py +++ b/dm_pix/_src/color_conversion.py @@ -199,17 +199,21 @@ def rgb_to_hsl( mask = c_min == c_max - rc = (c_max - red) / c_diff - gc = (c_max - green) / c_diff - bc = (c_max - blue) / c_diff + # Mask denominators before division so achromatic branches also have + # finite gradients, rather than masking undefined quotients afterward. + safe_diff = jnp.where(mask, 1., c_diff) + rc = (c_max - red) / safe_diff + gc = (c_max - green) / safe_diff + bc = (c_max - blue) / safe_diff eps = jnp.finfo(jnp.float32).eps h = jnp.where( mask, 0, (jnp.where(red == c_max, bc - gc, jnp.where(green == c_max, 2 + rc - bc, 4 + gc - rc)) / 6) % 1) - s = jnp.where(mask, 0, (c_diff + eps) / - (2 * eps + jnp.where(c_sum <= 1, c_sum, 2 - c_sum))) + saturation_denominator = jnp.where( + mask, 1., 2 * eps + jnp.where(c_sum <= 1, c_sum, 2 - c_sum)) + s = jnp.where(mask, 0, (c_diff + eps) / saturation_denominator) l = c_sum / 2 return jnp.stack([h, s, l], axis=-1) diff --git a/dm_pix/_src/color_conversion_test.py b/dm_pix/_src/color_conversion_test.py index 8f9a41e..eeac27e 100644 --- a/dm_pix/_src/color_conversion_test.py +++ b/dm_pix/_src/color_conversion_test.py @@ -283,6 +283,70 @@ def test_grayscale(self, test_images, keep_dims, channel_last): grayscale_jax, grayscale_tf, atol=1E-5, rtol=1E-5) +class HslGradientsTest(parameterized.TestCase): + + @parameterized.product( + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + compiled=(False, True), + ) + def test_greyscale_path_has_finite_lightness_gradients(self, dtype, compiled): + def grey(level): + return color_conversion.rgb_to_hsl(jnp.repeat(level[None], 3)) + + reverse = jax.jacrev(grey) + forward = jax.jacfwd(grey) + rgb_jacobian = jax.jacrev(color_conversion.rgb_to_hsl) + if compiled: + grey, reverse, forward, rgb_jacobian = map( + jax.jit, (grey, reverse, forward, rgb_jacobian) + ) + for level in (0.0, 0.25, 0.5, 1.0): + x = jnp.array(level, dtype=dtype) + np.testing.assert_array_equal(grey(x), np.array([0.0, 0.0, level])) + np.testing.assert_allclose(reverse(x), [0.0, 0.0, 1.0], atol=1e-3) + np.testing.assert_allclose(forward(x), [0.0, 0.0, 1.0], atol=1e-3) + jacobian = rgb_jacobian(jnp.repeat(x[None], 3)) + self.assertTrue(np.isfinite(jacobian).all()) + # The existing achromatic hue/saturation branches are constants. + np.testing.assert_array_equal(jacobian[:2], np.zeros((2, 3))) + + @parameterized.parameters(False, True) + def test_mixed_image_values_and_gradients(self, compiled): + image = jnp.array( + [[[0.0, 0.0, 0.0], [0.2, 0.6, 0.9]], [[0.5, 0.5, 0.5], [0.8, 0.3, 0.1]]] + ) + f = color_conversion.rgb_to_hsl + gradient = jax.grad(lambda x: jnp.sum(f(x))) + if compiled: + f, gradient = jax.jit(f), jax.jit(gradient) + expected = np.array( + [ + colorsys.rgb_to_hls(*pixel) + for pixel in np.asarray(image).reshape(-1, 3) + ] + ) + expected = expected[:, [0, 2, 1]].reshape(image.shape) + np.testing.assert_allclose(f(image), expected, rtol=2e-6, atol=2e-7) + self.assertTrue(np.isfinite(gradient(image)).all()) + + def test_chromatic_gradients_match_finite_differences(self): + x = jnp.array([0.2, 0.6, 0.9]) + numerical = np.zeros((3, 3)) + x64 = np.asarray(x).astype(np.float64) + for index in range(3): + delta = np.zeros(3) + delta[index] = 1e-5 + upper = np.array(colorsys.rgb_to_hls(*(x64 + delta)))[[0, 2, 1]] + lower = np.array(colorsys.rgb_to_hls(*(x64 - delta)))[[0, 2, 1]] + numerical[:, index] = (upper - lower) / (2e-5) + np.testing.assert_allclose( + jax.jacrev(color_conversion.rgb_to_hsl)(x), + numerical, + rtol=2e-5, + atol=1e-6, + ) + + if __name__ == "__main__": os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" tf.test.main()