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/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")