KerasCV augmentations throws error in `tf.data` with Keras V3

[awsaf49]

Current Behavior:

KerasCV augmentations in tf.data throw error. After digging deeper, I noticed this error occurs with Keras V3 but not with KerasCore.

Here is the error message

NotImplementedError: in user code:

    File "/tmp/ipykernel_26/3817339408.py", line 22, in apply_augmentation  *
        return (augmenter(images), labels)
    File "/opt/conda/lib/python3.10/site-packages/keras_cv/src/layers/augmenter.py", line 44, in __call__  *
        inputs = layer(inputs)
    File "/opt/conda/lib/python3.10/site-packages/keras/src/utils/traceback_utils.py", line 123, in error_handler  **
        raise e.with_traceback(filtered_tb) from None
    File "/opt/conda/lib/python3.10/site-packages/keras_cv/src/layers/preprocessing/vectorized_base_image_augmentation_layer.py", line 447, in call
        outputs = tree.map_structure(
    File "/opt/conda/lib/python3.10/site-packages/tree/__init__.py", line 435, in map_structure
        [func(*args) for args in zip(*map(flatten, structures))])
    File "/opt/conda/lib/python3.10/site-packages/tree/__init__.py", line 435, in <listcomp>
        [func(*args) for args in zip(*map(flatten, structures))])
    File "/opt/conda/lib/python3.10/site-packages/keras_cv/src/layers/preprocessing/vectorized_base_image_augmentation_layer.py", line 450, in <lambda>
        lambda x: ops.convert_to_tensor(x) if x is not None else x,
    File "/opt/conda/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 2134, in array
        leaves = [_convert_to_array_if_dtype_fails(leaf) for leaf in leaves]
    File "/opt/conda/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 2134, in <listcomp>
        leaves = [_convert_to_array_if_dtype_fails(leaf) for leaf in leaves]
    File "/opt/conda/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 2177, in _convert_to_array_if_dtype_fails
        return np.asarray(x)

    NotImplementedError: Exception encountered when calling RandomFlip.call().
    
    Cannot convert a symbolic tf.Tensor (SelectV2_1:0) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported.
    
    Arguments received by RandomFlip.call():
      • inputs=<tf.Tensor 'args_0:0' shape=(None, 256, 256, 3) dtype=float32>

Expected Behavior:

Steps To Reproduce:

We can reproduce this error using this Kaggle Notebook. In Version 5 of this notebook, KerasCV augmentations work seamlessly with tf.data that uses KerasCore. However, in Version 6, when I switch to Keras V3, it results in an error. Notably, in all cases, I'm using the jax backend.

Version:

• 0.8.1

Anything else:

Initially, I thought this issue was similar to #1942. However, when I used the workaround for this issue (use List instead of keras_cv.layers.Augmenter), it did not work for me.

[awsaf49]

I have also tried Semantic Segmentation with KerasCV Example code by @divyashreepathihalli and @ianstenbit with jax backend, it also throw the same error. Here is the Colab Notebook

[awsaf49]

cc: @martin-gorner

[martin-gorner]

More specifically @awsaf49 was only able to use MixUp augmentation. The CutOut augmentation could not be used.

[awsaf49]

More specifically @awsaf49 was only able to use MixUp augmentation. The CutOut augmentation could not be used.

Yes, @martin-gorner is right. To be more precise, augmentations built upon BaseImageAugmentationLayer, such as MixUp, CutMix, etc., work, while those built upon VectorizedBaseImageAugmentationLayer, such as RandomFlip or RandomCutOut, do not work."

[divyashreepathihalli]

Thanks for filing the issue, we will take a look.

[sampathweb]

I will work on this today.

[sampathweb]

Thanks for reporting this issue. The issue is now fixed. Please use keras-cv-nightly until the next release for the fix.

[sampathweb]

@awsaf49 - A new patch release was done that includes this update. You could use the latest keras-cv now that has this fix.

[awsaf49]

@sampathweb Great work! Thanks.