From 506ed2d5cce05ee085248eb22aec75050282f0d5 Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Mon, 5 Oct 2026 10:21:52 +0100 Subject: [PATCH 1/2] Keep achromatic HSL gradients finite Signed-off-by: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> --- dm_pix/_src/color_conversion.py | 14 +++-- dm_pix/_src/hsl_gradients_test.py | 87 +++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 5 deletions(-) create mode 100644 dm_pix/_src/hsl_gradients_test.py 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/hsl_gradients_test.py b/dm_pix/_src/hsl_gradients_test.py new file mode 100644 index 0000000..7aef6a8 --- /dev/null +++ b/dm_pix/_src/hsl_gradients_test.py @@ -0,0 +1,87 @@ +# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Regression tests for achromatic HSL gradient handling.""" + +import colorsys +from absl.testing import absltest +from absl.testing import parameterized +from dm_pix._src import color_conversion +import jax +import jax.numpy as jnp +import numpy as np + + +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) + if compiled: + grey, reverse, forward = map(jax.jit, (grey, reverse, forward)) + 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 = jax.jacrev(color_conversion.rgb_to_hsl)(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__': + absltest.main() From 52a5c62862f10dc80e67eed1041b11bb1f93bf3b Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:50:02 +0100 Subject: [PATCH 2/2] Move HSL gradient tests into color conversion suite --- dm_pix/_src/color_conversion_test.py | 64 ++++++++++++++++++++ dm_pix/_src/hsl_gradients_test.py | 87 ---------------------------- 2 files changed, 64 insertions(+), 87 deletions(-) delete mode 100644 dm_pix/_src/hsl_gradients_test.py 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() diff --git a/dm_pix/_src/hsl_gradients_test.py b/dm_pix/_src/hsl_gradients_test.py deleted file mode 100644 index 7aef6a8..0000000 --- a/dm_pix/_src/hsl_gradients_test.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Regression tests for achromatic HSL gradient handling.""" - -import colorsys -from absl.testing import absltest -from absl.testing import parameterized -from dm_pix._src import color_conversion -import jax -import jax.numpy as jnp -import numpy as np - - -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) - if compiled: - grey, reverse, forward = map(jax.jit, (grey, reverse, forward)) - 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 = jax.jacrev(color_conversion.rgb_to_hsl)(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__': - absltest.main()