Repository: keras-team/keras-cv

Issue:
Error in `bounding_box.converts._xyxy_to_rel_xyxy`:
Currently this breaks when operating on ragged images.  This can be see in the `examples/layers/random_shear_demo.py` demo.

More info available here: https://github.com/keras-team/keras-cv/pull/1537#issuecomment-1479032369

Perhaps we can unit test this behavior for all of our converters.

Discussion:

Patch:
diff --git a/keras_cv/bounding_box/converters.py b/keras_cv/bounding_box/converters.py
index a07bf4e..9a93fb2 100644
--- a/keras_cv/bounding_box/converters.py
+++ b/keras_cv/bounding_box/converters.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 The KerasCV Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -511,9 +511,8 @@ def _image_shape(images, image_shape, boxes):
             width = tf.reshape(
                 tf.reduce_max(images.row_lengths(axis=2), 1), (-1, 1)
             )
-            if isinstance(boxes, tf.RaggedTensor):
-                height = tf.expand_dims(height, axis=-1)
-                width = tf.expand_dims(width, axis=-1)
+            height = tf.expand_dims(height, axis=-1)
+            width = tf.expand_dims(width, axis=-1)
     else:
         height, width = image_shape[0], image_shape[1]
     return tf.cast(height, boxes.dtype), tf.cast(width, boxes.dtype)
diff --git a/keras_cv/layers/__init__.py b/keras_cv/layers/__init__.py
index 68c07f3..5399d74 100644
--- a/keras_cv/layers/__init__.py
+++ b/keras_cv/layers/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 The KerasCV Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -83,6 +83,9 @@ from keras_cv.layers.preprocessing.repeated_augmentation import (
 from keras_cv.layers.preprocessing.rescaling import Rescaling
 from keras_cv.layers.preprocessing.resizing import Resizing
 from keras_cv.layers.preprocessing.solarization import Solarization
+from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (
+    VectorizedBaseImageAugmentationLayer,
+)
 from keras_cv.layers.preprocessing_3d.frustum_random_dropping_points import (
     FrustumRandomDroppingPoints,
 )
diff --git a/keras_cv/layers/preprocessing/__init__.py b/keras_cv/layers/preprocessing/__init__.py
index 02ea9e4..564ebcc 100644
--- a/keras_cv/layers/preprocessing/__init__.py
+++ b/keras_cv/layers/preprocessing/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 The KerasCV Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -31,11 +31,13 @@ from keras_cv.layers.preprocessing.equalization import Equalization
 from keras_cv.layers.preprocessing.fourier_mix import FourierMix
 from keras_cv.layers.preprocessing.grayscale import Grayscale
 from keras_cv.layers.preprocessing.grid_mask import GridMask
+from keras_cv.layers.preprocessing.jittered_resize import JitteredResize
 from keras_cv.layers.preprocessing.maybe_apply import MaybeApply
 from keras_cv.layers.preprocessing.mix_up import MixUp
 from keras_cv.layers.preprocessing.mosaic import Mosaic
 from keras_cv.layers.preprocessing.posterization import Posterization
 from keras_cv.layers.preprocessing.rand_augment import RandAugment
+from keras_cv.layers.preprocessing.random_aspect_ratio import RandomAspectRatio
 from keras_cv.layers.preprocessing.random_augmentation_pipeline import (
     RandomAugmentationPipeline,
 )
@@ -75,3 +77,6 @@ from keras_cv.layers.preprocessing.repeated_augmentation import (
 from keras_cv.layers.preprocessing.rescaling import Rescaling
 from keras_cv.layers.preprocessing.resizing import Resizing
 from keras_cv.layers.preprocessing.solarization import Solarization
+from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (
+    VectorizedBaseImageAugmentationLayer,
+)
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 1c9505a..5e642e2 100644
--- a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
+++ b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
@@ -33,6 +33,7 @@ BATCHED = "batched"
 USE_TARGETS = "use_targets"
 
 
+@keras.utils.register_keras_serializable(package="keras_cv")
 class VectorizedBaseImageAugmentationLayer(
     keras.__internal__.layers.BaseRandomLayer
 ):
@@ -263,6 +264,7 @@ class VectorizedBaseImageAugmentationLayer(
 
     def _batch_augment(self, inputs):
         images = inputs.get(IMAGES, None)
+        raw_images = images
         labels = inputs.get(LABELS, None)
         bounding_boxes = inputs.get(BOUNDING_BOXES, None)
         keypoints = inputs.get(KEYPOINTS, None)
@@ -305,6 +307,7 @@ class VectorizedBaseImageAugmentationLayer(
                 transformations=transformations,
                 bounding_boxes=bounding_boxes,
                 images=images,
+                raw_images=raw_images,
             )
             result[LABELS] = labels
 
@@ -314,6 +317,7 @@ class VectorizedBaseImageAugmentationLayer(
                 transformations=transformations,
                 labels=labels,
                 images=images,
+                raw_images=raw_images,
             )
             bounding_boxes = bounding_box.to_ragged(bounding_boxes)
             result[BOUNDING_BOXES] = bounding_boxes
@@ -325,6 +329,7 @@ class VectorizedBaseImageAugmentationLayer(
                 labels=labels,
                 bounding_boxes=bounding_boxes,
                 images=images,
+                raw_images=raw_images,
             )
             result[KEYPOINTS] = keypoints
         if segmentation_masks is not None:
@@ -334,6 +339,7 @@ class VectorizedBaseImageAugmentationLayer(
                 labels=labels,
                 bounding_boxes=bounding_boxes,
                 images=images,
+                raw_images=raw_images,
             )
             result[SEGMENTATION_MASKS] = segmentation_masks
 
diff --git a/keras_cv/metrics/object_detection/sample_boxes.npz b/keras_cv/metrics/object_detection/test_data/sample_boxes.npz
similarity index 100%
rename from keras_cv/metrics/object_detection/sample_boxes.npz
rename to keras_cv/metrics/object_detection/test_data/sample_boxes.npz
diff --git a/keras_cv/models/object_detection_3d/center_pillar.py b/keras_cv/models/object_detection_3d/center_pillar.py
index d892ca3..987b668 100644
--- a/keras_cv/models/object_detection_3d/center_pillar.py
+++ b/keras_cv/models/object_detection_3d/center_pillar.py
@@ -26,7 +26,7 @@ class MultiClassDetectionHead(keras.layers.Layer):
 
     def __init__(
         self,
-        num_class: int,
+        num_classes: int,
         num_head_bin: Sequence[int],
         share_head: bool = False,
         name: str = "detection_head",
@@ -36,9 +36,9 @@ class MultiClassDetectionHead(keras.layers.Layer):
         self._heads = {}
         self._head_names = []
         self._per_class_prediction_size = []
-        self._num_class = num_class
+        self._num_classes = num_classes
         self._num_head_bin = num_head_bin
-        for i in range(num_class):
+        for i in range(num_classes):
             self._head_names.append(f"class_{i + 1}")
             size = 0
             # 0:1 outputs is for classification
@@ -52,7 +52,7 @@ class MultiClassDetectionHead(keras.layers.Layer):
             self._per_class_prediction_size.append(size)
 
         if not share_head:
-            for i in range(num_class):
+            for i in range(num_classes):
                 # 1x1 conv for each voxel/pixel.
                 self._heads[self._head_names[i]] = keras.layers.Conv2D(
                     filters=self._per_class_prediction_size[i],
@@ -65,7 +65,7 @@ class MultiClassDetectionHead(keras.layers.Layer):
                 kernel_size=(1, 1),
                 name="shared_head",
             )
-            for i in range(num_class):
+            for i in range(num_classes):
                 self._heads[self._head_names[i]] = shared_layer
 
     def call(self, feature: tf.Tensor, training: bool) -> List[tf.Tensor]:
@@ -79,7 +79,7 @@ class MultiClassDetectionHead(keras.layers.Layer):
 class MultiClassHeatmapDecoder(keras.layers.Layer):
     def __init__(
         self,
-        num_class,
+        num_classes,
         num_head_bin: Sequence[int],
         anchor_size: Sequence[Sequence[float]],
         max_pool_size: Sequence[int],
@@ -90,8 +90,8 @@ class MultiClassHeatmapDecoder(keras.layers.Layer):
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.num_class = num_class
-        self.class_ids = list(range(1, num_class + 1))
+        self.num_classes = num_classes
+        self.class_ids = list(range(1, num_classes + 1))
         self.num_head_bin = num_head_bin
         self.anchor_size = anchor_size
         self.max_pool_size = max_pool_size
@@ -200,8 +200,7 @@ class MultiHeadCenterPillar(keras.Model):
                 propagated to the `keras.Model` class.
         """
         losses = {}
-        # TODO(ianstenbit): Rename `num_class` to `num_classes` in this model.
-        for i in range(self._multiclass_head._num_class):
+        for i in range(self._multiclass_head._num_classes):
             losses[f"heatmap_class_{i+1}"] = heatmap_loss
             losses[f"box_class_{i+1}"] = box_loss
 


Test patch:
diff --git a/keras_cv/bounding_box/converters_test.py b/keras_cv/bounding_box/converters_test.py
index d16096f..56266ae 100644
--- a/keras_cv/bounding_box/converters_test.py
+++ b/keras_cv/bounding_box/converters_test.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 The KerasCV Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -58,6 +58,8 @@ ragged_images = tf.ragged.constant(
 
 images = tf.ones([2, 1000, 1000, 3])
 
+ragged_classes = tf.ragged.constant([[0], [0]], dtype=tf.float32)
+
 boxes = {
     "xyxy": xyxy_box,
     "center_xywh": center_xywh_box,
@@ -106,8 +108,8 @@ class ConvertersTestCase(tf.test.TestCase, parameterized.TestCase):
 
     @parameterized.named_parameters(*test_image_ragged)
     def test_converters_ragged_images(self, source, target):
-        source_box = boxes_ragged_images[source]
-        target_box = boxes_ragged_images[target]
+        source_box = _raggify(boxes_ragged_images[source])
+        target_box = _raggify(boxes_ragged_images[target])
         self.assertAllClose(
             bounding_box.convert_format(
                 source_box, source=source, target=target, images=ragged_images
@@ -190,6 +192,26 @@ class ConvertersTestCase(tf.test.TestCase, parameterized.TestCase):
             target_box,
         )
 
+    @parameterized.named_parameters(*test_image_ragged)
+    def test_dense_bounding_box_with_ragged_images(self, source, target):
+        source_box = _raggify(boxes_ragged_images[source])
+        target_box = _raggify(boxes_ragged_images[target])
+        source_bounding_boxes = {"boxes": source_box, "classes": ragged_classes}
+        source_bounding_boxes = bounding_box.to_dense(source_bounding_boxes)
+
+        result_bounding_boxes = bounding_box.convert_format(
+            source_bounding_boxes,
+            source=source,
+            target=target,
+            images=ragged_images,
+        )
+        result_bounding_boxes = bounding_box.to_ragged(result_bounding_boxes)
+
+        self.assertAllClose(
+            result_bounding_boxes["boxes"],
+            target_box,
+        )
+
 
 def _raggify(tensor):
     tensor = tf.squeeze(tensor, axis=0)
diff --git a/keras_cv/layers/object_detection/multi_class_non_max_suppression_test.py b/keras_cv/layers/object_detection/multi_class_non_max_suppression_test.py
index 8f14994..80df1c6 100644
--- a/keras_cv/layers/object_detection/multi_class_non_max_suppression_test.py
+++ b/keras_cv/layers/object_detection/multi_class_non_max_suppression_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-
 import tensorflow as tf
 
 from keras_cv import layers as cv_layers
@@ -51,18 +49,6 @@ class NmsPredictionDecoderTest(tf.test.TestCase):
         self.assertEqual(result["confidence"].shape, [8, 100])
 
 
-@unittest.expectedFailure
-class NmsPredictionDecoderTestWithXLA(tf.test.TestCase):
-    def test_decode_predictions_output_shapes(self):
-        xla_function = tf.function(
-            decode_predictions_output_shapes, jit_compile=True
-        )
-        result = xla_function()
-        self.assertEqual(result["boxes"].shape, [8, 100, 4])
-        self.assertEqual(result["classes"].shape, [8, 100])
-        self.assertEqual(result["confidence"].shape, [8, 100])
-
-
 class NmsPredictionDecoderTestWithXLAMlirBridge(tf.test.TestCase):
     def setUp(self):
         tf.config.experimental.enable_mlir_bridge()
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 7639334..193c186 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
@@ -122,6 +122,7 @@ class VectorizedAssertionLayer(VectorizedBaseImageAugmentationLayer):
         transformations=None,
         bounding_boxes=None,
         images=None,
+        raw_images=None,
         **kwargs
     ):
         assert isinstance(labels, TF_ALL_TENSOR_TYPES)
@@ -129,6 +130,7 @@ class VectorizedAssertionLayer(VectorizedBaseImageAugmentationLayer):
         assert isinstance(bounding_boxes["boxes"], TF_ALL_TENSOR_TYPES)
         assert isinstance(bounding_boxes["classes"], TF_ALL_TENSOR_TYPES)
         assert isinstance(images, TF_ALL_TENSOR_TYPES)
+        assert isinstance(raw_images, TF_ALL_TENSOR_TYPES)
         return labels
 
     def augment_bounding_boxes(
@@ -137,6 +139,7 @@ class VectorizedAssertionLayer(VectorizedBaseImageAugmentationLayer):
         transformations=None,
         labels=None,
         images=None,
+        raw_images=None,
         **kwargs
     ):
         assert isinstance(bounding_boxes["boxes"], TF_ALL_TENSOR_TYPES)
@@ -144,6 +147,7 @@ class VectorizedAssertionLayer(VectorizedBaseImageAugmentationLayer):
         assert isinstance(transformations, TF_ALL_TENSOR_TYPES)
         assert isinstance(labels, TF_ALL_TENSOR_TYPES)
         assert isinstance(images, TF_ALL_TENSOR_TYPES)
+        assert isinstance(raw_images, TF_ALL_TENSOR_TYPES)
         return bounding_boxes
 
     def augment_keypoints(
@@ -153,6 +157,7 @@ class VectorizedAssertionLayer(VectorizedBaseImageAugmentationLayer):
         labels=None,
         bounding_boxes=None,
         images=None,
+        raw_images=None,
         **kwargs
     ):
         assert isinstance(keypoints, TF_ALL_TENSOR_TYPES)
@@ -161,6 +166,7 @@ class VectorizedAssertionLayer(VectorizedBaseImageAugmentationLayer):
         assert isinstance(bounding_boxes["boxes"], TF_ALL_TENSOR_TYPES)
         assert isinstance(bounding_boxes["classes"], TF_ALL_TENSOR_TYPES)
         assert isinstance(images, TF_ALL_TENSOR_TYPES)
+        assert isinstance(raw_images, TF_ALL_TENSOR_TYPES)
         return keypoints
 
     def augment_segmentation_masks(
@@ -170,6 +176,7 @@ class VectorizedAssertionLayer(VectorizedBaseImageAugmentationLayer):
         labels=None,
         bounding_boxes=None,
         images=None,
+        raw_images=None,
         **kwargs
     ):
         assert isinstance(segmentation_masks, TF_ALL_TENSOR_TYPES)
@@ -178,6 +185,7 @@ class VectorizedAssertionLayer(VectorizedBaseImageAugmentationLayer):
         assert isinstance(bounding_boxes["boxes"], TF_ALL_TENSOR_TYPES)
         assert isinstance(bounding_boxes["classes"], TF_ALL_TENSOR_TYPES)
         assert isinstance(images, TF_ALL_TENSOR_TYPES)
+        assert isinstance(raw_images, TF_ALL_TENSOR_TYPES)
         return segmentation_masks
 
 
diff --git a/keras_cv/metrics/object_detection/box_coco_metrics_test.py b/keras_cv/metrics/object_detection/box_coco_metrics_test.py
index 19885c6..7841c1a 100644
--- a/keras_cv/metrics/object_detection/box_coco_metrics_test.py
+++ b/keras_cv/metrics/object_detection/box_coco_metrics_test.py
@@ -19,7 +19,9 @@ import tensorflow as tf
 from keras_cv import bounding_box
 from keras_cv.metrics import BoxCOCOMetrics
 
-SAMPLE_FILE = os.path.dirname(os.path.abspath(__file__)) + "/sample_boxes.npz"
+SAMPLE_FILE = (
+    os.path.dirname(os.path.abspath(__file__)) + "/test_data/sample_boxes.npz"
+)
 
 
 def load_samples(fname):
diff --git a/keras_cv/models/object_detection_3d/center_pillar_test.py b/keras_cv/models/object_detection_3d/center_pillar_test.py
index 3f4b99b..f6ca85b 100644
--- a/keras_cv/models/object_detection_3d/center_pillar_test.py
+++ b/keras_cv/models/object_detection_3d/center_pillar_test.py
@@ -71,7 +71,7 @@ class CenterPillarTest(tf.test.TestCase):
         # dimensions computed from voxel_net
         unet = self.build_centerpillar_unet([400, 400, 20])
         decoder = MultiClassHeatmapDecoder(
-            num_class=2,
+            num_classes=2,
             num_head_bin=[2, 2],
             anchor_size=[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
             max_pool_size=[3, 3],
@@ -81,7 +81,7 @@ class CenterPillarTest(tf.test.TestCase):
             spatial_size=voxel_net._spatial_size,
         )
         multiclass_head = MultiClassDetectionHead(
-            num_class=2,
+            num_classes=2,
             num_head_bin=[2, 2],
         )
         model = MultiHeadCenterPillar(
@@ -114,7 +114,7 @@ class CenterPillarTest(tf.test.TestCase):
         # dimensions computed from voxel_net
         unet = self.build_centerpillar_unet([400, 400, 20])
         decoder = MultiClassHeatmapDecoder(
-            num_class=2,
+            num_classes=2,
             num_head_bin=[2, 2],
             anchor_size=[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
             max_pool_size=[3, 3],
@@ -124,7 +124,7 @@ class CenterPillarTest(tf.test.TestCase):
             spatial_size=voxel_net._spatial_size,
         )
         multiclass_head = MultiClassDetectionHead(
-            num_class=2,
+            num_classes=2,
             num_head_bin=[2, 2],
         )
         model = MultiHeadCenterPillar(