feat(layers): Add 3D elastic deformation layer

Author: ashpakshaikh26732

keras_hub/src/layers/preprocessing/random_elastic_deformation_3d.py

@@ -21,6 +21,7 @@ def __init__(self,
         self.sigma = sigma
         self.data_format = data_format
         if data_format not in ["channels_last", "channels_first"]:
+
             message = (
                 "`data_format` must be one of 'channels_last' or "
                 f"'channels_first'. Received: {self.data_format}"

keras_hub/src/layers/preprocessing/random_elastic_deformation_3d_test.py

@@ -1,4 +1,4 @@
-# Add keras.utils for the random seed
+
 import os
 
 import keras
@@ -16,7 +16,7 @@
 
 class RandomElasticDeformation3DTest(TestCase):
     def test_layer_basics(self):
-        # --- BEST PRACTICE: Add a seed for reproducibility ---
+
         utils.set_random_seed(0)
         layer = RandomElasticDeformation3D(
             grid_size=(4, 4, 4),
@@ -32,8 +32,6 @@ def test_layer_basics(self):
         self.assertEqual(label.dtype, output_label.dtype)
 
     def test_serialization(self):
-        # --- BEST PRACTICE: Add a seed for reproducibility ---
-        utils.set_random_seed(0)
         layer = RandomElasticDeformation3D(
             grid_size=(3, 3, 3),
             alpha=50.0,
@@ -42,34 +40,39 @@ def test_serialization(self):
         image_data = ops.ones((2, 16, 16, 16, 3), dtype="float32")
         label_data = ops.ones((2, 16, 16, 16, 1), dtype="int32")
         input_data = (image_data, label_data)
+
         image_input = Input(shape=(16, 16, 16, 3), dtype="float32")
         label_input = Input(shape=(16, 16, 16, 1), dtype="int32")
         outputs = layer((image_input, label_input))
         model = Model(inputs=[image_input, label_input], outputs=outputs)
+
+
+        utils.set_random_seed(0)
         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 ---
+        path = os.path.join(self.get_temp_dir(), "model.keras")
         model.save(path)
-
         loaded_model = keras.models.load_model(
-            path,
-            custom_objects={
-                "RandomElasticDeformation3D": RandomElasticDeformation3D
-            },
+            path, custom_objects={"RandomElasticDeformation3D": RandomElasticDeformation3D}
         )
+
+
+        utils.set_random_seed(0)
         loaded_output_image, loaded_output_label = loaded_model(input_data)
+
+
         np.testing.assert_allclose(
             ops.convert_to_numpy(original_output_image),
             ops.convert_to_numpy(loaded_output_image),
+            atol=1e-6 
         )
         np.testing.assert_array_equal(
             ops.convert_to_numpy(original_output_label),
             ops.convert_to_numpy(loaded_output_label),
         )
 
     def test_label_values_are_preserved(self):
-        # --- BEST PRACTICE: Add a seed for reproducibility ---
+ 
         utils.set_random_seed(0)
         image = ops.zeros(shape=(1, 16, 16, 16, 1), dtype="float32")
         label_arange = ops.arange(16**3, dtype="int32")