Bug: VectorizedBaseImageAugmentationLayer fails to process unbatched bounding_boxes

[james77777778]

I find that current implementation of VectorizedBaseImageAugmentationLayer fails to correctly process unbatched bounding_boxes

Here is a standalone script:

import tensorflow as tf
from keras_cv.layers import preprocessing

if __name__ == "__main__":
    # construct unbatched input (images, bounding_boxes)
    images = tf.zeros([8, 8, 3])
    bounding_boxes = {
        "boxes": tf.ragged.constant(
            [[0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]],
            dtype=tf.float32,
        ),
        "classes": tf.RaggedTensor.from_tensor(tf.zeros([2, 1])),
    }
    inputs = {"images": images, "bounding_boxes": bounding_boxes}
    # any vectorized layers, take RandomZoom as example
    layer = preprocessing.RandomZoom(0.5, 0.5)
    outputs = layer(inputs, training=True)  # raises ValueError

Reason: VectorizedBaseImageAugmentationLayer tries to tf.expand_dims (tf.squeeze) all values of inputs (outputs) including bounding_boxes which is a dict

keras-cv/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py

Lines 368 to 370 in 4fd3a84

	if inputs["images"].shape.rank == 3:
	for key in list(inputs.keys()):
	inputs[key] = tf.expand_dims(inputs[key], axis=0)

keras-cv/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py

Lines 392 to 394 in 4fd3a84

	if not metadata[BATCHED]:
	for key in list(output.keys()):
	output[key] = tf.squeeze(output[key], axis=0)

The possible solution:

            # _format_inputs
            for key in list(inputs.keys()):
                if key == BOUNDING_BOXES:
                    inputs[BOUNDING_BOXES]["boxes"] = tf.expand_dims(
                        inputs[BOUNDING_BOXES]["boxes"], axis=0
                    )
                    inputs[BOUNDING_BOXES]["classes"] = tf.expand_dims(
                        inputs[BOUNDING_BOXES]["classes"], axis=0
                    )
                else:
                    inputs[key] = tf.expand_dims(inputs[key], axis=0)
                ...
            # _format_output
            for key in list(output.keys()):
                if key == BOUNDING_BOXES:
                    output[BOUNDING_BOXES]["boxes"] = tf.squeeze(
                        output[BOUNDING_BOXES]["boxes"], axis=0
                    )
                    output[BOUNDING_BOXES]["classes"] = tf.squeeze(
                        output[BOUNDING_BOXES]["classes"], axis=0
                    )
                else:
                    output[key] = tf.squeeze(output[key], axis=0)

It needs to be fixed if we want to implement augment_bounding_boxes for vectorized layers.

I can open the PR once approved.

[james77777778]

This bug should influence #1439
@soma2000-lang

The unit test should fail at VectorizedBaseImageAugmentationLayer part

keras-cv/keras_cv/layers/preprocessing/random_crop_and_resize_test.py

Lines 191 to 205 in 4fd3a84

	def test_augment_bounding_box_single(self):
	image = tf.zeros([20, 20, 3])
	boxes = {
	"boxes": tf.convert_to_tensor([[0, 0, 1, 1]]),
	"classes": tf.convert_to_tensor([0]),
	}
	input = {"images": image, "bounding_boxes": boxes}
	layer = preprocessing.RandomCropAndResize(
	target_size=(10, 10),
	crop_area_factor=(0.5**2, 0.5**2),
	aspect_ratio_factor=(1.0, 1.0),
	bounding_box_format="rel_xyxy",
	)
	output = layer(input, training=True)

[soma2000-lang]

@james77777778 thanks for flagging ,yes some of the tests are failing due to this.

[LukeWood]

Go ahead @james77777778 - thanks for filing!