From af112902cfd31bea74d6cc2e22601cf718bb90ee Mon Sep 17 00:00:00 2001 From: Krakken190 Date: Wed, 21 Oct 2026 17:20:51 +0530 Subject: [PATCH 1/3] Add channel_axis to rotate and affine_transform. --- dm_pix/_src/augment.py | 31 +++++++++++++++++++++++------- dm_pix/_src/augment_test.py | 38 +++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 7 deletions(-) diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py index a9ead1c..8889d43 100644 --- a/dm_pix/_src/augment.py +++ b/dm_pix/_src/augment.py @@ -618,6 +618,7 @@ def affine_transform( order: int = 1, mode: str = "nearest", cval: float = 0.0, + channel_axis: int = -1, ) -> chex.Array: """Applies an affine transformation given by matrix. @@ -634,8 +635,10 @@ def affine_transform( image: a JAX array representing an image. Assumes that the image is either HWC or CHW. matrix: the inverse coordinate transformation matrix, mapping output - coordinates to input coordinates. If ndim is the number of dimensions of - input, the given matrix must have one of the following shapes: + coordinates to input coordinates. The matrix is always expressed in HWC + axis order, independently of `channel_axis`. If ndim is the number of + dimensions of input, the given matrix must have one of the following + shapes: - (ndim, ndim): the linear transformation matrix for each output coordinate. @@ -649,7 +652,7 @@ def affine_transform( offset: the offset into the array where the transform is applied. If a float, offset is the same for each axis. If an array, offset should - contain one value for each axis. + contain one value for each axis, in HWC order. order: the order of the spline interpolation, default is 1. The order has to be in the range [0-1]. Note that PIX interpolation will only be used for order=1, for other values we use `jax.scipy.ndimage.map_coordinates`. @@ -661,6 +664,7 @@ def affine_transform( accelerators, so use it with care. cval: value to fill past edges of input if mode is 'constant'. Default is 0.0. + channel_axis: the index of the channel axis. Returns: The input image transformed by the given matrix. @@ -742,6 +746,9 @@ def affine_transform( f"The affine matrix provided has shape {matrix.shape}.") raise ValueError(error_msg) + if channel_axis != -1: + image = jnp.moveaxis(image, source=channel_axis, destination=-1) + meshgrid = jnp.meshgrid(*[jnp.arange(size) for size in image.shape], indexing="ij") indices = jnp.concatenate( @@ -763,7 +770,12 @@ def affine_transform( order=order, cval=cval, ) - return interpolate_function(image, coordinates) + transformed_image = interpolate_function(image, coordinates) + + if channel_axis != -1: # Set channel axis back to original index. + transformed_image = jnp.moveaxis( + transformed_image, source=-1, destination=channel_axis) + return transformed_image def rotate( @@ -773,6 +785,7 @@ def rotate( order: int = 1, mode: str = "nearest", cval: float = 0.0, + channel_axis: int = -1, ) -> chex.Array: """Rotates an image around its center using interpolation. @@ -786,6 +799,7 @@ def rotate( its boundaries. Default is 'nearest'. See `affine_transform` for details. cval: value to fill past edges of input if mode is 'constant'. Default is 0.0. + channel_axis: the index of the channel axis. Returns: The rotated image. @@ -797,12 +811,15 @@ def rotate( s = jnp.sin(angle) matrix = jnp.array([[c, s, 0], [-s, c, 0], [0, 0, 1]]) - # Use the offset to place the rotation at the image center. - image_center = (jnp.asarray(image.shape) - 1.) / 2. + # Use the offset to place the rotation at the image center. The matrix and + # offset are expressed in HWC order, so the center is computed in that order. + hwc_shape = list(image.shape) + hwc_shape.append(hwc_shape.pop(channel_axis)) + image_center = (jnp.asarray(hwc_shape) - 1.) / 2. offset = image_center - matrix @ image_center return affine_transform(image, matrix, offset=offset, order=order, mode=mode, - cval=cval) + cval=cval, channel_axis=channel_axis) def random_flip_left_right( diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py index 8de2143..e56f94b 100644 --- a/dm_pix/_src/augment_test.py +++ b/dm_pix/_src/augment_test.py @@ -522,6 +522,44 @@ def _test_fn(self, images_list, jax_fn, reference_fn): class TestCustom(parameterized.TestCase): """Tests custom logic that is not covered by reference functions.""" + @parameterized.product( + mode_and_order=[("nearest", 0), ("nearest", 1), ("mirror", 1), + ("constant", 1)], + channel_axis=(0, -3), + ) + def test_affine_transform_channel_axis(self, mode_and_order, channel_axis): + mode, order = mode_and_order + image = np.random.RandomState(0).rand(10, 12, 3).astype(np.float32) + # Matrix and offset are always given in HWC order. + matrix = jnp.array([[0.8, 0.3, 0.], [-0.2, 1.1, 0.], [0., 0., 1.]]) + offset = jnp.array([1.5, -2., 0.]) + expected = augment.affine_transform( + image, matrix, offset=offset, mode=mode, order=order, cval=0.5) + result = augment.affine_transform( + np.moveaxis(image, -1, 0), matrix, offset=offset, mode=mode, + order=order, cval=0.5, channel_axis=channel_axis) + np.testing.assert_allclose( + result, np.moveaxis(expected, -1, 0), atol=1e-5) + + @parameterized.product( + mode_and_order=[("nearest", 0), ("nearest", 1), ("mirror", 1), + ("constant", 1)], + angle=(np.pi / 2, -np.pi / 4), + channel_axis=(0, -3), + ) + def test_rotate_channel_axis(self, mode_and_order, angle, channel_axis): + mode, order = mode_and_order + image = np.random.RandomState(0).rand(10, 12, 3).astype(np.float32) + expected = augment.rotate(image, angle, mode=mode, order=order, cval=0.5) + rotate_fn = jax.jit( + augment.rotate, + static_argnames=("mode", "order", "cval", "channel_axis")) + result = rotate_fn( + np.moveaxis(image, -1, 0), angle, mode=mode, order=order, cval=0.5, + channel_axis=channel_axis) + np.testing.assert_allclose( + result, np.moveaxis(expected, -1, 0), atol=1e-5) + @parameterized.product( images_list=(_RAND_FLOATS_IN_RANGE, _RAND_FLOATS_OUT_OF_RANGE), height=(250, 200), From e76e336da62ec9bae11873b8d24c2a1d0c3d5ee7 Mon Sep 17 00:00:00 2001 From: Krakken190 Date: Sun, 25 Oct 2026 15:29:18 +0530 Subject: [PATCH 2/3] Fix decorator syntax in augment_test.py --- dm_pix/_src/augment_test.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py index 887cf7b..8928ea9 100644 --- a/dm_pix/_src/augment_test.py +++ b/dm_pix/_src/augment_test.py @@ -559,6 +559,8 @@ def test_rotate_channel_axis(self, mode_and_order, angle, channel_axis): channel_axis=channel_axis) np.testing.assert_allclose( result, np.moveaxis(expected, -1, 0), atol=1e-5) + + @parameterized.product( shape_and_axis=( ((5, 7, 3), 2), ((5, 7, 3), -1), From 7dab6fafef4a1c988951a4019d16c1c0b95906df Mon Sep 17 00:00:00 2001 From: Krakken190 Date: Thu, 29 Oct 2026 10:46:09 +0530 Subject: [PATCH 3/3] Test channel_axis in existing affine_transform and rotate tests --- dm_pix/_src/augment_test.py | 114 ++++++++++++++++++------------------ 1 file changed, 56 insertions(+), 58 deletions(-) diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py index 1790e69..6624fbc 100644 --- a/dm_pix/_src/augment_test.py +++ b/dm_pix/_src/augment_test.py @@ -202,22 +202,46 @@ def test_random_flip_left_right_channel_axis(self): # Due to a bug in scipy we cannot test all available modes, refer to these # issues for more information: https://github.com/jax-ml/jax/issues/11097, # https://github.com/jax-ml/jax/issues/11097 - @parameterized.named_parameters( - ("in_range_nearest_0", _RAND_FLOATS_IN_RANGE, "nearest", 0), - ("in_range_nearest_1", _RAND_FLOATS_IN_RANGE, "nearest", 1), - ("in_range_mirror_1", _RAND_FLOATS_IN_RANGE, "mirror", 1), - ("out_of_range_nearest_0", _RAND_FLOATS_OUT_OF_RANGE, "nearest", 0), - ("out_of_range_nearest_1", _RAND_FLOATS_OUT_OF_RANGE, "nearest", 1), - ("out_of_range_mirror_1", _RAND_FLOATS_OUT_OF_RANGE, "mirror", 1), + @parameterized.product( + parameters_base=[ + ("in_range", "nearest", 0), + ("in_range", "nearest", 1), + ("in_range", "mirror", 1), + ("out_of_range", "nearest", 0), + ("out_of_range", "nearest", 1), + ("out_of_range", "mirror", 1), + ], + channel_axis=(-1, 0), ) - def test_affine_transform(self, images_list, mode, order): + def test_affine_transform(self, parameters_base, channel_axis): + images_list_type, mode, order = parameters_base + if images_list_type == "in_range": + images_list = _RAND_FLOATS_IN_RANGE + elif images_list_type == "out_of_range": + images_list = _RAND_FLOATS_OUT_OF_RANGE + else: + raise ValueError(f"{images_list_type} not a valid image list for tests.") + + def reference_fn(image, **kwargs): + # The matrix and offset are given in HWC order, so the SciPy reference + # is applied to the HWC version of the image. + if channel_axis == -1: + return scipy.ndimage.affine_transform(image, **kwargs) + transformed = scipy.ndimage.affine_transform( + np.moveaxis(image, channel_axis, -1), **kwargs) + return np.moveaxis(transformed, -1, channel_axis) + + if channel_axis != -1: + images_list = [np.moveaxis(x, -1, channel_axis) for x in images_list] + # (ndim, ndim) no offset self._test_fn( images_list, jax_fn=functools.partial( - augment.affine_transform, matrix=np.eye(3), mode=mode, order=order), + augment.affine_transform, matrix=np.eye(3), mode=mode, order=order, + channel_axis=channel_axis), reference_fn=functools.partial( - scipy.ndimage.affine_transform, + reference_fn, matrix=np.eye(3), order=order, mode=mode)) @@ -232,9 +256,10 @@ def test_affine_transform(self, images_list, mode, order): matrix=matrix, mode=mode, offset=offset, - order=order), + order=order, + channel_axis=channel_axis), reference_fn=functools.partial( - scipy.ndimage.affine_transform, + reference_fn, matrix=matrix, offset=offset, order=order, @@ -246,9 +271,10 @@ def test_affine_transform(self, images_list, mode, order): self._test_fn( images_list, jax_fn=functools.partial( - augment.affine_transform, matrix=matrix, mode=mode, order=order), + augment.affine_transform, matrix=matrix, mode=mode, order=order, + channel_axis=channel_axis), reference_fn=functools.partial( - scipy.ndimage.affine_transform, + reference_fn, matrix=matrix, order=order, mode=mode)) @@ -258,9 +284,10 @@ def test_affine_transform(self, images_list, mode, order): self._test_fn( images_list, jax_fn=functools.partial( - augment.affine_transform, matrix=matrix, mode=mode, order=order), + augment.affine_transform, matrix=matrix, mode=mode, order=order, + channel_axis=channel_axis), reference_fn=functools.partial( - scipy.ndimage.affine_transform, + reference_fn, matrix=matrix, order=order, mode=mode)) @@ -270,9 +297,10 @@ def test_affine_transform(self, images_list, mode, order): self._test_fn( images_list, jax_fn=functools.partial( - augment.affine_transform, matrix=matrix, mode=mode, order=order), + augment.affine_transform, matrix=matrix, mode=mode, order=order, + channel_axis=channel_axis), reference_fn=functools.partial( - scipy.ndimage.affine_transform, + reference_fn, matrix=matrix, order=order, mode=mode)) @@ -290,8 +318,9 @@ def test_affine_transform(self, images_list, mode, order): ], cval=(0.0, 1.0, -2.0), angle=(0.0, np.pi / 4, -np.pi / 4), + channel_axis=(-1, 0), ) - def test_rotate(self, parameters_base, cval, angle): + def test_rotate(self, parameters_base, cval, angle, channel_axis): images_list_type, mode, order = parameters_base if images_list_type == "in_range": images_list = _RAND_FLOATS_IN_RANGE @@ -299,13 +328,20 @@ def test_rotate(self, parameters_base, cval, angle): images_list = _RAND_FLOATS_OUT_OF_RANGE else: raise ValueError(f"{images_list_type} not a valid image list for tests.") + if channel_axis == -1: + axes = (1, 0) # SciPy's default, rotating the H and W axes of HWC. + else: + images_list = [np.moveaxis(x, -1, channel_axis) for x in images_list] + axes = (2, 1) # The H and W axes of CHW. self._test_fn( images_list, jax_fn=functools.partial( - augment.rotate, angle=angle, mode=mode, order=order, cval=cval), + augment.rotate, angle=angle, mode=mode, order=order, cval=cval, + channel_axis=channel_axis), reference_fn=functools.partial( scipy.ndimage.rotate, angle=angle * 180 / np.pi, # SciPy uses degrees. + axes=axes, order=order, mode=mode, cval=cval, @@ -522,44 +558,6 @@ def _test_fn(self, images_list, jax_fn, reference_fn): class TestCustom(parameterized.TestCase): """Tests custom logic that is not covered by reference functions.""" - @parameterized.product( - mode_and_order=[("nearest", 0), ("nearest", 1), ("mirror", 1), - ("constant", 1)], - channel_axis=(0, -3), - ) - def test_affine_transform_channel_axis(self, mode_and_order, channel_axis): - mode, order = mode_and_order - image = np.random.RandomState(0).rand(10, 12, 3).astype(np.float32) - # Matrix and offset are always given in HWC order. - matrix = jnp.array([[0.8, 0.3, 0.], [-0.2, 1.1, 0.], [0., 0., 1.]]) - offset = jnp.array([1.5, -2., 0.]) - expected = augment.affine_transform( - image, matrix, offset=offset, mode=mode, order=order, cval=0.5) - result = augment.affine_transform( - np.moveaxis(image, -1, 0), matrix, offset=offset, mode=mode, - order=order, cval=0.5, channel_axis=channel_axis) - np.testing.assert_allclose( - result, np.moveaxis(expected, -1, 0), atol=1e-5) - - @parameterized.product( - mode_and_order=[("nearest", 0), ("nearest", 1), ("mirror", 1), - ("constant", 1)], - angle=(np.pi / 2, -np.pi / 4), - channel_axis=(0, -3), - ) - def test_rotate_channel_axis(self, mode_and_order, angle, channel_axis): - mode, order = mode_and_order - image = np.random.RandomState(0).rand(10, 12, 3).astype(np.float32) - expected = augment.rotate(image, angle, mode=mode, order=order, cval=0.5) - rotate_fn = jax.jit( - augment.rotate, - static_argnames=("mode", "order", "cval", "channel_axis")) - result = rotate_fn( - np.moveaxis(image, -1, 0), angle, mode=mode, order=order, cval=0.5, - channel_axis=channel_axis) - np.testing.assert_allclose( - result, np.moveaxis(expected, -1, 0), atol=1e-5) - @parameterized.product( shape_and_axis=( ((5, 7, 3), 2),