Add elastic_transform processing for image.py

keras/src/backend/jax/image.py

@@ -5,6 +5,7 @@
 
 from keras.src import backend
 from keras.src.backend.jax.core import convert_to_tensor
+from keras.src.random.seed_generator import draw_seed
 
 RESIZE_INTERPOLATIONS = (
     "bilinear",
@@ -729,3 +730,116 @@ def _get_gaussian_kernel2d(size, sigma):
         blurred_images = blurred_images.squeeze(axis=0)
 
     return blurred_images
+
+
+def elastic_transform(
+    images,
+    alpha=20.0,
+    sigma=5.0,
+    interpolation="bilinear",
+    seed=None,
+    data_format=None,
+):
+    if len(images.shape) not in (3, 4):
+        raise ValueError(
+            "Invalid images rank: expected rank 3 (single image) "
+            "or rank 4 (batch of images). Received input with shape: "
+            f"images.shape={images.shape}"
+        )
+
+    images = convert_to_tensor(images)
+    alpha = convert_to_tensor(alpha)
+    sigma = convert_to_tensor(sigma)
+    input_dtype = images.dtype
+    kernel_size = (int(6 * sigma) | 1, int(6 * sigma) | 1)
+
+    need_squeeze = False
+    if len(images.shape) == 3:
+        images = jnp.expand_dims(images, axis=0)
+        need_squeeze = True
+
+    if data_format == "channels_last":
+        batch_size, height, width, channels = images.shape
+        channel_axis = -1
+    else:
+        batch_size, channels, height, width = images.shape
+        channel_axis = 1
+
+    seed = draw_seed(seed)
+    dx = (
+        jax.random.normal(
+            seed, shape=(batch_size, height, width), dtype=input_dtype
+        )
+        * sigma
+    )
+    dy = (
+        jax.random.normal(
+            seed, shape=(batch_size, height, width), dtype=input_dtype
+        )
+        * sigma
+    )
+
+    dx = gaussian_blur(
+        jnp.expand_dims(dx, axis=channel_axis),
+        kernel_size=kernel_size,
+        sigma=(sigma, sigma),
+        data_format=data_format,
+    )
+    dy = gaussian_blur(
+        jnp.expand_dims(dy, axis=channel_axis),
+        kernel_size=kernel_size,
+        sigma=(sigma, sigma),
+        data_format=data_format,
+    )
+
+    dx = jnp.squeeze(dx)
+    dy = jnp.squeeze(dy)
+
+    x, y = jnp.meshgrid(jnp.arange(width), jnp.arange(height))
+    x, y = x[None, :, :], y[None, :, :]
+
+    distorted_x = x + alpha * dx
+    distorted_y = y + alpha * dy
+
+    transformed_images = jnp.zeros_like(images)
+
+    if data_format == "channels_last":
+        for i in range(channels):
+            transformed_images = transformed_images.at[..., i].set(
+                jnp.stack(
+                    [
+                        map_coordinates(
+                            images[b, ..., i],
+                            [distorted_y[b], distorted_x[b]],
+                            order=AFFINE_TRANSFORM_INTERPOLATIONS[
+                                interpolation
+                            ],
+                            fill_mode="reflect",
+                        )
+                        for b in range(batch_size)
+                    ]
+                )
+            )
+    else:
+        for i in range(channels):
+            transformed_images = transformed_images.at[:, i, :, :].set(
+                jnp.stack(
+                    [
+                        map_coordinates(
+                            images[b, i, ...],
+                            [distorted_y[b], distorted_x[b]],
+                            order=AFFINE_TRANSFORM_INTERPOLATIONS[
+                                interpolation
+                            ],
+                            fill_mode="reflect",
+                        )
+                        for b in range(batch_size)
+                    ]
+                )
+            )
+
+    if need_squeeze:
+        transformed_images = jnp.squeeze(transformed_images, axis=0)
+    transformed_images = transformed_images.astype(input_dtype)
+
+    return transformed_images

keras/src/backend/numpy/image.py

@@ -4,6 +4,7 @@
 
 from keras.src import backend
 from keras.src.backend.numpy.core import convert_to_tensor
+from keras.src.random.seed_generator import draw_seed
 from keras.src.utils.module_utils import scipy
 
 RESIZE_INTERPOLATIONS = (
@@ -878,3 +879,111 @@ def _get_gaussian_kernel2d(size, sigma):
         blurred_images = np.squeeze(blurred_images, axis=0)
 
     return blurred_images
+
+
+def elastic_transform(
+    images,
+    alpha=20.0,
+    sigma=5.0,
+    interpolation="bilinear",
+    seed=None,
+    data_format=None,
+):
+    if len(images.shape) not in (3, 4):
+        raise ValueError(
+            "Invalid images rank: expected rank 3 (single image) "
+            "or rank 4 (batch of images). Received input with shape: "
+            f"images.shape={images.shape}"
+        )
+
+    images = convert_to_tensor(images)
+    input_dtype = images.dtype
+
+    alpha = convert_to_tensor(alpha, dtype=input_dtype)
+    sigma = convert_to_tensor(sigma, dtype=input_dtype)
+
+    kernel_size = (int(6 * sigma) | 1, int(6 * sigma) | 1)
+
+    need_squeeze = False
+    if len(images.shape) == 3:
+        images = np.expand_dims(images, axis=0)
+        need_squeeze = True
+
+    if data_format == "channels_last":
+        batch_size, height, width, channels = images.shape
+        channel_axis = -1
+    else:
+        batch_size, channels, height, width = images.shape
+        channel_axis = 1
+
+    seed = draw_seed(seed)
+    rng = np.random.default_rng(seed)
+    dx = (
+        rng.normal(size=(batch_size, height, width), loc=0.0, scale=1.0).astype(
+            input_dtype
+        )
+        * sigma
+    )
+    dy = (
+        rng.normal(size=(batch_size, height, width), loc=0.0, scale=1.0).astype(
+            input_dtype
+        )
+        * sigma
+    )
+
+    dx = gaussian_blur(
+        np.expand_dims(dx, axis=channel_axis),
+        kernel_size=kernel_size,
+        sigma=(sigma, sigma),
+        data_format=data_format,
+    )
+    dy = gaussian_blur(
+        np.expand_dims(dy, axis=channel_axis),
+        kernel_size=kernel_size,
+        sigma=(sigma, sigma),
+        data_format=data_format,
+    )
+
+    dx = np.squeeze(dx)
+    dy = np.squeeze(dy)
+
+    x, y = np.meshgrid(np.arange(width), np.arange(height))
+    x, y = x[None, :, :], y[None, :, :]
+
+    distorted_x = x + alpha * dx
+    distorted_y = y + alpha * dy
+
+    transformed_images = np.zeros_like(images)
+
+    if data_format == "channels_last":
+        for i in range(channels):
+            transformed_images[..., i] = np.stack(
+                [
+                    map_coordinates(
+                        images[b, ..., i],
+                        [distorted_y[b], distorted_x[b]],
+                        order=AFFINE_TRANSFORM_INTERPOLATIONS[interpolation],
+                        fill_mode="reflect",
+                    )
+                    for b in range(batch_size)
+                ]
+            )
+    else:
+        for i in range(channels):
+            transformed_images[:, i, :, :] = np.stack(
+                [
+                    map_coordinates(
+                        images[b, i, ...],
+                        [distorted_y[b], distorted_x[b]],
+                        order=AFFINE_TRANSFORM_INTERPOLATIONS["bilinear"],
+                        fill_mode="reflect",
+                    )
+                    for b in range(batch_size)
+                ]
+            )
+
+    if need_squeeze:
+        transformed_images = np.squeeze(transformed_images, axis=0)
+    transformed_images = transformed_images.astype(input_dtype)
+
+    return transformed_images

keras/src/backend/tensorflow/image.py

@@ -6,6 +6,7 @@
 
 from keras.src import backend
 from keras.src.backend.tensorflow.core import convert_to_tensor
+from keras.src.random.seed_generator import draw_seed
 
 RESIZE_INTERPOLATIONS = (
     "bilinear",
@@ -776,3 +777,125 @@ def _get_gaussian_kernel2d(size, sigma):
         blurred_images = tf.squeeze(blurred_images, axis=0)
 
     return blurred_images
+
+
+def elastic_transform(
+    images,
+    alpha=20.0,
+    sigma=5.0,
+    interpolation="bilinear",
+    seed=None,
+    data_format=None,
+):
+    if len(images.shape) not in (3, 4):
+        raise ValueError(
+            "Invalid images rank: expected rank 3 (single image) "
+            "or rank 4 (batch of images). Received input with shape: "
+            f"images.shape={images.shape}"
+        )
+
+    images = convert_to_tensor(images)
+    input_dtype = images.dtype
+
+    alpha = convert_to_tensor(alpha, dtype=input_dtype)
+    sigma = convert_to_tensor(sigma, dtype=input_dtype)
+    kernel_size = (int(6 * sigma) | 1, int(6 * sigma) | 1)
+
+    need_squeeze = False
+    if len(images.shape) == 3:
+        images = tf.expand_dims(images, axis=0)
+        need_squeeze = True
+
+    if data_format == "channels_last":
+        batch_size, height, width, channels = images.shape
+        channel_axis = -1
+    else:
+        batch_size, channels, height, width = images.shape
+        channel_axis = 1
+
+    seed = draw_seed(seed)
+    dx = tf.random.stateless_normal(
+        shape=(batch_size, height, width),
+        mean=0.0,
+        stddev=1.0,
+        dtype=input_dtype,
+        seed=seed,
+    )
+    dy = tf.random.stateless_normal(
+        shape=(batch_size, height, width),
+        mean=0.0,
+        stddev=1.0,
+        dtype=input_dtype,
+        seed=seed,
+    )
+
+    dx = gaussian_blur(
+        tf.expand_dims(dx, axis=channel_axis),
+        kernel_size=kernel_size,
+        sigma=(sigma, sigma),
+        data_format=data_format,
+    )
+    dy = gaussian_blur(
+        tf.expand_dims(dy, axis=channel_axis),
+        kernel_size=kernel_size,
+        sigma=(sigma, sigma),
+        data_format=data_format,
+    )
+
+    dx = tf.squeeze(dx, axis=channel_axis)
+    dy = tf.squeeze(dy, axis=channel_axis)
+
+    x, y = tf.meshgrid(
+        tf.range(width, dtype=input_dtype),
+        tf.range(height, dtype=input_dtype),
+        indexing="xy",
+    )
+    x = tf.expand_dims(x, axis=0)
+    y = tf.expand_dims(y, axis=0)
+
+    distorted_x = x + alpha * dx
+    distorted_y = y + alpha * dy
+
+    channel_outputs = []
+    if data_format == "channels_last":
+        for i in range(channels):
+            channel_transformed = tf.stack(
+                [
+                    map_coordinates(
+                        images[b, ..., i],
+                        [distorted_y[b], distorted_x[b]],
+                        order=AFFINE_TRANSFORM_INTERPOLATIONS.index(
+                            interpolation
+                        ),
+                        fill_mode="reflect",
+                    )
+                    for b in range(batch_size)
+                ],
+                axis=0,
+            )
+            channel_outputs.append(channel_transformed)
+        transformed_images = tf.stack(channel_outputs, axis=-1)
+    else:
+        for i in range(channels):
+            channel_transformed = tf.stack(
+                [
+                    map_coordinates(
+                        images[b, i, ...],
+                        [distorted_y[b], distorted_x[b]],
+                        order=AFFINE_TRANSFORM_INTERPOLATIONS.index(
+                            interpolation
+                        ),
+                        fill_mode="reflect",
+                    )
+                    for b in range(batch_size)
+                ],
+                axis=0,
+            )
+            channel_outputs.append(channel_transformed)
+        transformed_images = tf.stack(channel_outputs, axis=1)
+
+    if need_squeeze:
+        transformed_images = tf.squeeze(transformed_images, axis=0)
+    transformed_images = tf.cast(transformed_images, input_dtype)
+
+    return transformed_images