From 60723e8f741941410a489b5c7b2f0d1582e06dbc Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Tue, 6 Oct 2026 00:31:25 +0100
Subject: [PATCH 1/2] Align center cropping with centered padding for odd size
 differences

Signed-off-by: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com>
---
 dm_pix/_src/augment.py             | 10 ++--
 dm_pix/_src/crop_alignment_test.py | 85 ++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 6 deletions(-)
 create mode 100644 dm_pix/_src/crop_alignment_test.py

diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py
index a9ead1c..6e3aa1a 100644
--- a/dm_pix/_src/augment.py
+++ b/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):
diff --git a/dm_pix/_src/crop_alignment_test.py b/dm_pix/_src/crop_alignment_test.py
new file mode 100644
index 0000000..f1781cd
--- /dev/null
+++ b/dm_pix/_src/crop_alignment_test.py
@@ -0,0 +1,85 @@
+# 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.
+"""Center cropping should use the same pixel alignment as centered padding."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import dm_pix
+import jax
+import jax.numpy as jnp
+import numpy as np
+import tensorflow as tf
+
+
+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: dm_pix.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 = dm_pix.pad_to_size(image, ph, pw)
+    recovered = dm_pix.center_crop(padded, h, w)
+    np.testing.assert_array_equal(recovered, image)
+    gradient = jax.grad(
+        lambda x: dm_pix.center_crop(dm_pix.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 = dm_pix.resize_with_crop_or_pad(jnp.asarray(image), th, tw)
+    np.testing.assert_array_equal(actual, expected)
+
+
+if __name__ == "__main__":
+  absltest.main()

From 80ca8242e0c6ea35621af37170259feb16eac8cd Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Mon, 19 Oct 2026 16:50:05 +0100
Subject: [PATCH 2/2] Move crop alignment tests into augmentation suite

---
 dm_pix/_src/augment_test.py        | 63 ++++++++++++++++++++++
 dm_pix/_src/crop_alignment_test.py | 85 ------------------------------
 2 files changed, 63 insertions(+), 85 deletions(-)
 delete mode 100644 dm_pix/_src/crop_alignment_test.py

diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py
index 8de2143..f0634e6 100644
--- a/dm_pix/_src/augment_test.py
+++ b/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")
diff --git a/dm_pix/_src/crop_alignment_test.py b/dm_pix/_src/crop_alignment_test.py
deleted file mode 100644
index f1781cd..0000000
--- a/dm_pix/_src/crop_alignment_test.py
+++ /dev/null
@@ -1,85 +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.
-"""Center cropping should use the same pixel alignment as centered padding."""
-
-from absl.testing import absltest
-from absl.testing import parameterized
-import dm_pix
-import jax
-import jax.numpy as jnp
-import numpy as np
-import tensorflow as tf
-
-
-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: dm_pix.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 = dm_pix.pad_to_size(image, ph, pw)
-    recovered = dm_pix.center_crop(padded, h, w)
-    np.testing.assert_array_equal(recovered, image)
-    gradient = jax.grad(
-        lambda x: dm_pix.center_crop(dm_pix.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 = dm_pix.resize_with_crop_or_pad(jnp.asarray(image), th, tw)
-    np.testing.assert_array_equal(actual, expected)
-
-
-if __name__ == "__main__":
-  absltest.main()