Keep achromatic HSL gradients finite

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)

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()