feat(layers): Add 3D elastic deformation layer

Author: ashpakshaikh26732

keras_hub/api/layers/__init__.py

@@ -147,7 +147,3 @@
 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

@@ -8,6 +8,7 @@ class RandomElasticDeformation3D(layers.Layer):
     """
     A high-performance 3D elastic deformation layer optimized for TPUs.
     """
+
     def __init__(self,
                  grid_size=(4, 4, 4),
                  alpha=35.0,
@@ -20,7 +21,11 @@ def __init__(self,
         self.sigma = sigma
         self.data_format = data_format
         if data_format not in ["channels_last", "channels_first"]:
-            raise ValueError(f"`data_format` must be one of 'channels_last' or 'channels_first'. Received: {data_format}")
+            message = (
+                "`data_format` must be one of 'channels_last' or "
+                f"'channels_first'. Received: {self.data_format}"
+            )
+            raise ValueError(message)
 
     def build(self, input_shape):
         self._alpha_tensor = ops.convert_to_tensor(self.alpha, dtype=self.compute_dtype)

keras_hub/src/layers/preprocessing/random_elastic_deformation_3d_test.py

@@ -1,12 +1,16 @@
 # Add keras.utils for the random seed
-from keras import utils
 import os
-import numpy as np
+
 import keras
-from keras import Model
+import numpy as np
 from keras import Input
+from keras import Model
 from keras import ops
-from keras_hub.src.layers.preprocessing.random_elastic_deformation_3d import RandomElasticDeformation3D
+from keras import utils
+
+from keras_hub.src.layers.preprocessing.random_elastic_deformation_3d import (
+    RandomElasticDeformation3D,
+)
 from keras_hub.src.tests.test_case import TestCase
 
 
@@ -44,12 +48,15 @@ def test_serialization(self):
         model = Model(inputs=[image_input, label_input], outputs=outputs)
         original_output_image, original_output_label = model(input_data)
         path = os.path.join(self.get_temp_dir(), "model.keras")
-        
+
         # --- FIX: Remove the deprecated save_format argument ---
         model.save(path)
-        
+
         loaded_model = keras.models.load_model(
-            path, custom_objects={"RandomElasticDeformation3D": RandomElasticDeformation3D}
+            path,
+            custom_objects={
+                "RandomElasticDeformation3D": RandomElasticDeformation3D
+            },
         )
         loaded_output_image, loaded_output_label = loaded_model(input_data)
         np.testing.assert_allclose(
@@ -71,4 +78,4 @@ def test_label_values_are_preserved(self):
         _, output_label = layer((image, label))
         output_values = set(np.unique(ops.convert_to_numpy(output_label)))
         expected_values = {0, 1, 2, 3}
-        self.assertLessEqual(output_values, expected_values)
+        self.assertLessEqual(output_values, expected_values)