Support for extracting volume patches (#21759)

* support for volume patching * correct `call` * fix pydoc * fix `test_extract_volume_patches_basic` casting residual * fix `test_extract_volume_patches_same_padding` casting * fix `test_extract_volume_patches_overlapping` casting * add extra testing * Update keras/src/ops/image.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * fix dimension orderiing * Update keras/src/ops/image.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * fix dimensional ordering + add test * fix docstring * docstring corrections * add validation checks for operation * comment * set default data format in tests, add a few parametrized tests to validate channels_first * add back test for channels first/last in dilation * fix dilation test * Update keras/src/ops/image.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * delete redundant prop set * Update keras/src/ops/image_test.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update keras/src/ops/image_test.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * rename per francois feedback * fix pydoc to reflect method * fix name of operation + test --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: danielenricocahall

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

@@ -8,6 +8,7 @@
 from keras.src.ops.image import crop_images as crop_images
 from keras.src.ops.image import elastic_transform as elastic_transform
 from keras.src.ops.image import extract_patches as extract_patches
+from keras.src.ops.image import extract_patches_3d as extract_patches_3d
 from keras.src.ops.image import gaussian_blur as gaussian_blur
 from keras.src.ops.image import hsv_to_rgb as hsv_to_rgb
 from keras.src.ops.image import map_coordinates as map_coordinates

keras/api/ops/image/__init__.py

@@ -8,6 +8,7 @@
 from keras.src.ops.image import crop_images as crop_images
 from keras.src.ops.image import elastic_transform as elastic_transform
 from keras.src.ops.image import extract_patches as extract_patches
+from keras.src.ops.image import extract_patches_3d as extract_patches_3d
 from keras.src.ops.image import gaussian_blur as gaussian_blur
 from keras.src.ops.image import hsv_to_rgb as hsv_to_rgb
 from keras.src.ops.image import map_coordinates as map_coordinates

keras/src/ops/image.py

@@ -712,6 +712,187 @@ def _extract_patches(
  return patches
 
 
+class ExtractPatches3D(Operation):
+ def __init__(
+ self,
+ size,
+ strides=None,
+ dilation_rate=1,
+ padding="valid",
+ data_format=None,
+ *,
+ name=None,
+ ):
+ super().__init__(name=name)
+ if isinstance(size, int):
+ size = (size, size, size)
+ elif len(size) != 3:
+ raise TypeError(
+ "Invalid `size` argument. Expected an "
+ f"int or a tuple of length 3. Received: size={size}"
+ )
+ self.size = size
+ if strides is not None:
+ if isinstance(strides, int):
+ strides = (strides, strides, strides)
+ elif len(strides) != 3:
+ raise ValueError(f"Invalid `strides` argument. Got: {strides}")
+ else:
+ strides = size
+ self.strides = strides
+ self.dilation_rate = dilation_rate
+ self.padding = padding
+ self.data_format = backend.standardize_data_format(data_format)
+
+ def call(self, volumes):
+ return _extract_patches_3d(
+ volumes,
+ self.size,
+ self.strides,
+ self.dilation_rate,
+ self.padding,
+ self.data_format,
+ )
+
+ def compute_output_spec(self, volumes):
+ volumes_shape = list(volumes.shape)
+ original_ndim = len(volumes_shape)
+ strides = self.strides
+ if self.data_format == "channels_last":
+ channels_in = volumes_shape[-1]
+ else:
+ channels_in = volumes_shape[-4]
+ if original_ndim == 4:
+ volumes_shape = [1] + volumes_shape
+ filters = self.size[0] * self.size[1] * self.size[2] * channels_in
+ kernel_size = (self.size[0], self.size[1], self.size[2])
+ out_shape = compute_conv_output_shape(
+ volumes_shape,
+ filters,
+ kernel_size,
+ strides=strides,
+ padding=self.padding,
+ data_format=self.data_format,
+ dilation_rate=self.dilation_rate,
+ )
+ if original_ndim == 4:
+ out_shape = out_shape[1:]
+ return KerasTensor(shape=out_shape, dtype=volumes.dtype)
+
+
+def _extract_patches_3d(
+ volumes,
+ size,
+ strides=None,
+ dilation_rate=1,
+ padding="valid",
+ data_format=None,
+):
+ if isinstance(size, int):
+ patch_d = patch_h = patch_w = size
+ elif len(size) == 3:
+ patch_d, patch_h, patch_w = size
+ else:
+ raise TypeError(
+ "Invalid `size` argument. Expected an "
+ f"int or a tuple of length 3. Received: size={size}"
+ )
+ if strides is None:
+ strides = size
+ if isinstance(strides, int):
+ strides = (strides, strides, strides)
+ if len(strides) != 3:
+ raise ValueError(f"Invalid `strides` argument. Got: {strides}")
+ data_format = backend.standardize_data_format(data_format)
+ if data_format == "channels_last":
+ channels_in = volumes.shape[-1]
+ elif data_format == "channels_first":
+ channels_in = volumes.shape[-4]
+ out_dim = patch_d * patch_w * patch_h * channels_in
+ kernel = backend.numpy.eye(out_dim, dtype=volumes.dtype)
+ kernel = backend.numpy.reshape(
+ kernel, (patch_d, patch_h, patch_w, channels_in, out_dim)
+ )
+ _unbatched = False
+ if len(volumes.shape) == 4:
+ _unbatched = True
+ volumes = backend.numpy.expand_dims(volumes, axis=0)
+ patches = backend.nn.conv(
+ inputs=volumes,
+ kernel=kernel,
+ strides=strides,
+ padding=padding,
+ data_format=data_format,
+ dilation_rate=dilation_rate,
+ )
+ if _unbatched:
+ patches = backend.numpy.squeeze(patches, axis=0)
+ return patches
+
+
+@keras_export("keras.ops.image.extract_patches_3d")
+def extract_patches_3d(
+ volumes,
+ size,
+ strides=None,
+ dilation_rate=1,
+ padding="valid",
+ data_format=None,
+):
+ """Extracts patches from the volume(s).
+
+ Args:
+ volumes: Input volume or batch of volumes. Must be 4D or 5D.
+ size: Patch size int or tuple (patch_depth, patch_height, patch_width)
+ strides: strides along depth, height, and width. If not specified, or
+ if `None`, it defaults to the same value as `size`.
+ dilation_rate: This is the input stride, specifying how far two
+ consecutive patch samples are in the input. Note that using
+ `dilation_rate > 1` is not supported in conjunction with
+ `strides > 1` on the TensorFlow backend.
+ padding: The type of padding algorithm to use: `"same"` or `"valid"`.
+ 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, depth, height, width, channels)`, while `"channels_first"`
+ corresponds to inputs with shape
+ `(batch, channels, depth, height, width)`. If not specified,
+ the value will default to `keras.config.image_data_format()`.
+
+ Returns:
+ Extracted patches 4D (if not batched) or 5D (if batched)
+
+ Examples:
+
+ >>> import numpy as np
+ >>> import keras
+ >>> # Batched case
+ >>> volumes = np.random.random(
+ ... (2, 10, 10, 10, 3)
+ ... ).astype("float32") # batch of 2 volumes
+ >>> patches = keras.ops.image.extract_patches_3d(volumes, (3, 3, 3))
+ >>> patches.shape
+ (2, 3, 3, 3, 81)
+ >>> # Unbatched case
+ >>> volume = np.random.random((10, 10, 10, 3)).astype("float32") # 1 volume
+ >>> patches = keras.ops.image.extract_patches_3d(volume, (3, 3, 3))
+ >>> patches.shape
+ (3, 3, 3, 81)
+ """
+ if any_symbolic_tensors((volumes,)):
+ return ExtractPatches3D(
+ size=size,
+ strides=strides,
+ dilation_rate=dilation_rate,
+ padding=padding,
+ data_format=data_format,
+ ).symbolic_call(volumes)
+
+ return _extract_patches_3d(
+ volumes, size, strides, dilation_rate, padding, data_format=data_format
+ )
+
+
 class MapCoordinates(Operation):
  def __init__(self, order, fill_mode="constant", fill_value=0, *, name=None):
  super().__init__(name=name)