Add channel_axis to rotate and affine_transform

dm_pix/_src/augment.py

@@ -620,6 +620,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.
 
@@ -636,8 +637,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.
@@ -651,7 +654,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`.
@@ -663,6 +666,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.
@@ -744,6 +748,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(
@@ -765,7 +772,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(
@@ -775,6 +787,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.
 
@@ -788,6 +801,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.
@@ -799,12 +813,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(

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(
       shape_and_axis=(
           ((5, 7, 3), 2),