diff --git a/keras_cv/layers/preprocessing/aug_mix_test.py b/keras_cv/layers/preprocessing/aug_mix_test.py
index a6b3a6d4ab..2adba4eab5 100644
--- a/keras_cv/layers/preprocessing/aug_mix_test.py
+++ b/keras_cv/layers/preprocessing/aug_mix_test.py
@@ -27,16 +27,16 @@ def test_return_shapes(self):
xs = layer(xs)
ys_segmentation_masks = tf.ones((2, 512, 512, 3))
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
# greyscale
xs = tf.ones((2, 512, 512, 1))
xs = layer(xs)
ys_segmentation_masks = tf.ones((2, 512, 512, 1))
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [2, 512, 512, 1])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 1])
+ self.assertEqual(xs.shape, (2, 512, 512, 1))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 1))
def test_in_single_image_and_mask(self):
layer = preprocessing.AugMix([0, 255])
@@ -54,8 +54,8 @@ def test_in_single_image_and_mask(self):
)
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [512, 512, 3])
- self.assertEqual(ys_segmentation_masks.shape, [512, 512, 3])
+ self.assertEqual(xs.shape, (512, 512, 3))
+ self.assertEqual(ys_segmentation_masks.shape, (512, 512, 3))
# greyscale
xs = tf.cast(
@@ -69,8 +69,8 @@ def test_in_single_image_and_mask(self):
dtype=tf.float32,
)
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [512, 512, 1])
- self.assertEqual(ys_segmentation_masks.shape, [512, 512, 1])
+ self.assertEqual(xs.shape, (512, 512, 1))
+ self.assertEqual(ys_segmentation_masks.shape, (512, 512, 1))
def test_non_square_images_and_masks(self):
layer = preprocessing.AugMix([0, 255])
@@ -80,16 +80,16 @@ def test_non_square_images_and_masks(self):
xs = layer(xs)
ys_segmentation_masks = tf.ones((2, 256, 512, 3))
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [2, 256, 512, 3])
- self.assertEqual(ys_segmentation_masks.shape, [2, 256, 512, 3])
+ self.assertEqual(xs.shape, (2, 256, 512, 3))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 256, 512, 3))
# greyscale
xs = tf.ones((2, 256, 512, 1))
xs = layer(xs)
ys_segmentation_masks = tf.ones((2, 256, 512, 1))
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [2, 256, 512, 1])
- self.assertEqual(ys_segmentation_masks.shape, [2, 256, 512, 1])
+ self.assertEqual(xs.shape, (2, 256, 512, 1))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 256, 512, 1))
def test_single_input_args(self):
layer = preprocessing.AugMix([0, 255])
@@ -99,16 +99,16 @@ def test_single_input_args(self):
xs = layer(xs)
ys_segmentation_masks = tf.ones((2, 512, 512, 3))
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
# greyscale
xs = tf.ones((2, 512, 512, 1))
xs = layer(xs)
ys_segmentation_masks = tf.ones((2, 512, 512, 1))
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [2, 512, 512, 1])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 1])
+ self.assertEqual(xs.shape, (2, 512, 512, 1))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 1))
def test_many_augmentations(self):
layer = preprocessing.AugMix([0, 255], chain_depth=[25, 26])
@@ -118,13 +118,13 @@ def test_many_augmentations(self):
xs = layer(xs)
ys_segmentation_masks = tf.ones((2, 512, 512, 3))
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
# greyscale
xs = tf.ones((2, 512, 512, 1))
xs = layer(xs)
ys_segmentation_masks = tf.ones((2, 512, 512, 1))
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [2, 512, 512, 1])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 1])
+ self.assertEqual(xs.shape, (2, 512, 512, 1))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 1))
diff --git a/keras_cv/layers/preprocessing/auto_contrast_test.py b/keras_cv/layers/preprocessing/auto_contrast_test.py
index 9b87c2f100..bc34d5fae9 100644
--- a/keras_cv/layers/preprocessing/auto_contrast_test.py
+++ b/keras_cv/layers/preprocessing/auto_contrast_test.py
@@ -12,53 +12,53 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+import numpy as np
-import tensorflow as tf
-
+from keras_cv.backend import ops
from keras_cv.layers import preprocessing
from keras_cv.tests.test_case import TestCase
class AutoContrastTest(TestCase):
def test_constant_channels_dont_get_nanned(self):
- img = tf.constant([1, 1], dtype=tf.float32)
- img = tf.expand_dims(img, axis=-1)
- img = tf.expand_dims(img, axis=-1)
- img = tf.expand_dims(img, axis=0)
+ img = np.array([1, 1], dtype=np.float32)
+ img = np.expand_dims(img, axis=-1)
+ img = np.expand_dims(img, axis=-1)
+ img = np.expand_dims(img, axis=0)
layer = preprocessing.AutoContrast(value_range=(0, 255))
ys = layer(img)
- self.assertTrue(tf.math.reduce_any(ys[0] == 1.0))
- self.assertTrue(tf.math.reduce_any(ys[0] == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 1.0))
def test_auto_contrast_expands_value_range(self):
- img = tf.constant([0, 128], dtype=tf.float32)
- img = tf.expand_dims(img, axis=-1)
- img = tf.expand_dims(img, axis=-1)
- img = tf.expand_dims(img, axis=0)
+ img = np.array([0, 128], dtype=np.float32)
+ img = np.expand_dims(img, axis=-1)
+ img = np.expand_dims(img, axis=-1)
+ img = np.expand_dims(img, axis=0)
layer = preprocessing.AutoContrast(value_range=(0, 255))
ys = layer(img)
- self.assertTrue(tf.math.reduce_any(ys[0] == 0.0))
- self.assertTrue(tf.math.reduce_any(ys[0] == 255.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 0.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 255.0))
def test_auto_contrast_different_values_per_channel(self):
- img = tf.constant(
+ img = np.array(
[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]],
- dtype=tf.float32,
+ dtype=np.float32,
)
- img = tf.expand_dims(img, axis=0)
+ img = np.expand_dims(img, axis=0)
layer = preprocessing.AutoContrast(value_range=(0, 255))
ys = layer(img)
- self.assertTrue(tf.math.reduce_any(ys[0, ..., 0] == 0.0))
- self.assertTrue(tf.math.reduce_any(ys[0, ..., 1] == 0.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 0]) == 0.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 1]) == 0.0))
- self.assertTrue(tf.math.reduce_any(ys[0, ..., 0] == 255.0))
- self.assertTrue(tf.math.reduce_any(ys[0, ..., 1] == 255.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 0]) == 255.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 1]) == 255.0))
self.assertAllClose(
ys,
@@ -71,25 +71,25 @@ def test_auto_contrast_different_values_per_channel(self):
)
def test_auto_contrast_expands_value_range_uint8(self):
- img = tf.constant([0, 128], dtype=tf.uint8)
- img = tf.expand_dims(img, axis=-1)
- img = tf.expand_dims(img, axis=-1)
- img = tf.expand_dims(img, axis=0)
+ img = np.array([0, 128], dtype=np.uint8)
+ img = np.expand_dims(img, axis=-1)
+ img = np.expand_dims(img, axis=-1)
+ img = np.expand_dims(img, axis=0)
layer = preprocessing.AutoContrast(value_range=(0, 255))
ys = layer(img)
- self.assertTrue(tf.math.reduce_any(ys[0] == 0.0))
- self.assertTrue(tf.math.reduce_any(ys[0] == 255.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 0.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 255.0))
def test_auto_contrast_properly_converts_value_range(self):
- img = tf.constant([0, 0.5], dtype=tf.float32)
- img = tf.expand_dims(img, axis=-1)
- img = tf.expand_dims(img, axis=-1)
- img = tf.expand_dims(img, axis=0)
+ img = np.array([0, 0.5], dtype=np.float32)
+ img = np.expand_dims(img, axis=-1)
+ img = np.expand_dims(img, axis=-1)
+ img = np.expand_dims(img, axis=0)
layer = preprocessing.AutoContrast(value_range=(0, 1))
ys = layer(img)
- self.assertTrue(tf.math.reduce_any(ys[0] == 0.0))
- self.assertTrue(tf.math.reduce_any(ys[0] == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 0.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 1.0))
diff --git a/keras_cv/layers/preprocessing/base_image_augmentation_layer.py b/keras_cv/layers/preprocessing/base_image_augmentation_layer.py
index 0a365891b7..2eba5aaa74 100644
--- a/keras_cv/layers/preprocessing/base_image_augmentation_layer.py
+++ b/keras_cv/layers/preprocessing/base_image_augmentation_layer.py
@@ -14,6 +14,7 @@
import keras
import tensorflow as tf
+import tree
if hasattr(keras, "src"):
keras_backend = keras.src.backend
@@ -23,8 +24,8 @@
from keras_cv import bounding_box
from keras_cv.api_export import keras_cv_export
from keras_cv.backend import keras
+from keras_cv.backend import ops
from keras_cv.backend import scope
-from keras_cv.backend.config import keras_3
from keras_cv.utils import preprocessing
# In order to support both unbatched and batched inputs, the horizontal
@@ -42,15 +43,8 @@
USE_TARGETS = "use_targets"
-base_class = (
- keras.src.layers.preprocessing.tf_data_layer.TFDataLayer
- if keras_3()
- else keras.layers.Layer
-)
-
-
@keras_cv_export("keras_cv.layers.BaseImageAugmentationLayer")
-class BaseImageAugmentationLayer(base_class):
+class BaseImageAugmentationLayer(keras.layers.Layer):
"""Abstract base layer for image augmentation.
This layer contains base functionalities for preprocessing layers which
@@ -415,6 +409,19 @@ def get_random_transformation(
return None
def call(self, inputs):
+ # try to convert a given backend native tensor to TensorFlow tensor
+ # before passing it over to TFDataScope
+ contains_ragged = lambda y: any(
+ tree.map_structure(
+ lambda x: isinstance(x, (tf.RaggedTensor, tf.SparseTensor)),
+ tree.flatten(y),
+ )
+ )
+ inputs_contain_ragged = contains_ragged(inputs)
+ if not inputs_contain_ragged:
+ inputs = tree.map_structure(
+ lambda x: tf.convert_to_tensor(x), inputs
+ )
with scope.TFDataScope():
inputs = self._ensure_inputs_are_compute_dtype(inputs)
inputs, metadata = self._format_inputs(inputs)
@@ -431,7 +438,20 @@ def call(self, inputs):
"rank 3 (HWC) or 4D (NHWC) tensors. Got shape: "
f"{images.shape}"
)
- return outputs
+ # convert the outputs to backend native tensors if none of them
+ # contain RaggedTensors. Note that if the user passed in Raggeds
+ # but the outputs are dense, we still don't want to convert to
+ # backend native tensors. This is to avoid breaking TF data
+ # pipelines that can't easily be ported to become backend
+ # agnostic.
+ if not inputs_contain_ragged and not contains_ragged(outputs):
+ outputs = tree.map_structure(
+ # some layers return None, handle that case when
+ # converting to tensors
+ lambda x: ops.convert_to_tensor(x) if x is not None else x,
+ outputs,
+ )
+ return outputs
def _augment(self, inputs):
raw_image = inputs.get(IMAGES, None)
diff --git a/keras_cv/layers/preprocessing/base_image_augmentation_layer_test.py b/keras_cv/layers/preprocessing/base_image_augmentation_layer_test.py
index 50cf8b5049..aed4dd3af0 100644
--- a/keras_cv/layers/preprocessing/base_image_augmentation_layer_test.py
+++ b/keras_cv/layers/preprocessing/base_image_augmentation_layer_test.py
@@ -12,9 +12,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
+import pytest
import tensorflow as tf
from keras_cv import bounding_box
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
BaseImageAugmentationLayer,
)
@@ -78,17 +80,17 @@ def test_augment_dict_return_type(self):
def test_augment_casts_dtypes(self):
add_layer = RandomAddLayer(fixed_value=2.0)
- images = tf.ones((2, 8, 8, 3), dtype="uint8")
+ images = np.ones((2, 8, 8, 3), dtype="uint8")
output = add_layer(images)
self.assertAllClose(
- tf.ones((2, 8, 8, 3), dtype="float32") * 3.0, output
+ np.ones((2, 8, 8, 3), dtype="float32") * 3.0, output
)
def test_augment_batch_images(self):
add_layer = RandomAddLayer()
images = np.random.random(size=(2, 8, 8, 3)).astype("float32")
- output = add_layer(images)
+ output = ops.convert_to_numpy(add_layer(images))
diff = output - images
# Make sure the first image and second image get different augmentation
@@ -118,8 +120,8 @@ def test_augment_batch_images_and_targets(self):
targets = np.random.random(size=(2, 1)).astype("float32")
output = add_layer({"images": images, "targets": targets})
- image_diff = output["images"] - images
- label_diff = output["targets"] - targets
+ image_diff = ops.convert_to_numpy(output["images"]) - images
+ label_diff = ops.convert_to_numpy(output["targets"]) - targets
# Make sure the first image and second image get different augmentation
self.assertNotAllClose(image_diff[0], image_diff[1])
self.assertNotAllClose(label_diff[0], label_diff[1])
@@ -225,6 +227,7 @@ def test_augment_batch_image_and_localization_data(self):
segmentation_mask_diff[0], segmentation_mask_diff[1]
)
+ @pytest.mark.tf_only
def test_augment_all_data_in_tf_function(self):
add_layer = RandomAddLayer()
images = np.random.random(size=(2, 8, 8, 3)).astype("float32")
diff --git a/keras_cv/layers/preprocessing/channel_shuffle_test.py b/keras_cv/layers/preprocessing/channel_shuffle_test.py
index a14f138ee1..e608e5584d 100644
--- a/keras_cv/layers/preprocessing/channel_shuffle_test.py
+++ b/keras_cv/layers/preprocessing/channel_shuffle_test.py
@@ -15,6 +15,7 @@
import pytest
import tensorflow as tf
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.channel_shuffle import ChannelShuffle
from keras_cv.tests.test_case import TestCase
@@ -25,7 +26,7 @@ def test_return_shapes(self):
layer = ChannelShuffle(groups=3)
xs = layer(xs, training=True)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
def test_channel_shuffle_call_results_one_channel(self):
xs = tf.cast(
@@ -38,8 +39,8 @@ def test_channel_shuffle_call_results_one_channel(self):
layer = ChannelShuffle(groups=1)
xs = layer(xs, training=True)
- self.assertTrue(tf.math.reduce_any(xs[0] == 3.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 3.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
def test_channel_shuffle_call_results_multi_channel(self):
xs = tf.cast(
@@ -52,8 +53,8 @@ def test_channel_shuffle_call_results_multi_channel(self):
layer = ChannelShuffle(groups=5)
xs = layer(xs, training=True)
- self.assertTrue(tf.math.reduce_any(xs[0] == 3.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 3.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
def test_non_square_image(self):
xs = tf.cast(
@@ -66,9 +67,10 @@ def test_non_square_image(self):
layer = ChannelShuffle(groups=1)
xs = layer(xs, training=True)
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+ @pytest.mark.tf_only
def test_in_tf_function(self):
xs = tf.cast(
tf.stack(
@@ -84,8 +86,8 @@ def augment(x):
return layer(x, training=True)
xs = augment(xs)
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
def test_in_single_image(self):
xs = tf.cast(
@@ -95,7 +97,7 @@ def test_in_single_image(self):
layer = ChannelShuffle(groups=1)
xs = layer(xs, training=True)
- self.assertTrue(tf.math.reduce_any(xs == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs) == 1.0))
@pytest.mark.skip(reason="flaky")
def test_channel_shuffle_on_batched_images_independently(self):
@@ -116,9 +118,11 @@ def test_config_with_custom_name(self):
def test_output_dtypes(self):
inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
layer = ChannelShuffle(groups=1)
- self.assertAllEqual(layer(inputs).dtype, "float32")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+ )
layer = ChannelShuffle(groups=1, dtype="uint8")
- self.assertAllEqual(layer(inputs).dtype, "uint8")
+ self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
def test_config(self):
layer = ChannelShuffle(groups=5)
diff --git a/keras_cv/layers/preprocessing/cut_mix_test.py b/keras_cv/layers/preprocessing/cut_mix_test.py
index 09f398328d..4dfa3120e8 100644
--- a/keras_cv/layers/preprocessing/cut_mix_test.py
+++ b/keras_cv/layers/preprocessing/cut_mix_test.py
@@ -11,8 +11,11 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import numpy as np
+import pytest
import tensorflow as tf
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.cut_mix import CutMix
from keras_cv.tests.test_case import TestCase
@@ -51,9 +54,9 @@ def test_return_shapes(self):
outputs["segmentation_masks"],
)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
- self.assertEqual(ys_labels.shape, [2, 10])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
+ self.assertEqual(ys_labels.shape, (2, 10))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
def test_cut_mix_call_results_with_labels(self):
xs = tf.cast(
@@ -70,10 +73,10 @@ def test_cut_mix_call_results_with_labels(self):
xs, ys = outputs["images"], outputs["labels"]
# At least some pixels should be replaced in the CutMix operation
- self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
# No labels should still be close to their original values
self.assertNotAllClose(ys, 1.0)
self.assertNotAllClose(ys, 0.0)
@@ -93,10 +96,10 @@ def test_cut_mix_call_results_one_channel_with_labels(self):
xs, ys = outputs["images"], outputs["labels"]
# At least some pixels should be replaced in the CutMix operation
- self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
# No labels should still be close to their original values
self.assertNotAllClose(ys, 1.0)
self.assertNotAllClose(ys, 0.0)
@@ -128,15 +131,23 @@ def test_cut_mix_call_results_with_dense_encoded_segmentation_masks(self):
)
# At least some pixels should be replaced in the images
- self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
# At least some pixels should be replaced in the segmentation_masks
- self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[0] == 1.0))
- self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[1] == 1.0))
- self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[1] == 2.0))
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(ys_segmentation_masks[0]) == 1.0)
+ )
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(ys_segmentation_masks[0]) == 2.0)
+ )
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(ys_segmentation_masks[1]) == 1.0)
+ )
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(ys_segmentation_masks[1]) == 2.0)
+ )
def test_cut_mix_call_results_with_one_hot_encoded_segmentation_masks(self):
xs = tf.cast(
@@ -166,24 +177,33 @@ def test_cut_mix_call_results_with_one_hot_encoded_segmentation_masks(self):
)
# At least some pixels should be replaced in the images
- self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
# At least some pixels should be replaced in the segmentation_masks
self.assertTrue(
- tf.math.reduce_any(ys_segmentation_masks[0][:, :, 2] == 1.0)
+ np.any(
+ ops.convert_to_numpy(ys_segmentation_masks[0][:, :, 2]) == 1.0
+ )
)
self.assertTrue(
- tf.math.reduce_any(ys_segmentation_masks[0][:, :, 2] == 0.0)
+ np.any(
+ ops.convert_to_numpy(ys_segmentation_masks[0][:, :, 2]) == 0.0
+ )
)
self.assertTrue(
- tf.math.reduce_any(ys_segmentation_masks[1][:, :, 1] == 1.0)
+ np.any(
+ ops.convert_to_numpy(ys_segmentation_masks[1][:, :, 1]) == 1.0
+ )
)
self.assertTrue(
- tf.math.reduce_any(ys_segmentation_masks[1][:, :, 1] == 0.0)
+ np.any(
+ ops.convert_to_numpy(ys_segmentation_masks[1][:, :, 1]) == 0.0
+ )
)
+ @pytest.mark.tf_only
def test_in_tf_function(self):
xs = tf.cast(
tf.stack(
@@ -203,10 +223,10 @@ def augment(x, y):
xs, ys = outputs["images"], outputs["labels"]
# At least some pixels should be replaced in the CutMix operation
- self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
# No labels should still be close to their original values
self.assertNotAllClose(ys, 1.0)
self.assertNotAllClose(ys, 0.0)
diff --git a/keras_cv/layers/preprocessing/equalization_test.py b/keras_cv/layers/preprocessing/equalization_test.py
index b6849ca142..937f19ca28 100644
--- a/keras_cv/layers/preprocessing/equalization_test.py
+++ b/keras_cv/layers/preprocessing/equalization_test.py
@@ -12,23 +12,24 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+import numpy as np
import pytest
-import tensorflow as tf
from absl.testing import parameterized
from keras_cv.backend import keras
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.equalization import Equalization
from keras_cv.tests.test_case import TestCase
class EqualizationTest(TestCase):
def test_return_shapes(self):
- xs = 255 * tf.ones((2, 512, 512, 3), dtype=tf.int32)
+ xs = 255 * np.ones((2, 512, 512, 3), dtype=np.int32)
layer = Equalization(value_range=(0, 255))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
- self.assertAllEqual(xs, 255 * tf.ones((2, 512, 512, 3)))
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
+ self.assertAllEqual(xs, 255 * np.ones((2, 512, 512, 3)))
@pytest.mark.tf_keras_only
def test_return_shapes_inside_model(self):
@@ -40,28 +41,34 @@ def test_return_shapes_inside_model(self):
self.assertEqual(model.output_shape, (None, 512, 512, 5))
def test_equalizes_to_all_bins(self):
- xs = tf.random.uniform((2, 512, 512, 3), 0, 255, dtype=tf.float32)
+ xs = np.random.uniform(size=(2, 512, 512, 3), low=0, high=255).astype(
+ np.float32
+ )
layer = Equalization(value_range=(0, 255))
xs = layer(xs)
for i in range(0, 256):
- self.assertTrue(tf.math.reduce_any(xs == i))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs) == i))
@parameterized.named_parameters(
- ("float32", tf.float32), ("int32", tf.int32), ("int64", tf.int64)
+ ("float32", np.float32), ("int32", np.int32), ("int64", np.int64)
)
def test_input_dtypes(self, dtype):
- xs = tf.random.uniform((2, 512, 512, 3), 0, 255, dtype=dtype)
+ xs = np.random.uniform(size=(2, 512, 512, 3), low=0, high=255).astype(
+ dtype
+ )
layer = Equalization(value_range=(0, 255))
- xs = layer(xs)
+ xs = ops.convert_to_numpy(layer(xs))
for i in range(0, 256):
- self.assertTrue(tf.math.reduce_any(xs == i))
+ self.assertTrue(np.any(xs == i))
self.assertAllInRange(xs, 0, 255)
@parameterized.named_parameters(("0_255", 0, 255), ("0_1", 0, 1))
def test_output_range(self, lower, upper):
- xs = tf.random.uniform((2, 512, 512, 3), lower, upper, dtype=tf.float32)
+ xs = np.random.uniform(
+ size=(2, 512, 512, 3), low=lower, high=upper
+ ).astype(np.float32)
layer = Equalization(value_range=(lower, upper))
- xs = layer(xs)
+ xs = ops.convert_to_numpy(layer(xs))
self.assertAllInRange(xs, lower, upper)
diff --git a/keras_cv/layers/preprocessing/fourier_mix_test.py b/keras_cv/layers/preprocessing/fourier_mix_test.py
index d033aa83db..cc97c84017 100644
--- a/keras_cv/layers/preprocessing/fourier_mix_test.py
+++ b/keras_cv/layers/preprocessing/fourier_mix_test.py
@@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import pytest
import tensorflow as tf
from keras_cv.layers.preprocessing.fourier_mix import FourierMix
@@ -50,9 +51,9 @@ def test_return_shapes(self):
outputs["segmentation_masks"],
)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
- self.assertEqual(ys.shape, [2, 10])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
+ self.assertEqual(ys.shape, (2, 10))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
def test_fourier_mix_call_results_with_labels(self):
xs = tf.cast(
@@ -110,6 +111,7 @@ def test_mix_up_call_results_with_masks(self):
self.assertNotAllClose(ys_segmentation_masks, 1.0)
self.assertNotAllClose(ys_segmentation_masks, 0.0)
+ @pytest.mark.tf_only
def test_in_tf_function(self):
xs = tf.cast(
tf.stack(
diff --git a/keras_cv/layers/preprocessing/grayscale_test.py b/keras_cv/layers/preprocessing/grayscale_test.py
index 07072475e6..1a52d34019 100644
--- a/keras_cv/layers/preprocessing/grayscale_test.py
+++ b/keras_cv/layers/preprocessing/grayscale_test.py
@@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import pytest
import tensorflow as tf
from keras_cv.layers import preprocessing
@@ -31,9 +32,10 @@ def test_return_shapes(self):
)
xs2 = layer(xs, training=True)
- self.assertEqual(xs1.shape, [2, 52, 24, 1])
- self.assertEqual(xs2.shape, [2, 52, 24, 3])
+ self.assertEqual(xs1.shape, (2, 52, 24, 1))
+ self.assertEqual(xs2.shape, (2, 52, 24, 3))
+ @pytest.mark.tf_only
def test_in_tf_function(self):
xs = tf.cast(
tf.stack([2 * tf.ones((10, 10, 3)), tf.ones((10, 10, 3))], axis=0),
@@ -62,8 +64,8 @@ def augment(x):
xs2 = augment(xs)
- self.assertEqual(xs1.shape, [2, 10, 10, 1])
- self.assertEqual(xs2.shape, [2, 10, 10, 3])
+ self.assertEqual(xs1.shape, (2, 10, 10, 1))
+ self.assertEqual(xs2.shape, (2, 10, 10, 3))
def test_non_square_image(self):
xs = tf.cast(
@@ -81,8 +83,8 @@ def test_non_square_image(self):
)
xs2 = layer(xs, training=True)
- self.assertEqual(xs1.shape, [2, 52, 24, 1])
- self.assertEqual(xs2.shape, [2, 52, 24, 3])
+ self.assertEqual(xs1.shape, (2, 52, 24, 1))
+ self.assertEqual(xs2.shape, (2, 52, 24, 3))
def test_in_single_image(self):
xs = tf.cast(
@@ -100,5 +102,5 @@ def test_in_single_image(self):
)
xs2 = layer(xs, training=True)
- self.assertEqual(xs1.shape, [52, 24, 1])
- self.assertEqual(xs2.shape, [52, 24, 3])
+ self.assertEqual(xs1.shape, (52, 24, 1))
+ self.assertEqual(xs2.shape, (52, 24, 3))
diff --git a/keras_cv/layers/preprocessing/grid_mask_test.py b/keras_cv/layers/preprocessing/grid_mask_test.py
index 3226a0ddd7..d06f9df9b8 100644
--- a/keras_cv/layers/preprocessing/grid_mask_test.py
+++ b/keras_cv/layers/preprocessing/grid_mask_test.py
@@ -12,10 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-
+import numpy as np
+import pytest
import tensorflow as tf
import keras_cv
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.grid_mask import GridMask
from keras_cv.tests.test_case import TestCase
@@ -27,7 +29,7 @@ def test_return_shapes(self):
layer = GridMask(ratio_factor=0.1, rotation_factor=(-0.2, 0.3))
xs = layer(xs, training=True)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
def test_gridmask_call_results_one_channel(self):
xs = tf.cast(
@@ -48,10 +50,14 @@ def test_gridmask_call_results_one_channel(self):
xs = layer(xs, training=True)
# Some pixels should be replaced with fill_value
- self.assertTrue(tf.math.reduce_any(xs[0] == float(fill_value)))
- self.assertTrue(tf.math.reduce_any(xs[0] == 3.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == float(fill_value)))
- self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(xs[0]) == float(fill_value))
+ )
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 3.0))
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(xs[1]) == float(fill_value))
+ )
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
def test_non_square_image(self):
xs = tf.cast(
@@ -72,11 +78,16 @@ def test_non_square_image(self):
xs = layer(xs, training=True)
# Some pixels should be replaced with fill_value
- self.assertTrue(tf.math.reduce_any(xs[0] == float(fill_value)))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == float(fill_value)))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(xs[0]) == float(fill_value))
+ )
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(xs[1]) == float(fill_value))
+ )
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+ @pytest.mark.tf_only
def test_in_tf_function(self):
xs = tf.cast(
tf.stack(
@@ -100,10 +111,14 @@ def augment(x):
xs = augment(xs)
# Some pixels should be replaced with fill_value
- self.assertTrue(tf.math.reduce_any(xs[0] == float(fill_value)))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == float(fill_value)))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(xs[0]) == float(fill_value))
+ )
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(xs[1]) == float(fill_value))
+ )
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
def test_in_single_image(self):
xs = tf.cast(
@@ -115,5 +130,5 @@ def test_in_single_image(self):
ratio_factor=(0.5, 0.5), fill_mode="constant", fill_value=0.0
)
xs = layer(xs, training=True)
- self.assertTrue(tf.math.reduce_any(xs == 0.0))
- self.assertTrue(tf.math.reduce_any(xs == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs) == 0.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs) == 1.0))
diff --git a/keras_cv/layers/preprocessing/jittered_resize_test.py b/keras_cv/layers/preprocessing/jittered_resize_test.py
index 3bf2f1ad8f..dfa6a2ed69 100644
--- a/keras_cv/layers/preprocessing/jittered_resize_test.py
+++ b/keras_cv/layers/preprocessing/jittered_resize_test.py
@@ -17,6 +17,7 @@
from keras_cv import bounding_box
from keras_cv import core
from keras_cv import layers
+from keras_cv.backend import ops
from keras_cv.tests.test_case import TestCase
@@ -220,13 +221,15 @@ def test_output_dtypes(self):
target_size=self.target_size,
scale_factor=(3 / 4, 4 / 3),
)
- self.assertAllEqual(layer(inputs).dtype, "float32")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+ )
layer = layers.JitteredResize(
target_size=self.target_size,
scale_factor=(3 / 4, 4 / 3),
dtype="uint8",
)
- self.assertAllEqual(layer(inputs).dtype, "uint8")
+ self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
def test_config(self):
layer = layers.JitteredResize(
diff --git a/keras_cv/layers/preprocessing/mix_up_test.py b/keras_cv/layers/preprocessing/mix_up_test.py
index ec1e0af1a6..332e0e226c 100644
--- a/keras_cv/layers/preprocessing/mix_up_test.py
+++ b/keras_cv/layers/preprocessing/mix_up_test.py
@@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import pytest
import tensorflow as tf
from keras_cv.layers.preprocessing.mix_up import MixUp
@@ -60,11 +61,11 @@ def test_return_shapes(self):
outputs["segmentation_masks"],
)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
- self.assertEqual(ys_labels.shape, [2, 10])
- self.assertEqual(ys_bounding_boxes["boxes"].shape, [2, 6, 4])
- self.assertEqual(ys_bounding_boxes["classes"].shape, [2, 6])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
+ self.assertEqual(ys_labels.shape, (2, 10))
+ self.assertEqual(ys_bounding_boxes["boxes"].shape, (2, 6, 4))
+ self.assertEqual(ys_bounding_boxes["classes"].shape, (2, 6))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
def test_mix_up_call_results_with_labels(self):
xs = tf.cast(
@@ -122,6 +123,7 @@ def test_mix_up_call_results_with_masks(self):
self.assertNotAllClose(ys_segmentation_masks, 1.0)
self.assertNotAllClose(ys_segmentation_masks, 0.0)
+ @pytest.mark.tf_only
def test_in_tf_function(self):
xs = tf.cast(
tf.stack(
diff --git a/keras_cv/layers/preprocessing/mosaic_test.py b/keras_cv/layers/preprocessing/mosaic_test.py
index c5f5c9d1c1..5d2561731c 100644
--- a/keras_cv/layers/preprocessing/mosaic_test.py
+++ b/keras_cv/layers/preprocessing/mosaic_test.py
@@ -11,6 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import pytest
import tensorflow as tf
from keras_cv.layers.preprocessing.mosaic import Mosaic
@@ -59,6 +60,7 @@ def test_return_shapes(self):
self.assertEqual(ys_bounding_boxes["classes"].shape, [2, None])
self.assertEqual(ys_segmentation_masks.shape, input_shape)
+ @pytest.mark.tf_only
def test_in_tf_function(self):
xs = tf.cast(
tf.stack(
diff --git a/keras_cv/layers/preprocessing/posterization_test.py b/keras_cv/layers/preprocessing/posterization_test.py
index 0825ffe3a4..8ea37b5565 100644
--- a/keras_cv/layers/preprocessing/posterization_test.py
+++ b/keras_cv/layers/preprocessing/posterization_test.py
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
+import pytest
import tensorflow as tf
from keras_cv.layers.preprocessing.posterization import Posterization
@@ -74,6 +75,7 @@ def test_batched_input(self):
self.assertAllEqual(output, expected_output)
+ @pytest.mark.tf_only
def test_works_with_xla(self):
dummy_input = self.rng.uniform(shape=(2, 224, 224, 3))
layer = Posterization(bits=4, value_range=[0, 1])
diff --git a/keras_cv/layers/preprocessing/rand_augment_test.py b/keras_cv/layers/preprocessing/rand_augment_test.py
index 0f9759cc42..5f93349c29 100644
--- a/keras_cv/layers/preprocessing/rand_augment_test.py
+++ b/keras_cv/layers/preprocessing/rand_augment_test.py
@@ -16,6 +16,7 @@
from absl.testing import parameterized
from keras_cv import layers
+from keras_cv.backend import ops
from keras_cv.tests.test_case import TestCase
@@ -56,10 +57,8 @@ def test_runs_with_value_range(self, low, high):
value_range=(low, high),
)
xs = tf.random.uniform((2, 512, 512, 3), low, high, dtype=tf.float32)
- ys = rand_augment(xs)
- self.assertTrue(
- tf.math.reduce_all(tf.logical_and(ys >= low, ys <= high))
- )
+ ys = ops.convert_to_numpy(rand_augment(xs))
+ self.assertTrue(np.all(np.logical_and(ys >= low, ys <= high)))
@parameterized.named_parameters(
("float32", "float32"),
@@ -85,9 +84,9 @@ def test_standard_policy_respects_value_range(self, lower, upper):
layers=my_layers, augmentations_per_image=3
)
xs = tf.random.uniform((2, 512, 512, 3), lower, upper, dtype=tf.float32)
- ys = rand_augment(xs)
- self.assertLessEqual(tf.math.reduce_max(ys), upper)
- self.assertGreaterEqual(tf.math.reduce_min(ys), lower)
+ ys = ops.convert_to_numpy(rand_augment(xs))
+ self.assertLessEqual(np.max(ys), upper)
+ self.assertGreaterEqual(np.min(ys), lower)
def test_runs_unbatched(self):
rand_augment = layers.RandAugment(
diff --git a/keras_cv/layers/preprocessing/random_apply_test.py b/keras_cv/layers/preprocessing/random_apply_test.py
index 86a62eefe6..693e788be0 100644
--- a/keras_cv/layers/preprocessing/random_apply_test.py
+++ b/keras_cv/layers/preprocessing/random_apply_test.py
@@ -11,11 +11,12 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
-
+import pytest
import tensorflow as tf
from absl.testing import parameterized
from keras_cv import layers
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
BaseImageAugmentationLayer,
)
@@ -49,7 +50,7 @@ def test_works_with_batched_input(self):
dummy_inputs = self.rng.uniform(shape=(batch_size, 224, 224, 3))
layer = RandomApply(rate=0.5, layer=ZeroOut(), seed=1234)
- outputs = layer(dummy_inputs)
+ outputs = ops.convert_to_numpy(layer(dummy_inputs))
num_zero_inputs = self._num_zero_batches(dummy_inputs)
num_zero_outputs = self._num_zero_batches(outputs)
@@ -107,6 +108,7 @@ def test_can_modify_label(self):
self.assertAllEqual(outputs["labels"], tf.zeros_like(dummy_labels))
+ @pytest.mark.tf_only
def test_works_with_xla(self):
dummy_inputs = self.rng.uniform(shape=(32, 224, 224, 3))
# auto_vectorize=True will crash XLA
diff --git a/keras_cv/layers/preprocessing/random_channel_shift_test.py b/keras_cv/layers/preprocessing/random_channel_shift_test.py
index 9f1f724eb5..0b2b8b1086 100644
--- a/keras_cv/layers/preprocessing/random_channel_shift_test.py
+++ b/keras_cv/layers/preprocessing/random_channel_shift_test.py
@@ -13,8 +13,10 @@
# limitations under the License.
import numpy as np
+import pytest
import tensorflow as tf
+from keras_cv.backend import ops
from keras_cv.layers import preprocessing
from keras_cv.tests.test_case import TestCase
@@ -27,7 +29,7 @@ def test_return_shapes(self):
)
xs = layer(xs, training=True)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
def test_non_square_image(self):
xs = tf.cast(
@@ -42,9 +44,10 @@ def test_non_square_image(self):
)
xs = layer(xs, training=True)
- self.assertFalse(tf.math.reduce_any(xs[0] == 2.0))
- self.assertFalse(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertFalse(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertFalse(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+ @pytest.mark.tf_only
def test_in_tf_function(self):
xs = tf.cast(
tf.stack(
@@ -61,8 +64,8 @@ def augment(x):
return layer(x, training=True)
xs = augment(xs)
- self.assertFalse(tf.math.reduce_any(xs[0] == 2.0))
- self.assertFalse(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertFalse(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertFalse(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
def test_5_channels(self):
xs = tf.cast(
@@ -73,7 +76,7 @@ def test_5_channels(self):
factor=0.4, channels=5, value_range=(0, 255)
)
xs = layer(xs, training=True)
- self.assertFalse(tf.math.reduce_any(xs == 1.0))
+ self.assertFalse(np.any(ops.convert_to_numpy(xs) == 1.0))
def test_1_channel(self):
xs = tf.cast(
@@ -84,7 +87,7 @@ def test_1_channel(self):
factor=0.4, channels=1, value_range=(0, 255)
)
xs = layer(xs, training=True)
- self.assertFalse(tf.math.reduce_any(xs == 1.0))
+ self.assertFalse(np.any(ops.convert_to_numpy(xs) == 1.0))
def test_in_single_image(self):
xs = tf.cast(
@@ -95,7 +98,7 @@ def test_in_single_image(self):
factor=0.4, value_range=(0, 255)
)
xs = layer(xs, training=True)
- self.assertFalse(tf.math.reduce_any(xs == 1.0))
+ self.assertFalse(np.any(ops.convert_to_numpy(xs) == 1.0))
def test_config(self):
layer = preprocessing.RandomChannelShift(
diff --git a/keras_cv/layers/preprocessing/random_color_degeneration_test.py b/keras_cv/layers/preprocessing/random_color_degeneration_test.py
index 975a95b8ae..a27a7b67fd 100644
--- a/keras_cv/layers/preprocessing/random_color_degeneration_test.py
+++ b/keras_cv/layers/preprocessing/random_color_degeneration_test.py
@@ -14,6 +14,7 @@
import numpy as np
import tensorflow as tf
+from keras_cv.backend import ops
from keras_cv.layers import preprocessing
from keras_cv.tests.test_case import TestCase
@@ -39,7 +40,7 @@ def test_color_degeneration_full_factor(self):
xs = tf.concat([r, g, b], axis=-1)
layer = preprocessing.RandomColorDegeneration(factor=(1, 1))
- ys = layer(xs)
+ ys = ops.convert_to_numpy(layer(xs))
# Color degeneration uses standard luma conversion for RGB->Grayscale.
# The formula for luma is result= 0.2989*r + 0.5870*g + 0.1140*b
@@ -54,7 +55,7 @@ def test_color_degeneration_70p_factor(self):
xs = tf.concat([r, g, b], axis=-1)
layer = preprocessing.RandomColorDegeneration(factor=(0.7, 0.7))
- ys = layer(xs)
+ ys = ops.convert_to_numpy(layer(xs))
# Color degeneration uses standard luma conversion for RGB->Grayscale.
# The formula for luma is result= 0.2989*r + 0.5870*g + 0.1140*b
diff --git a/keras_cv/layers/preprocessing/random_color_jitter_test.py b/keras_cv/layers/preprocessing/random_color_jitter_test.py
index 17fa2f8b60..24387d962a 100644
--- a/keras_cv/layers/preprocessing/random_color_jitter_test.py
+++ b/keras_cv/layers/preprocessing/random_color_jitter_test.py
@@ -13,6 +13,7 @@
# limitations under the License.
import numpy as np
+import pytest
import tensorflow as tf
from keras_cv.layers import preprocessing
@@ -37,9 +38,9 @@ def test_return_shapes(self):
non_square_batch_output = layer(non_square_batch_input, training=True)
unbatch_output = layer(unbatch_input, training=True)
- self.assertEqual(batch_output.shape, [2, 512, 512, 3])
- self.assertEqual(non_square_batch_output.shape, [2, 1024, 512, 3])
- self.assertEqual(unbatch_output.shape, [512, 512, 3])
+ self.assertEqual(batch_output.shape, (2, 512, 512, 3))
+ self.assertEqual(non_square_batch_output.shape, (2, 1024, 512, 3))
+ self.assertEqual(unbatch_output.shape, (512, 512, 3))
# Test 2: Check if the factor ranges are set properly.
def test_factor_range(self):
@@ -57,6 +58,7 @@ def test_factor_range(self):
self.assertEqual(layer.hue_factor, (0.5, 0.9))
# Test 3: Test if it is OK to run on graph mode.
+ @pytest.mark.tf_only
def test_in_tf_function(self):
inputs = np.ones((2, 512, 512, 3))
diff --git a/keras_cv/layers/preprocessing/random_crop_and_resize_test.py b/keras_cv/layers/preprocessing/random_crop_and_resize_test.py
index dc8c2b76b8..34d4613f93 100644
--- a/keras_cv/layers/preprocessing/random_crop_and_resize_test.py
+++ b/keras_cv/layers/preprocessing/random_crop_and_resize_test.py
@@ -16,6 +16,7 @@
from absl.testing import parameterized
from keras_cv import bounding_box
+from keras_cv.backend import ops
from keras_cv.layers import preprocessing
from keras_cv.tests.test_case import TestCase
@@ -147,7 +148,9 @@ def test_augment_sparse_segmentation_mask(self):
seed=self.seed,
)
output = layer(inputs, training=True)
- self.assertAllInSet(output["segmentation_masks"], [0, 7])
+ self.assertAllInSet(
+ ops.convert_to_numpy(output["segmentation_masks"]), [0, 7]
+ )
def test_augment_one_hot_segmentation_mask(self):
num_classes = 8
diff --git a/keras_cv/layers/preprocessing/random_crop_test.py b/keras_cv/layers/preprocessing/random_crop_test.py
index 9de5fdd2ad..bd00b0e21b 100644
--- a/keras_cv/layers/preprocessing/random_crop_test.py
+++ b/keras_cv/layers/preprocessing/random_crop_test.py
@@ -20,6 +20,7 @@
from absl.testing import parameterized
from keras_cv import layers as cv_layers
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.random_crop import RandomCrop
from keras_cv.tests.test_case import TestCase
@@ -185,6 +186,7 @@ def test_augment_bounding_boxes_resize(self):
)
self.assertAllClose(expected_output, output["bounding_boxes"]["boxes"])
+ @pytest.mark.tf_only
def test_in_tf_function(self):
np.random.seed(1337)
inp = np.random.random((20, 16, 16, 3))
@@ -252,9 +254,11 @@ def test_config_with_custom_name(self):
def test_output_dtypes(self):
inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
layer = RandomCrop(2, 2)
- self.assertAllEqual(layer(inputs).dtype, "float32")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+ )
layer = RandomCrop(2, 2, dtype="uint8")
- self.assertAllEqual(layer(inputs).dtype, "uint8")
+ self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
def test_config(self):
layer = RandomCrop(height=2, width=3, bounding_box_format="xyxy")
diff --git a/keras_cv/layers/preprocessing/random_cutout_test.py b/keras_cv/layers/preprocessing/random_cutout_test.py
index 14930b6fd1..b8d549b0d0 100644
--- a/keras_cv/layers/preprocessing/random_cutout_test.py
+++ b/keras_cv/layers/preprocessing/random_cutout_test.py
@@ -12,8 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
+import pytest
import tensorflow as tf
+from keras_cv.backend import ops
from keras_cv.layers import preprocessing
from keras_cv.tests.test_case import TestCase
@@ -38,10 +40,10 @@ def _run_test(self, height_factor, width_factor):
xs = layer(xs)
# Some pixels should be replaced with fill value
- self.assertTrue(tf.math.reduce_any(xs[0] == fill_value))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == fill_value))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == fill_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == fill_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
def test_return_shapes(self):
xs = np.ones((2, 512, 512, 3))
@@ -53,8 +55,8 @@ def test_return_shapes(self):
xs = layer(xs)
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
def test_return_shapes_single_element(self):
xs = np.ones((512, 512, 3))
@@ -66,8 +68,8 @@ def test_return_shapes_single_element(self):
xs = layer(xs)
ys_segmentation_masks = layer(ys_segmentation_masks)
- self.assertEqual(xs.shape, [512, 512, 3])
- self.assertEqual(ys_segmentation_masks.shape, [512, 512, 3])
+ self.assertEqual(xs.shape, (512, 512, 3))
+ self.assertEqual(ys_segmentation_masks.shape, (512, 512, 3))
def test_random_cutout_single_float(self):
self._run_test(0.5, 0.5)
@@ -103,10 +105,10 @@ def test_random_cutout_call_results_one_channel(self):
xs = layer(xs)
# Some pixels should be replaced with fill value
- self.assertTrue(tf.math.reduce_any(xs[0] == patch_value))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == patch_value))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == patch_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == patch_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
def test_random_cutout_call_tiny_image(self):
img_shape = (4, 4, 3)
@@ -127,11 +129,12 @@ def test_random_cutout_call_tiny_image(self):
xs = layer(xs)
# Some pixels should be replaced with fill value
- self.assertTrue(tf.math.reduce_any(xs[0] == fill_value))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == fill_value))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == fill_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == fill_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+ @pytest.mark.tf_only
def test_in_tf_function(self):
xs = tf.cast(
tf.stack(
@@ -156,7 +159,7 @@ def augment(x):
xs = augment(xs)
# Some pixels should be replaced with fill value
- self.assertTrue(tf.math.reduce_any(xs[0] == patch_value))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == patch_value))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == patch_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == patch_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
diff --git a/keras_cv/layers/preprocessing/random_flip_test.py b/keras_cv/layers/preprocessing/random_flip_test.py
index c6f16e0da5..5b130091e3 100644
--- a/keras_cv/layers/preprocessing/random_flip_test.py
+++ b/keras_cv/layers/preprocessing/random_flip_test.py
@@ -17,6 +17,7 @@
import tensorflow as tf
from keras_cv import bounding_box
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.random_flip import HORIZONTAL_AND_VERTICAL
from keras_cv.layers.preprocessing.random_flip import RandomFlip
from keras_cv.tests.test_case import TestCase
@@ -141,9 +142,11 @@ def test_random_flip_unbatched_image(self):
def test_output_dtypes(self):
inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
layer = RandomFlip()
- self.assertAllEqual(layer(inputs).dtype, "float32")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+ )
layer = RandomFlip(dtype="uint8")
- self.assertAllEqual(layer(inputs).dtype, "uint8")
+ self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
def test_augment_bounding_box_batched_input(self):
image = tf.zeros([20, 20, 3])
diff --git a/keras_cv/layers/preprocessing/random_gaussian_blur_test.py b/keras_cv/layers/preprocessing/random_gaussian_blur_test.py
index be917133cf..698c65d885 100644
--- a/keras_cv/layers/preprocessing/random_gaussian_blur_test.py
+++ b/keras_cv/layers/preprocessing/random_gaussian_blur_test.py
@@ -28,12 +28,12 @@ def test_return_shapes(self):
# RGB
xs = np.ones((2, 512, 512, 3))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
# greyscale
xs = np.ones((2, 512, 512, 1))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 512, 512, 1])
+ self.assertEqual(xs.shape, (2, 512, 512, 1))
def test_in_single_image(self):
layer = preprocessing.RandomGaussianBlur(
@@ -47,7 +47,7 @@ def test_in_single_image(self):
)
xs = layer(xs)
- self.assertEqual(xs.shape, [512, 512, 3])
+ self.assertEqual(xs.shape, (512, 512, 3))
# greyscale
xs = tf.cast(
@@ -56,7 +56,7 @@ def test_in_single_image(self):
)
xs = layer(xs)
- self.assertEqual(xs.shape, [512, 512, 1])
+ self.assertEqual(xs.shape, (512, 512, 1))
def test_non_square_images(self):
layer = preprocessing.RandomGaussianBlur(
@@ -66,12 +66,12 @@ def test_non_square_images(self):
# RGB
xs = np.ones((2, 256, 512, 3))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 256, 512, 3])
+ self.assertEqual(xs.shape, (2, 256, 512, 3))
# greyscale
xs = np.ones((2, 256, 512, 1))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 256, 512, 1])
+ self.assertEqual(xs.shape, (2, 256, 512, 1))
def test_single_input_args(self):
layer = preprocessing.RandomGaussianBlur(kernel_size=7, factor=2)
@@ -79,12 +79,12 @@ def test_single_input_args(self):
# RGB
xs = np.ones((2, 512, 512, 3))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
# greyscale
xs = np.ones((2, 512, 512, 1))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 512, 512, 1])
+ self.assertEqual(xs.shape, (2, 512, 512, 1))
def test_numerical(self):
layer = preprocessing.RandomGaussianBlur(
diff --git a/keras_cv/layers/preprocessing/random_hue_test.py b/keras_cv/layers/preprocessing/random_hue_test.py
index cad371f514..9977c7f2ae 100644
--- a/keras_cv/layers/preprocessing/random_hue_test.py
+++ b/keras_cv/layers/preprocessing/random_hue_test.py
@@ -11,10 +11,11 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
-import tensorflow as tf
+import numpy as np
from absl.testing import parameterized
from keras_cv import core
+from keras_cv.backend import ops
from keras_cv.layers import preprocessing
from keras_cv.tests.test_case import TestCase
@@ -22,7 +23,7 @@
class RandomHueTest(TestCase):
def test_preserves_output_shape(self):
image_shape = (4, 8, 8, 3)
- image = tf.random.uniform(shape=image_shape) * 255.0
+ image = np.random.uniform(size=image_shape) * 255.0
layer = preprocessing.RandomHue(factor=(0.3, 0.8), value_range=(0, 255))
output = layer(image)
@@ -32,7 +33,7 @@ def test_preserves_output_shape(self):
def test_adjust_no_op(self):
image_shape = (4, 8, 8, 3)
- image = tf.random.uniform(shape=image_shape) * 255.0
+ image = np.random.uniform(size=image_shape) * 255.0
layer = preprocessing.RandomHue(factor=(0.0, 0.0), value_range=(0, 255))
output = layer(image)
@@ -40,24 +41,24 @@ def test_adjust_no_op(self):
def test_adjust_full_opposite_hue(self):
image_shape = (4, 8, 8, 3)
- image = tf.random.uniform(shape=image_shape) * 255.0
+ image = np.random.uniform(size=image_shape) * 255.0
layer = preprocessing.RandomHue(factor=(1.0, 1.0), value_range=(0, 255))
- output = layer(image)
+ output = ops.convert_to_numpy(layer(image))
- channel_max = tf.math.reduce_max(output, axis=-1)
- channel_min = tf.math.reduce_min(output, axis=-1)
+ channel_max = np.max(output, axis=-1)
+ channel_min = np.min(output, axis=-1)
# Make sure the max and min channel are the same between input and
# output. In the meantime, and channel will swap between each other.
self.assertAllClose(
channel_max,
- tf.math.reduce_max(image, axis=-1),
+ np.max(image, axis=-1),
atol=1e-5,
rtol=1e-5,
)
self.assertAllClose(
channel_min,
- tf.math.reduce_min(image, axis=-1),
+ np.min(image, axis=-1),
atol=1e-5,
rtol=1e-5,
)
@@ -68,7 +69,7 @@ def test_adjust_full_opposite_hue(self):
def test_adjusts_all_values_for_factor(self, factor):
image_shape = (4, 8, 8, 3)
# Value range (0, 100)
- image = tf.random.uniform(shape=image_shape) * 100.0
+ image = np.random.uniform(size=image_shape) * 100.0
layer = preprocessing.RandomHue(
factor=(factor, factor), value_range=(0, 255)
@@ -79,7 +80,7 @@ def test_adjusts_all_values_for_factor(self, factor):
def test_adjustment_for_non_rgb_value_range(self):
image_shape = (4, 8, 8, 3)
# Value range (0, 100)
- image = tf.random.uniform(shape=image_shape) * 100.0
+ image = np.random.uniform(size=image_shape) * 100.0
layer = preprocessing.RandomHue(factor=(0.0, 0.0), value_range=(0, 255))
output = layer(image)
@@ -91,9 +92,7 @@ def test_adjustment_for_non_rgb_value_range(self):
def test_with_uint8(self):
image_shape = (4, 8, 8, 3)
- image = tf.cast(
- tf.random.uniform(shape=image_shape) * 255.0, dtype=tf.uint8
- )
+ image = (np.random.uniform(size=image_shape) * 255.0).astype(np.uint8)
layer = preprocessing.RandomHue(factor=(0.0, 0.0), value_range=(0, 255))
output = layer(image)
diff --git a/keras_cv/layers/preprocessing/random_jpeg_quality_test.py b/keras_cv/layers/preprocessing/random_jpeg_quality_test.py
index 69e9b42ebf..c52a6f2e52 100644
--- a/keras_cv/layers/preprocessing/random_jpeg_quality_test.py
+++ b/keras_cv/layers/preprocessing/random_jpeg_quality_test.py
@@ -26,12 +26,12 @@ def test_return_shapes(self):
# RGB
xs = np.ones((2, 512, 512, 3))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
# greyscale
xs = np.ones((2, 512, 512, 1))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 512, 512, 1])
+ self.assertEqual(xs.shape, (2, 512, 512, 1))
def test_in_single_image(self):
layer = preprocessing.RandomJpegQuality(factor=[0, 100])
@@ -43,7 +43,7 @@ def test_in_single_image(self):
)
xs = layer(xs)
- self.assertEqual(xs.shape, [512, 512, 3])
+ self.assertEqual(xs.shape, (512, 512, 3))
# greyscale
xs = tf.cast(
@@ -52,7 +52,7 @@ def test_in_single_image(self):
)
xs = layer(xs)
- self.assertEqual(xs.shape, [512, 512, 1])
+ self.assertEqual(xs.shape, (512, 512, 1))
def test_non_square_images(self):
layer = preprocessing.RandomJpegQuality(factor=[0, 100])
@@ -60,9 +60,9 @@ def test_non_square_images(self):
# RGB
xs = np.ones((2, 256, 512, 3))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 256, 512, 3])
+ self.assertEqual(xs.shape, (2, 256, 512, 3))
# greyscale
xs = np.ones((2, 256, 512, 1))
xs = layer(xs)
- self.assertEqual(xs.shape, [2, 256, 512, 1])
+ self.assertEqual(xs.shape, (2, 256, 512, 1))
diff --git a/keras_cv/layers/preprocessing/random_rotation_test.py b/keras_cv/layers/preprocessing/random_rotation_test.py
index 4581a1b5f7..0fa6e59704 100644
--- a/keras_cv/layers/preprocessing/random_rotation_test.py
+++ b/keras_cv/layers/preprocessing/random_rotation_test.py
@@ -15,6 +15,7 @@
import tensorflow as tf
from keras_cv import bounding_box
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.random_rotation import RandomRotation
from keras_cv.tests.test_case import TestCase
@@ -83,9 +84,11 @@ def test_augment_bounding_boxes(self):
def test_output_dtypes(self):
inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
layer = RandomRotation(0.5)
- self.assertAllEqual(layer(inputs).dtype, "float32")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+ )
layer = RandomRotation(0.5, dtype="uint8")
- self.assertAllEqual(layer(inputs).dtype, "uint8")
+ self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
def test_ragged_bounding_boxes(self):
input_image = tf.random.uniform((2, 512, 512, 3))
@@ -174,7 +177,9 @@ def test_augment_sparse_segmentation_mask(self):
factor=(0.125, 0.125), segmentation_classes=num_classes
)
outputs = layer(inputs)
- self.assertAllInSet(outputs["segmentation_masks"], [0, 7])
+ self.assertAllInSet(
+ ops.convert_to_numpy(outputs["segmentation_masks"]), [0, 7]
+ )
def test_augment_one_hot_segmentation_mask(self):
num_classes = 8
diff --git a/keras_cv/layers/preprocessing/random_saturation_test.py b/keras_cv/layers/preprocessing/random_saturation_test.py
index 441a2665f3..652ad7c6f9 100644
--- a/keras_cv/layers/preprocessing/random_saturation_test.py
+++ b/keras_cv/layers/preprocessing/random_saturation_test.py
@@ -17,6 +17,7 @@
from keras_cv import core
from keras_cv.backend import keras
+from keras_cv.backend import ops
from keras_cv.layers import preprocessing
from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
BaseImageAugmentationLayer,
@@ -133,9 +134,9 @@ def test_adjust_to_grayscale(self):
image = tf.random.uniform(shape=image_shape) * 255.0
layer = preprocessing.RandomSaturation(factor=(0.0, 0.0))
- output = layer(image)
+ output = ops.convert_to_numpy(layer(image))
- channel_mean = tf.math.reduce_mean(output, axis=-1)
+ channel_mean = np.mean(output, axis=-1)
channel_values = tf.unstack(output, axis=-1)
# Make sure all the pixel has the same value among the channel dim,
# which is a fully gray RGB.
@@ -149,9 +150,9 @@ def test_adjust_to_full_saturation(self):
image = tf.random.uniform(shape=image_shape) * 255.0
layer = preprocessing.RandomSaturation(factor=(1.0, 1.0))
- output = layer(image)
+ output = ops.convert_to_numpy(layer(image))
- channel_mean = tf.math.reduce_min(output, axis=-1)
+ channel_mean = np.min(output, axis=-1)
# Make sure at least one of the channel is 0.0 (fully saturated image)
self.assertAllClose(channel_mean, np.zeros((4, 8, 8)))
diff --git a/keras_cv/layers/preprocessing/random_shear_test.py b/keras_cv/layers/preprocessing/random_shear_test.py
index cb9e61060d..0d2ef32ad8 100644
--- a/keras_cv/layers/preprocessing/random_shear_test.py
+++ b/keras_cv/layers/preprocessing/random_shear_test.py
@@ -11,10 +11,12 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
+import numpy as np
import pytest
import tensorflow as tf
from keras_cv import bounding_box
+from keras_cv.backend import ops
from keras_cv.layers import preprocessing
from keras_cv.tests.test_case import TestCase
@@ -43,18 +45,22 @@ def test_aggressive_shear_fills_at_least_some_pixels(self):
ys_segmentation_masks = layer(ys_segmentation_masks)
# Some pixels should be replaced with fill value
- self.assertTrue(tf.math.reduce_any(xs[0] == fill_value))
- self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
- self.assertTrue(tf.math.reduce_any(xs[1] == fill_value))
- self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == fill_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == fill_value))
+ self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
self.assertTrue(
- tf.math.reduce_any(ys_segmentation_masks[0] == fill_value)
+ np.any(ops.convert_to_numpy(ys_segmentation_masks[0]) == fill_value)
)
- self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[0] == 2.0))
self.assertTrue(
- tf.math.reduce_any(ys_segmentation_masks[1] == fill_value)
+ np.any(ops.convert_to_numpy(ys_segmentation_masks[0]) == 2.0)
+ )
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(ys_segmentation_masks[1]) == fill_value)
+ )
+ self.assertTrue(
+ np.any(ops.convert_to_numpy(ys_segmentation_masks[1]) == 1.0)
)
- self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[1] == 1.0))
def test_return_shapes(self):
"""test return dict keys and value pairs"""
@@ -96,11 +102,11 @@ def test_return_shapes(self):
outputs["segmentation_masks"],
)
ys_bounding_boxes = bounding_box.to_dense(ys_bounding_boxes)
- self.assertEqual(xs.shape, [2, 512, 512, 3])
- self.assertEqual(ys_labels.shape, [2, 10])
- self.assertEqual(ys_bounding_boxes["boxes"].shape, [2, 3, 4])
- self.assertEqual(ys_bounding_boxes["classes"].shape, [2, 3])
- self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+ self.assertEqual(xs.shape, (2, 512, 512, 3))
+ self.assertEqual(ys_labels.shape, (2, 10))
+ self.assertEqual(ys_bounding_boxes["boxes"].shape, (2, 3, 4))
+ self.assertEqual(ys_bounding_boxes["classes"].shape, (2, 3))
+ self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
def test_single_image_input(self):
"""test for single image input"""
@@ -112,7 +118,7 @@ def test_single_image_input(self):
fill_mode="constant",
)
outputs = layer(inputs)
- self.assertEqual(outputs["images"].shape, [512, 512, 3])
+ self.assertEqual(outputs["images"].shape, (512, 512, 3))
@pytest.mark.skip(reason="Flaky")
def test_area(self):
@@ -155,6 +161,7 @@ def test_area(self):
)
self.assertTrue(tf.math.reduce_all(new_area > old_area))
+ @pytest.mark.tf_only
def test_in_tf_function(self):
"""test for class works with tf function"""
xs = tf.cast(
diff --git a/keras_cv/layers/preprocessing/random_translation_test.py b/keras_cv/layers/preprocessing/random_translation_test.py
index ed98920560..6a2f36a556 100644
--- a/keras_cv/layers/preprocessing/random_translation_test.py
+++ b/keras_cv/layers/preprocessing/random_translation_test.py
@@ -15,6 +15,7 @@
import numpy as np
import tensorflow as tf
+from keras_cv.backend import ops
from keras_cv.layers import preprocessing
from keras_cv.tests.test_case import TestCase
@@ -224,6 +225,8 @@ def test_unbatched_image(self):
def test_output_dtypes(self):
inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
layer = preprocessing.RandomTranslation(0.5, 0.6)
- self.assertAllEqual(layer(inputs).dtype, "float32")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+ )
layer = preprocessing.RandomTranslation(0.5, 0.6, dtype="uint8")
- self.assertAllEqual(layer(inputs).dtype, "uint8")
+ self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
diff --git a/keras_cv/layers/preprocessing/random_zoom_test.py b/keras_cv/layers/preprocessing/random_zoom_test.py
index 0fdcf6eec3..875402750f 100644
--- a/keras_cv/layers/preprocessing/random_zoom_test.py
+++ b/keras_cv/layers/preprocessing/random_zoom_test.py
@@ -16,6 +16,7 @@
import tensorflow as tf
from absl.testing import parameterized
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.random_zoom import RandomZoom
from keras_cv.tests.test_case import TestCase
@@ -163,6 +164,8 @@ def test_unbatched_image(self):
def test_output_dtypes(self):
inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
layer = RandomZoom(0.5, 0.5)
- self.assertAllEqual(layer(inputs).dtype, "float32")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+ )
layer = RandomZoom(0.5, 0.5, dtype="uint8")
- self.assertAllEqual(layer(inputs).dtype, "uint8")
+ self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
diff --git a/keras_cv/layers/preprocessing/rescaling_test.py b/keras_cv/layers/preprocessing/rescaling_test.py
index 0801c227d0..1ef2ed6c2f 100644
--- a/keras_cv/layers/preprocessing/rescaling_test.py
+++ b/keras_cv/layers/preprocessing/rescaling_test.py
@@ -14,6 +14,7 @@
import numpy as np
import tensorflow as tf
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.rescaling import Rescaling
from keras_cv.tests.test_case import TestCase
@@ -23,14 +24,17 @@ def test_rescaling_correctness_float(self):
layer = Rescaling(scale=1.0 / 127.5, offset=-1.0)
inputs = tf.random.uniform((2, 4, 5, 3))
outputs = layer(inputs)
- self.assertAllClose(outputs.numpy(), inputs.numpy() * (1.0 / 127.5) - 1)
+ self.assertAllClose(outputs, inputs * (1.0 / 127.5) - 1)
def test_rescaling_correctness_int(self):
layer = Rescaling(scale=1.0 / 127.5, offset=-1)
inputs = tf.random.uniform((2, 4, 5, 3), 0, 100, dtype="int32")
outputs = layer(inputs)
+ outputs = ops.convert_to_numpy(outputs)
self.assertEqual(outputs.dtype.name, "float32")
- self.assertAllClose(outputs.numpy(), inputs.numpy() * (1.0 / 127.5) - 1)
+ self.assertAllClose(
+ outputs, ops.convert_to_numpy(inputs) * (1.0 / 127.5) - 1
+ )
def test_config_with_custom_name(self):
layer = Rescaling(0.5, name="rescaling")
@@ -42,11 +46,15 @@ def test_unbatched_image(self):
layer = Rescaling(scale=1.0 / 127.5, offset=-1)
inputs = tf.random.uniform((4, 5, 3))
outputs = layer(inputs)
- self.assertAllClose(outputs.numpy(), inputs.numpy() * (1.0 / 127.5) - 1)
+ self.assertAllClose(outputs, inputs * (1.0 / 127.5) - 1)
def test_output_dtypes(self):
inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
layer = Rescaling(0.5)
- self.assertAllEqual(layer(inputs).dtype, "float32")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype.name, "float32"
+ )
layer = Rescaling(0.5, dtype="uint8")
- self.assertAllEqual(layer(inputs).dtype, "uint8")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype.name, "uint8"
+ )
diff --git a/keras_cv/layers/preprocessing/resizing_test.py b/keras_cv/layers/preprocessing/resizing_test.py
index 62267368e6..c094dee55a 100644
--- a/keras_cv/layers/preprocessing/resizing_test.py
+++ b/keras_cv/layers/preprocessing/resizing_test.py
@@ -17,6 +17,7 @@
from absl.testing import parameterized
from keras_cv import layers as cv_layers
+from keras_cv.backend import ops
from keras_cv.backend.config import keras_3
from keras_cv.tests.test_case import TestCase
@@ -183,9 +184,11 @@ def test_ragged_image(self, crop_to_aspect_ratio):
def test_output_dtypes(self):
inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
layer = cv_layers.Resizing(2, 2)
- self.assertAllEqual(layer(inputs).dtype, "float32")
+ self.assertAllEqual(
+ ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+ )
layer = cv_layers.Resizing(2, 2, dtype="uint8")
- self.assertAllEqual(layer(inputs).dtype, "uint8")
+ self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
@parameterized.named_parameters(
("batch_crop_to_aspect_ratio", True, False, True),
diff --git a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
index 8d3dacfa98..3d9fc8e52a 100644
--- a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
+++ b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
@@ -12,19 +12,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-import keras
import tensorflow as tf
-
-if hasattr(keras, "src"):
- keras_backend = keras.src.backend
-else:
- keras_backend = keras.backend
+import tree
from keras_cv import bounding_box
from keras_cv.api_export import keras_cv_export
from keras_cv.backend import keras
+from keras_cv.backend import ops
from keras_cv.backend import scope
-from keras_cv.backend.config import keras_3
from keras_cv.utils import preprocessing
H_AXIS = -3
@@ -42,15 +37,8 @@
USE_TARGETS = "use_targets"
-base_class = (
- keras.src.layers.preprocessing.tf_data_layer.TFDataLayer
- if keras_3()
- else keras.layers.Layer
-)
-
-
@keras_cv_export("keras_cv.layers.VectorizedBaseImageAugmentationLayer")
-class VectorizedBaseImageAugmentationLayer(base_class):
+class VectorizedBaseImageAugmentationLayer(keras.layers.Layer):
"""Abstract base layer for vectorized image augmentation.
This layer contains base functionalities for preprocessing layers which
@@ -422,6 +410,19 @@ def _batch_augment(self, inputs):
return result
def call(self, inputs):
+ # try to convert a given backend native tensor to TensorFlow tensor
+ # before passing it over to TFDataScope
+ contains_ragged = lambda y: any(
+ tree.map_structure(
+ lambda x: isinstance(x, (tf.RaggedTensor, tf.SparseTensor)),
+ tree.flatten(y),
+ )
+ )
+ inputs_contain_ragged = contains_ragged(inputs)
+ if not inputs_contain_ragged:
+ inputs = tree.map_structure(
+ lambda x: tf.convert_to_tensor(x), inputs
+ )
with scope.TFDataScope():
inputs = self._ensure_inputs_are_compute_dtype(inputs)
inputs, metadata = self._format_inputs(inputs)
@@ -436,7 +437,20 @@ def call(self, inputs):
"rank 3 (HWC) or 4D (NHWC) tensors. Got shape: "
f"{images.shape}"
)
- return outputs
+ # convert the outputs to backend native tensors if none of them
+ # contain RaggedTensors. Note that if the user passed in Raggeds
+ # but the outputs are dense, we still don't want to convert to
+ # backend native tensors. This is to avoid breaking TF data
+ # pipelines that can't easily be ported to become backend
+ # agnostic.
+ if not inputs_contain_ragged and not contains_ragged(outputs):
+ outputs = tree.map_structure(
+ # some layers return None, handle that case when
+ # converting to tensors
+ lambda x: ops.convert_to_tensor(x) if x is not None else x,
+ outputs,
+ )
+ return outputs
def _format_inputs(self, inputs):
metadata = {IS_DICT: True, USE_TARGETS: False}
diff --git a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer_test.py b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer_test.py
index aeb6d4d3d8..3ebdfdb820 100644
--- a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer_test.py
+++ b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer_test.py
@@ -16,6 +16,8 @@
import tensorflow as tf
from keras_cv import bounding_box
+from keras_cv.backend import keras
+from keras_cv.backend import ops
from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import ( # noqa: E501
VectorizedBaseImageAugmentationLayer,
)
@@ -208,11 +210,11 @@ def test_augment_dict_return_type(self):
def test_augment_casts_dtypes(self):
add_layer = VectorizedRandomAddLayer(fixed_value=2.0)
- images = tf.ones((2, 8, 8, 3), dtype="uint8")
+ images = np.ones((2, 8, 8, 3), dtype="uint8")
output = add_layer(images)
self.assertAllClose(
- tf.ones((2, 8, 8, 3), dtype="float32") * 3.0, output
+ np.ones((2, 8, 8, 3), dtype="float32") * 3.0, output
)
def test_augment_batch_images(self):
@@ -220,7 +222,7 @@ def test_augment_batch_images(self):
images = np.random.random(size=(2, 8, 8, 3)).astype("float32")
output = add_layer(images)
- diff = output - images
+ diff = ops.convert_to_numpy(output) - images
# Make sure the first image and second image get different augmentation
self.assertNotAllClose(diff[0], diff[1])
@@ -248,8 +250,8 @@ def test_augment_batch_images_and_targets(self):
targets = np.random.random(size=(2, 1)).astype("float32")
output = add_layer({"images": images, "targets": targets})
- image_diff = output["images"] - images
- label_diff = output["targets"] - targets
+ image_diff = ops.convert_to_numpy(output["images"]) - images
+ label_diff = ops.convert_to_numpy(output["targets"]) - targets
# Make sure the first image and second image get different augmentation
self.assertNotAllClose(image_diff[0], image_diff[1])
self.assertNotAllClose(label_diff[0], label_diff[1])
@@ -357,6 +359,13 @@ def test_augment_batch_image_and_localization_data(self):
segmentation_mask_diff[0], segmentation_mask_diff[1]
)
+ # the test finishes here for the non-tensorflow backends.
+ if (
+ getattr(keras.config, "backend", lambda: "tensorflow")()
+ != "tensorflow"
+ ):
+ return
+
@tf.function
def in_tf_function(inputs):
return add_layer(inputs)
@@ -383,6 +392,7 @@ def in_tf_function(inputs):
segmentation_mask_diff[0], segmentation_mask_diff[1]
)
+ @pytest.mark.tf_only
def test_augment_all_data_in_tf_function(self):
add_layer = VectorizedRandomAddLayer()
images = np.random.random(size=(2, 8, 8, 3)).astype("float32")
@@ -443,11 +453,11 @@ def test_augment_unbatched_all_data(self):
self.assertAllClose(output["keypoints"], keypoints + 2.0)
self.assertAllClose(
output["bounding_boxes"]["boxes"],
- tf.squeeze(bounding_boxes["boxes"]) + 2.0,
+ np.squeeze(bounding_boxes["boxes"]) + 2.0,
)
self.assertAllClose(
output["bounding_boxes"]["classes"],
- tf.squeeze(bounding_boxes["classes"]) + 2.0,
+ np.squeeze(bounding_boxes["classes"]) + 2.0,
)
self.assertAllClose(
output["segmentation_masks"], segmentation_masks + 2.0
@@ -478,7 +488,6 @@ def test_augment_all_data_for_assertion(self):
# assertion is at VectorizedAssertionLayer's methods
- @pytest.mark.skip(reason="disable temporarily")
def test_augment_all_data_with_ragged_images_for_assertion(self):
images = tf.ragged.stack(
[
@@ -497,15 +506,6 @@ def test_augment_all_data_with_ragged_images_for_assertion(self):
segmentation_masks = tf.random.uniform(shape=(2, 8, 8, 1))
assertion_layer = VectorizedAssertionLayer()
- print(
- {
- "images": type(images),
- "labels": type(labels),
- "bounding_boxes": type(bounding_boxes),
- "keypoints": type(keypoints),
- "segmentation_masks": type(segmentation_masks),
- }
- )
_ = assertion_layer(
{
"images": images,