From af112902cfd31bea74d6cc2e22601cf718bb90ee Mon Sep 17 00:00:00 2001 From: Krakken190 Date: Wed, 21 Oct 2026 17:20:51 +0530 Subject: [PATCH 1/2] 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/2] 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),