Add sobel_edges to keras.ops.image

keras/api/_tf_keras/keras/ops/image/__init__.py

@@ -18,3 +18,4 @@
 from keras.src.ops.image import rgb_to_grayscale as rgb_to_grayscale
 from keras.src.ops.image import rgb_to_hsv as rgb_to_hsv
 from keras.src.ops.image import scale_and_translate as scale_and_translate
+from keras.src.ops.image import sobel_edges as sobel_edges

keras/api/ops/image/__init__.py

@@ -18,3 +18,4 @@
 from keras.src.ops.image import rgb_to_grayscale as rgb_to_grayscale
 from keras.src.ops.image import rgb_to_hsv as rgb_to_hsv
 from keras.src.ops.image import scale_and_translate as scale_and_translate
+from keras.src.ops.image import sobel_edges as sobel_edges

keras/src/backend/jax/image.py

@@ -895,3 +895,54 @@ def scale_and_translate(
         method,
         antialias,
     )
+
+
+def sobel_edges(images, data_format=None):
+    data_format = backend.standardize_data_format(data_format)
+    input_dtype = backend.standardize_dtype(images.dtype)
+    compute_dtype = backend.result_type(input_dtype, "float32")
+    images = images.astype(compute_dtype)
+
+    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}"
+        )
+
+    need_squeeze = False
+    if len(images.shape) == 3:
+        images = jnp.expand_dims(images, axis=0)
+        need_squeeze = True
+
+    if data_format == "channels_first":
+        images = jnp.transpose(images, (0, 2, 3, 1))
+
+    images = jnp.pad(images, ((0, 0), (1, 1), (1, 1), (0, 0)), mode="edge")
+
+    dy = (
+        -1 * images[:, :-2, :-2, :]
+        + -2 * images[:, :-2, 1:-1, :]
+        + -1 * images[:, :-2, 2:, :]
+        + 1 * images[:, 2:, :-2, :]
+        + 2 * images[:, 2:, 1:-1, :]
+        + 1 * images[:, 2:, 2:, :]
+    )
+    dx = (
+        -1 * images[:, :-2, :-2, :]
+        + -2 * images[:, 1:-1, :-2, :]
+        + -1 * images[:, 2:, :-2, :]
+        + 1 * images[:, :-2, 2:, :]
+        + 2 * images[:, 1:-1, 2:, :]
+        + 1 * images[:, 2:, 2:, :]
+    )
+
+    result = jnp.stack([dy, dx], axis=-1)
+
+    if data_format == "channels_first":
+        result = jnp.transpose(result, (0, 3, 1, 2, 4))
+
+    if need_squeeze:
+        result = jnp.squeeze(result, axis=0)
+
+    return result

keras/src/backend/numpy/image.py

@@ -1203,3 +1203,60 @@ def scale_and_translate(
         kernel,
         antialias,
     )
+
+
+def sobel_edges(images, data_format=None):
+    images = convert_to_tensor(images)
+    data_format = backend.standardize_data_format(data_format)
+    input_dtype = backend.standardize_dtype(images.dtype)
+    compute_dtype = backend.result_type(input_dtype, "float32")
+    images = images.astype(compute_dtype)
+
+    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}"
+        )
+
+    need_squeeze = False
+    if len(images.shape) == 3:
+        images = np.expand_dims(images, axis=0)
+        need_squeeze = True
+
+    if data_format == "channels_first":
+        images = np.transpose(images, (0, 2, 3, 1))
+
+    # Pad with replicated edges
+    images = np.pad(images, ((0, 0), (1, 1), (1, 1), (0, 0)), mode="edge")
+
+    # Sobel kernels
+    # Vertical (dy)
+    dy = (
+        -1 * images[:, :-2, :-2, :]
+        + -2 * images[:, :-2, 1:-1, :]
+        + -1 * images[:, :-2, 2:, :]
+        + 1 * images[:, 2:, :-2, :]
+        + 2 * images[:, 2:, 1:-1, :]
+        + 1 * images[:, 2:, 2:, :]
+    )
+    # Horizontal (dx)
+    dx = (
+        -1 * images[:, :-2, :-2, :]
+        + -2 * images[:, 1:-1, :-2, :]
+        + -1 * images[:, 2:, :-2, :]
+        + 1 * images[:, :-2, 2:, :]
+        + 2 * images[:, 1:-1, 2:, :]
+        + 1 * images[:, 2:, 2:, :]
+    )
+
+    result = np.stack([dy, dx], axis=-1)
+
+    if data_format == "channels_first":
+        # (b, h, w, c, 2) -> (b, c, h, w, 2)
+        result = np.transpose(result, (0, 3, 1, 2, 4))
+
+    if need_squeeze:
+        result = np.squeeze(result, axis=0)
+
+    return result

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -312,6 +312,7 @@ ImageDatasetFromDirectoryTest::test_sample_count_grain
 ImageOpsBehaviorTests::test_affine_transform
 ImageOpsBehaviorTests::test_elastic_transform
 ImageOpsBehaviorTests::test_gaussian_blur
+ImageOpsBehaviorTests::test_sobel_edges
 ImageOpsBehaviorTests::test_map_coordinates
 ImageOpsBehaviorTests::test_perspective_transform
 ImageOpsBehaviorTests::test_resize
@@ -320,6 +321,7 @@ ImageOpsCorrectnessTest::test_crop_images
 ImageOpsCorrectnessTest::test_elastic_transform
 ImageOpsCorrectnessTest::test_extract_patches
 ImageOpsCorrectnessTest::test_gaussian_blur
+ImageOpsCorrectnessTest::test_sobel_edges
 ImageOpsCorrectnessTest::test_map_coordinates
 ImageOpsCorrectnessTest::test_pad_images
 ImageOpsCorrectnessTest::test_perspective_transform
@@ -328,6 +330,7 @@ ImageOpsCorrectnessTest::test_scale_and_translate
 ImageOpsDtypeTest::test_affine_transform
 ImageOpsDtypeTest::test_elastic_transform
 ImageOpsDtypeTest::test_gaussian_blur
+ImageOpsDtypeTest::test_sobel_edges
 ImageOpsDtypeTest::test_map_coordinates
 ImageOpsDtypeTest::test_perspective_transform
 ImageOpsDtypeTest::test_resize

keras/src/backend/openvino/image.py

@@ -291,3 +291,9 @@ def scale_and_translate(
     raise NotImplementedError(
         "`scale_and_translate` is not supported with openvino backend"
     )
+
+
+def sobel_edges(images, data_format=None):
+    raise NotImplementedError(
+        "`sobel_edges` is not supported with openvino backend"
+    )

keras/src/backend/tensorflow/image.py

@@ -1074,3 +1074,37 @@ def scale_and_translate(
         kernel,
         antialias,
     )
+
+
+def sobel_edges(images, data_format=None):
+    data_format = backend.standardize_data_format(data_format)
+    input_dtype = backend.standardize_dtype(images.dtype)
+    compute_dtype = backend.result_type(input_dtype, "float32")
+    images = tf.cast(images, compute_dtype)
+
+    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}"
+        )
+
+    need_squeeze = False
+    if len(images.shape) == 3:
+        images = tf.expand_dims(images, axis=0)
+        need_squeeze = True
+
+    if data_format == "channels_first":
+        images = tf.transpose(images, (0, 2, 3, 1))
+
+    # tf.image.sobel_edges expects channels_last (b, h, w, c)
+    # and returns (b, h, w, c, 2) where last dim is [dy, dx]
+    result = tf.image.sobel_edges(images)
+
+    if data_format == "channels_first":
+        result = tf.transpose(result, (0, 3, 1, 2, 4))
+
+    if need_squeeze:
+        result = tf.squeeze(result, axis=0)
+
+    return result

keras/src/backend/torch/image.py

@@ -1190,3 +1190,55 @@ def scale_and_translate(
         kernel,
         antialias,
     )
+
+
+def sobel_edges(images, data_format=None):
+    data_format = backend.standardize_data_format(data_format)
+    input_dtype = backend.standardize_dtype(images.dtype)
+    compute_dtype = backend.result_type(input_dtype, "float32")
+    images = cast(images, compute_dtype)
+
+    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}"
+        )
+
+    need_squeeze = False
+    if len(images.shape) == 3:
+        images = images.unsqueeze(0)
+        need_squeeze = True
+
+    if data_format == "channels_first":
+        images = images.permute(0, 2, 3, 1)
+
+    images = F.pad(images.permute(0, 3, 1, 2), (1, 1, 1, 1), mode="replicate")
+    images = images.permute(0, 2, 3, 1)
+
+    dy = (
+        -1 * images[:, :-2, :-2, :]
+        + -2 * images[:, :-2, 1:-1, :]
+        + -1 * images[:, :-2, 2:, :]
+        + 1 * images[:, 2:, :-2, :]
+        + 2 * images[:, 2:, 1:-1, :]
+        + 1 * images[:, 2:, 2:, :]
+    )
+    dx = (
+        -1 * images[:, :-2, :-2, :]
+        + -2 * images[:, 1:-1, :-2, :]
+        + -1 * images[:, 2:, :-2, :]
+        + 1 * images[:, :-2, 2:, :]
+        + 2 * images[:, 1:-1, 2:, :]
+        + 1 * images[:, 2:, 2:, :]
+    )
+
+    result = torch.stack([dy, dx], dim=-1)
+
+    if data_format == "channels_first":
+        result = result.permute(0, 3, 1, 2, 4)
+
+    if need_squeeze:
+        result = result.squeeze(0)
+
+    return result

keras/src/ops/image.py

@@ -1910,3 +1910,65 @@ def scale_and_translate(
         method,
         antialias,
     )
+
+
+class SobelEdges(Operation):
+    def __init__(self, data_format=None, *, name=None):
+        super().__init__(name=name)
+        self.data_format = backend.standardize_data_format(data_format)
+
+    def call(self, images):
+        return backend.image.sobel_edges(images, data_format=self.data_format)
+
+    def compute_output_spec(self, images):
+        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}"
+            )
+        return KerasTensor(
+            images.shape + (2,), dtype=backend.result_type(images.dtype, float)
+        )
+
+
+@keras_export("keras.ops.image.sobel_edges")
+def sobel_edges(images, data_format=None):
+    """Computes Sobel edge maps of the given images.
+
+    The Sobel operator highlights regions of high spatial frequency
+    (edges) by approximating the image gradient. For each pixel the
+    vertical (dy) and horizontal (dx) gradients are returned as the
+    last dimension of the output tensor.
+
+    Args:
+        images: Input image or batch of images. Must be 3D or 4D.
+        data_format: A string specifying the data format of the input
+            tensor. It can be either `"channels_last"` or
+            `"channels_first"`. `"channels_last"` corresponds to inputs
+            with shape `(batch, height, width, channels)`, while
+            `"channels_first"` corresponds to inputs with shape
+            `(batch, channels, height, width)`. If not specified, the
+            value will default to `keras.config.image_data_format`.
+
+    Returns:
+        A tensor with the same dtype (promoted to float) and one
+        additional trailing dimension of size 2 containing `[dy, dx]`.
+        For a 4D input of shape `(batch, height, width, channels)` the
+        output shape is `(batch, height, width, channels, 2)`.
+
+    Examples:
+
+    >>> x = np.random.random((2, 64, 80, 3))
+    >>> y = keras.ops.image.sobel_edges(x)
+    >>> y.shape
+    (2, 64, 80, 3, 2)
+
+    >>> x = np.random.random((64, 80, 3))
+    >>> y = keras.ops.image.sobel_edges(x)
+    >>> y.shape
+    (64, 80, 3, 2)
+    """
+    if any_symbolic_tensors((images,)):
+        return SobelEdges(data_format=data_format).symbolic_call(images)
+    return backend.image.sobel_edges(images, data_format=data_format)

keras/src/ops/image_test.py

@@ -221,6 +221,18 @@ def test_gaussian_blur(self):
         out = kimage.gaussian_blur(x)
         self.assertEqual(out.shape, (None, 3, 20, 20))
 
+    def test_sobel_edges(self):
+        # Test channels_last
+        x = KerasTensor([None, 20, 20, 3])
+        out = kimage.sobel_edges(x)
+        self.assertEqual(out.shape, (None, 20, 20, 3, 2))
+
+        # Test channels_first
+        backend.set_image_data_format("channels_first")
+        x = KerasTensor([None, 3, 20, 20])
+        out = kimage.sobel_edges(x)
+        self.assertEqual(out.shape, (None, 3, 20, 20, 2))
+
     def test_elastic_transform(self):
         # Test channels_last
         x = KerasTensor([None, 20, 20, 3])
@@ -1903,6 +1915,34 @@ def test_gaussian_blur(self):
         self.assertEqual(tuple(out.shape), tuple(ref_out.shape))
         self.assertAllClose(ref_out, out, atol=1e-2, rtol=1e-2)
 
+    def test_sobel_edges(self):
+        # Test channels_last
+        backend.set_image_data_format("channels_last")
+        np.random.seed(42)
+        x = np.random.uniform(size=(2, 10, 10, 3)).astype("float32")
+        out = kimage.sobel_edges(x, data_format="channels_last")
+        self.assertEqual(out.shape, (2, 10, 10, 3, 2))
+
+        # Test single image
+        x_single = np.random.uniform(size=(10, 10, 3)).astype("float32")
+        out_single = kimage.sobel_edges(x_single, data_format="channels_last")
+        self.assertEqual(out_single.shape, (10, 10, 3, 2))
+
+        # Test channels_first
+        backend.set_image_data_format("channels_first")
+        x_cf = np.random.uniform(size=(2, 3, 10, 10)).astype("float32")
+        out_cf = kimage.sobel_edges(x_cf, data_format="channels_first")
+        self.assertEqual(out_cf.shape, (2, 3, 10, 10, 2))
+
+        # Test edge detection on known pattern: vertical edge
+        backend.set_image_data_format("channels_last")
+        img = np.zeros((1, 5, 5, 1), dtype="float32")
+        img[0, :, 3:, 0] = 1.0
+        result = kimage.sobel_edges(img, data_format="channels_last")
+        dx = result[0, :, :, 0, 1]
+        # The horizontal gradient should be non-zero near the edge
+        self.assertTrue(np.any(np.abs(np.array(dx)) > 0))
+
     def test_gaussian_blur_even_kernel_size(self):
         """Test gaussian_blur with even kernel sizes"""
         # This test is specific to the numpy backend fix