Align center cropping with centered padding for odd size differences

dm_pix/_src/augment.py

@@ -266,8 +266,8 @@ def center_crop(
   Target height/width given can be greater than the current size of the image
   which results in being a no-op for that dimension.
 
-  In case of odd size along any dimension the bottom/right side gets the extra
-  pixel.
+  When an odd number of pixels is removed along a dimension, the bottom/right
+  side loses the extra pixel, matching the convention used by `pad_to_size`.
 
   Args:
     image: a JAX array representing an image. Assumes that the image is either
@@ -285,11 +285,9 @@ def center_crop(
   batch, current_height, current_width, channel = _get_dimension_values(
       image=image, channel_axis=channel_axis
   )
-  center_h, center_w = current_height // 2, current_width // 2
-
-  left = max(center_w - (width // 2), 0)  # pyrefly: ignore[unsupported-operation]
+  left = max((current_width - width) // 2, 0)  # pyrefly: ignore[unsupported-operation]
   right = min(left + width, current_width)
-  top = max(center_h - (height // 2), 0)  # pyrefly: ignore[unsupported-operation]
+  top = max((current_height - height) // 2, 0)  # pyrefly: ignore[unsupported-operation]
   bottom = min(top + height, current_height)
 
   if _channels_last(image, channel_axis):

dm_pix/_src/augment_test.py

@@ -695,6 +695,69 @@ def test_negative_scalar_sigma_for_gaussian_blur(self):
       augment.gaussian_blur(image, sigma=-1.0, kernel_size=3)
 
 
+class CropAlignmentTest(parameterized.TestCase):
+
+  @parameterized.product(
+      sizes=[
+          ((4, 6), (3, 5)),
+          ((5, 7), (2, 4)),
+          ((4, 4), (3, 3)),
+          ((3, 3), (2, 2)),
+          ((4, 5), (6, 3)),
+          ((1, 4), (1, 1)),
+      ],
+      batched=[False, True],
+      channels_first=[False, True],
+  )
+  def test_center_crop_matches_integer_margin_reference(
+      self, sizes, batched, channels_first
+  ):
+    (height, width), (out_h, out_w) = sizes
+    image = np.arange(height * width * 2, dtype=np.float32).reshape(
+        height, width, 2
+    )
+    top, left = max((height - out_h) // 2, 0), max((width - out_w) // 2, 0)
+    expected = image[
+        top : top + min(height, out_h), left : left + min(width, out_w)
+    ]
+    if batched:
+      image = np.stack([image, image + 100])
+      expected = np.stack([expected, expected + 100])
+    if channels_first:
+      image = np.moveaxis(image, -1, -3)
+      expected = np.moveaxis(expected, -1, -3)
+    axis = -3 if channels_first else -1
+    fn = lambda x: augment.center_crop(x, out_h, out_w, channel_axis=axis)
+    for evaluate in (fn, jax.jit(fn)):
+      np.testing.assert_array_equal(evaluate(jnp.asarray(image)), expected)
+
+  @parameterized.parameters(
+      (3, 5, 4, 6), (3, 4, 6, 7), (4, 3, 5, 6), (1, 1, 2, 2)
+  )
+  def test_padding_then_cropping_returns_every_original_pixel(
+      self, h, w, ph, pw
+  ):
+    image = jnp.arange(h * w * 3, dtype=jnp.float32).reshape(h, w, 3)
+    padded = augment.pad_to_size(image, ph, pw)
+    recovered = augment.center_crop(padded, h, w)
+    np.testing.assert_array_equal(recovered, image)
+    gradient = jax.grad(
+        lambda x: augment.center_crop(augment.pad_to_size(x, ph, pw), h, w).sum()
+    )(image)
+    np.testing.assert_array_equal(gradient, jnp.ones_like(image))
+
+  @parameterized.parameters((4, 6, 3, 5), (4, 5, 7, 2), (7, 4, 2, 7))
+  def test_resize_with_crop_or_pad_matches_tensorflow(self, h, w, th, tw):
+    image = np.arange(2 * h * w * 3, dtype=np.float32).reshape(2, h, w, 3)
+    expected = tf.image.resize_with_crop_or_pad(image, th, tw).numpy()
+    actual = augment.resize_with_crop_or_pad(jnp.asarray(image), th, tw)
+    np.testing.assert_array_equal(actual, expected)
+
+
+if __name__ == "__main__":
+  absltest.main()
+
+
 if __name__ == "__main__":
   os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
   jax.config.update("jax_default_matmul_precision", "float32")