Add channel_axis to rotate and affine_transform

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(

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,22 +318,30 @@ 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
     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.")
+    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,