Vectorize Random Shear.

[sebastian-sz]

Short Description
Random Shear can be vectorized, similar to other layers. It seems the performance benefit is quite large for layers that use ImageProjectiveTransformV3 (similar to: #1480).

EDIT: the below graphs were generated with slight error that favoured Vectorized implementation. These are updated graphs, without the error.
Vectorized implementation is faster than the original:

Known Issues

• augment_bounding_boxes is quite challenging. I did vectorize the existing code, but it doesn't work if the bounding box is a tf.RaggedTensor as it's impossible to tf.transpose ragged tensors (transposing bounding boxes is necessary for matmul).
  This breaks the examples/layers/preprocessing/bounding_box/random_shear_demo.py

• No XLA for ImageProjetiveTransformV3.

Let me know what you think!

[james77777778]

Hi @sebastian-sz
I face the same problem about tf.RaggedTensor in bounding_boxes as well and the unbatched inputs with bounding_boxes #1512 (the unit test of RandomShear doesn't catch this)

There is a workaround for augment_bounding_boxes by calling

    def augment_bounding_boxes(
        self, bounding_boxes, transformations, images, **kwargs
    ):
        ...
        bounding_boxes = bounding_box.to_dense(bounding_boxes)  # <--

to transform tf.RaggedTensor to dense tensor.

Side effect: it will fulfill dummy boxes to meet the max number of the boxes in the batch

VectorizedBaseImageAugmentationLayer will automatically turn it back at Line 317

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

Lines 310 to 318 in cd12204

	if bounding_boxes is not None:
	bounding_boxes = self.augment_bounding_boxes(
	bounding_boxes,
	transformations=transformations,
	labels=labels,
	images=images,
	)
	bounding_boxes = bounding_box.to_ragged(bounding_boxes)
	result[BOUNDING_BOXES] = bounding_boxes

P.S.
BaseImageAugmentationLayer exactly implements this behavior

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

Lines 450 to 461 in cd12204

	if bounding_boxes is not None:
	bounding_boxes = bounding_box.to_dense(bounding_boxes)
	bounding_boxes = self.augment_bounding_boxes(
	bounding_boxes,
	transformation=transformation,
	label=label,
	image=raw_image,
	)
	bounding_boxes = bounding_box.to_ragged(bounding_boxes)
	result[BOUNDING_BOXES] = bounding_boxes

I'm not sure if we are encouraged to augment bounding_boxes in tf.RaggedTensor type or we can convert it to dense tensor first?

Kindly ping @LukeWood @ianstenbit

[LukeWood]

Awesome! Go ahead and introduce this @sebastian-sz

[LukeWood]

augment_bounding_boxes is quite challenging. I did vectorize the existing code, but it doesn't work if the bounding box is a tf.RaggedTensor as it's impossible to tf.transpose ragged tensors (transposing bounding boxes is necessary for matmul).

Regarding this, we can always just use a tf.map_function inside of the augment_bounding_box_batch() batched version. For bounding boxes the performance overhead should be marginal.

[sebastian-sz]

@james77777778 thank you for the suggestion - it works with Ragged!

@LukeWood there seems to be a different issue on master branch with the non-vectorized implementation. These are some of the outputs that I'm getting from examples/layers/preprocessing/bounding_box/random_shear_demo.py:

The boxes look quite off in some of those examples - are we sure that the current algorithm to augment bounding boxes is correct?

[LukeWood]

Looks like it was just a bug in the demo: https://github.com/keras-team/keras-cv/pull/1534/files