feat(layers): Add TPU-optimized 3D Elastic Deformation layer

Author: ashpakshaikh26732

keras_hub/api/layers/__init__.py

@@ -147,3 +147,7 @@
 from keras_hub.src.models.xception.xception_image_converter import (
     XceptionImageConverter as XceptionImageConverter,
 )
+
+from keras_hub.src.layers.preprocessing.random_elastic_deformation_3d import (
+    RandomElasticDeformation3D,
+)

keras_hub/src/layers/preprocessing/random_elastic_deformation_3d.py

@@ -0,0 +1,127 @@
+import tensorflow as tf
+
+class RandomElasticDeformation3D(tf.keras.layers.Layer):
+    """
+    A high-performance 3D elastic deformation layer optimized for TPUs and GPUs.
+    ... (docstring is the same) ...
+    """
+    def __init__(self,
+                 grid_size=(4, 4, 4),
+                 alpha=35.0,
+                 sigma=2.5,
+                 data_format="DHWC",
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.grid_size = grid_size
+        self.alpha = tf.constant(alpha, dtype=tf.bfloat16)
+        self.sigma = tf.constant(sigma, dtype=tf.bfloat16)
+        if data_format not in ["DHWC", "HWDC"]:
+            raise ValueError("`data_format` must be one of 'DHWC' or 'HWDC'")
+        self.data_format = data_format
+
+    def _separable_gaussian_filter_3d(self, tensor, sigma):
+
+        kernel_size = tf.cast(2 * tf.round(3 * sigma) + 1, dtype=tf.int32)
+        ax = tf.range(-tf.cast(kernel_size // 2, tf.bfloat16) + 1.0,
+                      tf.cast(kernel_size // 2, tf.bfloat16) + 1.0)
+        kernel_1d = tf.exp(-(ax**2) / (2.0 * self.sigma**2))
+        kernel_1d = kernel_1d / tf.reduce_sum(kernel_1d)
+        filter_d = tf.cast(tf.reshape(kernel_1d, [-1, 1, 1, 1, 1]), dtype=tensor.dtype)
+        filter_h = tf.cast(tf.reshape(kernel_1d, [1, -1, 1, 1, 1]), dtype=tensor.dtype)
+        filter_w = tf.cast(tf.reshape(kernel_1d, [1, 1, -1, 1, 1]), dtype=tensor.dtype)
+        tensor = tf.nn.convolution(tensor, filter_d, strides=1, padding='SAME')
+        tensor = tf.nn.convolution(tensor, filter_h, strides=1, padding='SAME')
+        tensor = tf.nn.convolution(tensor, filter_w, strides=1, padding='SAME')
+        return tensor
+
+    def call(self, inputs):
+        image_volume, label_volume = inputs
+        original_image_dtype = image_volume.dtype
+
+        was_batched = True
+        if image_volume.shape.rank == 4:
+            was_batched = False
+            image_volume = tf.expand_dims(image_volume, axis=0)
+            label_volume = tf.expand_dims(label_volume, axis=0)
+
+        if self.data_format == "HWDC":
+            image_volume = tf.transpose(image_volume, perm=[0, 3, 1, 2, 4])
+            label_volume = tf.transpose(label_volume, perm=[0, 3, 1, 2, 4])
+
+        image_volume = tf.cast(image_volume, dtype=tf.bfloat16)
+        input_shape = tf.shape(image_volume)
+        B, D, H, W = input_shape[0], input_shape[1], input_shape[2], input_shape[3]
+
+        coarse_flow = tf.random.uniform(
+            shape=(B, self.grid_size[0], self.grid_size[1], self.grid_size[2], 3),
+            minval=-1, maxval=1, dtype=tf.bfloat16)
+
+        flow = tf.reshape(coarse_flow, [B * self.grid_size[0], self.grid_size[1], self.grid_size[2], 3])
+        flow = tf.image.resize(flow, size=[H, W], method='bicubic')
+        flow = tf.reshape(flow, [B, self.grid_size[0], H, W, 3])
+        flow = tf.transpose(flow, perm=[0, 2, 3, 1, 4])
+        flow = tf.reshape(flow, [B * H * W, self.grid_size[0], 3])
+        flow = tf.image.resize(tf.expand_dims(flow, axis=1), size=[1, D], method='bicubic')
+        flow = tf.squeeze(flow, axis=1)
+        flow = tf.reshape(flow, [B, H, W, D, 3])
+        flow = tf.transpose(flow, perm=[0, 3, 1, 2, 4])
+
+
+        flow = tf.cast(flow, dtype=tf.bfloat16)
+        
+        flow_components = tf.unstack(flow, axis=-1)
+        smoothed_components = []
+        for component in flow_components:
+            smoothed_component = self._separable_gaussian_filter_3d(
+                component[..., tf.newaxis], self.sigma
+            )
+            smoothed_components.append(smoothed_component[..., 0])
+        smoothed_flow = tf.stack(smoothed_components, axis=-1)
+        
+
+        flow = smoothed_flow * self.alpha
+
+        grid_d, grid_h, grid_w = tf.meshgrid(
+            tf.range(D, dtype=tf.bfloat16),
+            tf.range(H, dtype=tf.bfloat16),
+            tf.range(W, dtype=tf.bfloat16),
+            indexing='ij'
+        )
+        grid = tf.stack([grid_d, grid_h, grid_w], axis=-1)
+        
+
+        warp_grid = tf.expand_dims(grid, 0) + flow
+        
+        warp_grid_floor = tf.floor(warp_grid)
+        t = warp_grid - warp_grid_floor
+
+        d0 = tf.cast(warp_grid_floor[..., 0], tf.int32); h0 = tf.cast(warp_grid_floor[..., 1], tf.int32); w0 = tf.cast(warp_grid_floor[..., 2], tf.int32)
+        d1 = tf.clip_by_value(d0 + 1, 0, D - 1); h1 = tf.clip_by_value(h0 + 1, 0, H - 1); w1 = tf.clip_by_value(w0 + 1, 0, W - 1)
+        d0 = tf.clip_by_value(d0, 0, D - 1); h0 = tf.clip_by_value(h0, 0, H - 1); w0 = tf.clip_by_value(w0, 0, W - 1)
+
+        c000 = tf.gather_nd(image_volume, tf.stack([d0, h0, w0], axis=-1), batch_dims=1); c001 = tf.gather_nd(image_volume, tf.stack([d0, h0, w1], axis=-1), batch_dims=1)
+        c010 = tf.gather_nd(image_volume, tf.stack([d0, h1, w0], axis=-1), batch_dims=1); c011 = tf.gather_nd(image_volume, tf.stack([d0, h1, w1], axis=-1), batch_dims=1)
+        c100 = tf.gather_nd(image_volume, tf.stack([d1, h0, w0], axis=-1), batch_dims=1); c101 = tf.gather_nd(image_volume, tf.stack([d1, h0, w1], axis=-1), batch_dims=1)
+        c110 = tf.gather_nd(image_volume, tf.stack([d1, h1, w0], axis=-1), batch_dims=1); c111 = tf.gather_nd(image_volume, tf.stack([d1, h1, w1], axis=-1), batch_dims=1)
+
+        td, th, tw = t[..., 0:1], t[..., 1:2], t[..., 2:3]
+        c00 = c000*(1-tw) + c001*tw; c01 = c010*(1-tw) + c011*tw; c10 = c100*(1-tw) + c101*tw; c11 = c110*(1-tw) + c111*tw
+        c0 = c00*(1-th) + c01*th; c1 = c10*(1-th) + c11*th
+        deformed_image = c0*(1-td) + c1*td
+        deformed_image = tf.cast(deformed_image, original_image_dtype)
+
+        nearest_indices_float = tf.round(warp_grid)
+        nearest_d = tf.clip_by_value(tf.cast(nearest_indices_float[..., 0], tf.int32), 0, D - 1)
+        nearest_h = tf.clip_by_value(tf.cast(nearest_indices_float[..., 1], tf.int32), 0, H - 1)
+        nearest_w = tf.clip_by_value(tf.cast(nearest_indices_float[..., 2], tf.int32), 0, W - 1)
+        deformed_label = tf.gather_nd(label_volume, tf.stack([nearest_d, nearest_h, nearest_w], axis=-1), batch_dims=1)
+
+        if self.data_format == "HWDC":
+            deformed_image = tf.transpose(deformed_image, perm=[0, 2, 3, 1, 4])
+            deformed_label = tf.transpose(deformed_label, perm=[0, 2, 3, 1, 4])
+
+        if not was_batched:
+            deformed_image = tf.squeeze(deformed_image, axis=0)
+            deformed_label = tf.squeeze(deformed_label, axis=0)
+
+        return deformed_image, deformed_label

keras_hub/src/layers/preprocessing/random_elastic_deformation_3d_test.py

@@ -0,0 +1,70 @@
+import tensorflow as tf
+from tensorflow import keras
+from keras_hub.src.layers.preprocessing.random_elastic_deformation_3d import RandomElasticDeformation3D
+
+class RandomElasticDeformation3DTest(tf.test.TestCase):
+
+    def test_output_shape_is_same_as_input_dhwc(self):
+        input_image = tf.random.uniform(shape=(2, 32, 64, 64, 3), dtype=tf.float32)
+        input_label = tf.random.uniform(shape=(2, 32, 64, 64, 1), maxval=4, dtype=tf.int32)
+        layer = RandomElasticDeformation3D(data_format="DHWC")
+        output_image, output_label = layer((input_image, tf.cast(input_label, tf.float32)))
+        self.assertAllEqual(tf.shape(input_image), tf.shape(output_image))
+        self.assertAllEqual(tf.shape(input_label), tf.shape(output_label))
+
+    def test_output_shape_is_same_as_input_hwdc(self):
+        input_image = tf.random.uniform(shape=(2, 64, 64, 32, 3), dtype=tf.float32)
+        input_label = tf.random.uniform(shape=(2, 64, 64, 32, 1), maxval=4, dtype=tf.int32)
+        layer = RandomElasticDeformation3D(data_format="HWDC")
+        output_image, output_label = layer((input_image, tf.cast(input_label, tf.float32)))
+        self.assertAllEqual(tf.shape(input_image), tf.shape(output_image))
+        self.assertAllEqual(tf.shape(input_label), tf.shape(output_label))
+
+    def test_unbatched_input(self):
+        input_image = tf.random.uniform(shape=(32, 64, 64, 3), dtype=tf.float32)
+        input_label = tf.random.uniform(shape=(32, 64, 64, 1), maxval=4, dtype=tf.int32)
+        layer = RandomElasticDeformation3D(data_format="DHWC")
+        output_image, output_label = layer((input_image, tf.cast(input_label, tf.float32)))
+        self.assertAllEqual(tf.shape(input_image), tf.shape(output_image))
+        self.assertEqual(tf.rank(output_image), 4)
+
+    def test_dtype_preservation(self):
+        input_image = tf.random.uniform(shape=(2, 16, 16, 16, 3), dtype=tf.float32)
+        input_label = tf.random.uniform(shape=(2, 16, 16, 16, 1), maxval=4, dtype=tf.int32)
+        layer = RandomElasticDeformation3D()
+        output_image, output_label = layer((input_image, tf.cast(input_label, tf.float32)))
+        self.assertEqual(output_image.dtype, tf.float32)
+        self.assertEqual(output_label.dtype, tf.float32)
+
+    def test_label_values_are_preserved(self):
+        input_image = tf.zeros(shape=(1, 16, 16, 16, 1), dtype=tf.float32)
+        label_arange = tf.experimental.numpy.arange(16**3)
+        input_label = tf.reshape(label_arange, (1, 16, 16, 16, 1))
+        input_label = tf.cast(input_label, dtype=tf.float32) % 4
+        
+        layer = RandomElasticDeformation3D(alpha=80.0, sigma=8.0)
+        _, output_label = layer((input_image, input_label))
+        
+        unique_values_tensor = tf.unique(tf.reshape(output_label, [-1]))[0]
+        
+
+        expected_values = [0., 1., 2., 3.]
+        actual_values = unique_values_tensor.numpy().tolist()
+        self.assertContainsSubset(expected_values, actual_values)
+        
+    def test_config_serialization(self):
+        layer = RandomElasticDeformation3D(
+            grid_size=(3, 3, 3),
+            alpha=50.0,
+            sigma=5.0,
+            data_format="HWDC"
+        )
+        config = layer.get_config()
+        new_layer = RandomElasticDeformation3D.from_config(config)
+        self.assertEqual(new_layer.grid_size, (3, 3, 3))
+        self.assertAllClose(new_layer.alpha, tf.constant(50.0, dtype=tf.bfloat16))
+        self.assertAllClose(new_layer.sigma, tf.constant(5.0, dtype=tf.bfloat16))
+        self.assertEqual(new_layer.data_format, "HWDC")
+
+if __name__ == "__main__":
+    tf.test.main()