Fix format and Update Vectorized Base to fix tf.data with non-TF backends

benchmarks/vectorized_randomly_zoomed_crop.py

@@ -249,10 +249,10 @@ def from_config(cls, config):
                 config["zoom_factor"]
             )
         if isinstance(config["aspect_ratio_factor"], dict):
-            config[
-                "aspect_ratio_factor"
-            ] = keras.utils.deserialize_keras_object(
-                config["aspect_ratio_factor"]
+            config["aspect_ratio_factor"] = (
+                keras.utils.deserialize_keras_object(
+                    config["aspect_ratio_factor"]
+                )
             )
         return cls(**config)
 

keras_cv/layers/preprocessing/base_image_augmentation_layer.py

@@ -236,15 +236,15 @@ def _compute_output_signature(self, inputs):
         bounding_boxes = inputs.get(BOUNDING_BOXES, None)
 
         if bounding_boxes is not None:
-            fn_output_signature[
-                BOUNDING_BOXES
-            ] = self._compute_bounding_box_signature(bounding_boxes)
+            fn_output_signature[BOUNDING_BOXES] = (
+                self._compute_bounding_box_signature(bounding_boxes)
+            )
 
         segmentation_masks = inputs.get(SEGMENTATION_MASKS, None)
         if segmentation_masks is not None:
-            fn_output_signature[
-                SEGMENTATION_MASKS
-            ] = self.compute_image_signature(segmentation_masks)
+            fn_output_signature[SEGMENTATION_MASKS] = (
+                self.compute_image_signature(segmentation_masks)
+            )
 
         keypoints = inputs.get(KEYPOINTS, None)
         if keypoints is not None:

keras_cv/layers/preprocessing/random_crop_and_resize.py

@@ -272,10 +272,10 @@ def from_config(cls, config):
                 config["crop_area_factor"]
             )
         if isinstance(config["aspect_ratio_factor"], dict):
-            config[
-                "aspect_ratio_factor"
-            ] = keras.utils.deserialize_keras_object(
-                config["aspect_ratio_factor"]
+            config["aspect_ratio_factor"] = (
+                keras.utils.deserialize_keras_object(
+                    config["aspect_ratio_factor"]
+                )
             )
         return cls(**config)
 

keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py

@@ -17,6 +17,7 @@
 
 from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import config
 from keras_cv.backend import keras
 from keras_cv.backend import ops
 from keras_cv.backend import scope
@@ -412,14 +413,16 @@ def _batch_augment(self, inputs):
     def call(self, inputs):
         # try to convert a given backend native tensor to TensorFlow tensor
         # before passing it over to TFDataScope
+        is_tf_backend = config.backend() == "tensorflow"
+        is_in_tf_graph = not tf.executing_eagerly()
         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:
+        if not is_tf_backend and not inputs_contain_ragged:
             inputs = tree.map_structure(
                 lambda x: tf.convert_to_tensor(x), inputs
             )
@@ -443,13 +446,14 @@ def call(self, inputs):
         # 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,
-            )
+        if not is_tf_backend and not is_in_tf_graph:
+            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):

keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer_test.py

@@ -549,3 +549,15 @@ def test_converts_ragged_to_dense_segmentation_masks(self):
             {"images": images, "segmentation_masks": segmentation_masks}
         )
         self.assertTrue(isinstance(result["segmentation_masks"], tf.Tensor))
+
+    def test_in_tf_data_pipeline(self):
+        images = np.random.randn(4, 100, 100, 3).astype("float32")
+        train_ds = tf.data.Dataset.from_tensor_slices(images)
+        train_ds = train_ds.map(lambda x: {"images": x})
+        train_ds = train_ds.map(
+            VectorizedRandomAddLayer(fixed_value=2.0)
+        ).batch(4)
+        for output in train_ds.take(1):
+            pass
+        self.assertTrue(isinstance(output["images"], tf.Tensor))
+        self.assertAllClose(output["images"], images + 2.0)

keras_cv/layers/regularization/squeeze_excite.py

@@ -118,10 +118,10 @@ def get_config(self):
     @classmethod
     def from_config(cls, config):
         if isinstance(config["squeeze_activation"], dict):
-            config[
-                "squeeze_activation"
-            ] = keras.saving.deserialize_keras_object(
-                config["squeeze_activation"]
+            config["squeeze_activation"] = (
+                keras.saving.deserialize_keras_object(
+                    config["squeeze_activation"]
+                )
             )
         if isinstance(config["excite_activation"], dict):
             config["excite_activation"] = keras.saving.deserialize_keras_object(

keras_cv/layers/vit_det_layers.py

@@ -430,9 +430,9 @@ def __init__(
             key_dim=self.project_dim // self.num_heads,
             use_bias=use_bias,
             use_rel_pos=use_rel_pos,
-            input_size=input_size
-            if window_size == 0
-            else (window_size, window_size),
+            input_size=(
+                input_size if window_size == 0 else (window_size, window_size)
+            ),
         )
         self.mlp_block = MLP(
             mlp_dim,

keras_cv/metrics/object_detection/box_coco_metrics.py

@@ -212,9 +212,9 @@ def result_fn(self, force=False):
             )
             result = {}
             for i, key in enumerate(METRIC_NAMES):
-                result[
-                    self.name_prefix() + METRIC_MAPPING[key]
-                ] = py_func_result[i]
+                result[self.name_prefix() + METRIC_MAPPING[key]] = (
+                    py_func_result[i]
+                )
             return result
 
         obj.result = types.MethodType(result_fn, obj)

keras_cv/models/backbones/densenet/densenet_backbone.py

@@ -119,9 +119,9 @@ def __init__(
             name=f"conv{len(stackwise_num_repeats) + 1}",
         )
 
-        pyramid_level_inputs[
-            f"P{len(stackwise_num_repeats) + 1}"
-        ] = utils.get_tensor_input_name(x)
+        pyramid_level_inputs[f"P{len(stackwise_num_repeats) + 1}"] = (
+            utils.get_tensor_input_name(x)
+        )
         x = keras.layers.BatchNormalization(
             axis=BN_AXIS, epsilon=BN_EPSILON, name="bn"
         )(x)

keras_cv/models/backbones/resnet_v1/resnet_v1_backbone.py

@@ -130,9 +130,9 @@ def __init__(
                 first_shortcut=(block_type == "block" or stack_index > 0),
                 name=f"v2_stack_{stack_index}",
             )
-            pyramid_level_inputs[
-                f"P{stack_index + 2}"
-            ] = utils.get_tensor_input_name(x)
+            pyramid_level_inputs[f"P{stack_index + 2}"] = (
+                utils.get_tensor_input_name(x)
+            )
 
         # Create model.
         super().__init__(inputs=inputs, outputs=x, **kwargs)

keras_cv/models/backbones/resnet_v2/resnet_v2_backbone.py

@@ -136,9 +136,9 @@ def __init__(
                 first_shortcut=(block_type == "block" or stack_index > 0),
                 name=f"v2_stack_{stack_index}",
             )
-            pyramid_level_inputs[
-                f"P{stack_index + 2}"
-            ] = utils.get_tensor_input_name(x)
+            pyramid_level_inputs[f"P{stack_index + 2}"] = (
+                utils.get_tensor_input_name(x)
+            )
 
         x = keras.layers.BatchNormalization(
             axis=BN_AXIS, epsilon=BN_EPSILON, name="post_bn"

keras_cv/models/backbones/vit_det/vit_det_backbone.py

@@ -144,9 +144,9 @@ def __init__(
                 num_heads=num_heads,
                 use_bias=use_bias,
                 use_rel_pos=use_rel_pos,
-                window_size=window_size
-                if i not in global_attention_indices
-                else 0,
+                window_size=(
+                    window_size if i not in global_attention_indices else 0
+                ),
                 input_size=(img_size // patch_size, img_size // patch_size),
             )(x)
         x = keras.models.Sequential(

keras_cv/models/legacy/darknet.py

@@ -76,7 +76,6 @@
 
 @keras.utils.register_keras_serializable(package="keras_cv.models")
 class DarkNet(keras.Model):
-
     """Represents the DarkNet architecture.
 
     The DarkNet architecture is commonly used for detection tasks. It is

keras_cv/models/legacy/mlp_mixer.py

@@ -143,7 +143,6 @@ def apply_mixer_block(x, tokens_mlp_dim, channels_mlp_dim, name=None):
 
 @keras.utils.register_keras_serializable(package="keras_cv.models")
 class MLPMixer(keras.Model):
-
     """Instantiates the MLP Mixer architecture.
 
     Args:

keras_cv/models/object_detection/yolo_v8/yolo_v8_backbone.py

@@ -178,9 +178,9 @@ def __init__(
                     activation=activation,
                     name=f"{stack_name}_spp_fast",
                 )
-            pyramid_level_inputs[
-                f"P{stack_id + 2}"
-            ] = utils.get_tensor_input_name(x)
+            pyramid_level_inputs[f"P{stack_id + 2}"] = (
+                utils.get_tensor_input_name(x)
+            )
 
         super().__init__(inputs=inputs, outputs=x, **kwargs)
         self.pyramid_level_inputs = pyramid_level_inputs

keras_cv/models/object_detection/yolo_v8/yolo_v8_detector.py

@@ -663,9 +663,9 @@ def from_config(cls, config):
         if prediction_decoder is not None and isinstance(
             prediction_decoder, dict
         ):
-            config[
-                "prediction_decoder"
-            ] = keras.saving.deserialize_keras_object(prediction_decoder)
+            config["prediction_decoder"] = (
+                keras.saving.deserialize_keras_object(prediction_decoder)
+            )
         return cls(**config)
 
     @classproperty

keras_cv/models/stable_diffusion/noise_scheduler.py

@@ -54,9 +54,7 @@ def __init__(
         elif beta_schedule == "scaled_linear":
             # this schedule is very specific to the latent diffusion model.
             self.betas = (
-                ops.linspace(
-                    beta_start**0.5, beta_end**0.5, train_timesteps
-                )
+                ops.linspace(beta_start**0.5, beta_end**0.5, train_timesteps)
                 ** 2
             )
         else: