Repository: keras-team/keras-cv

Issue:
I've copy-pasted retinanet training example to colab. It fails
It's failing at the end of the first epoch here. It failed the same way mid epoch when I tried using my private dataset. 
Using YOLOv8 also results in a fail.
Logs:
```
Epoch 1/100
4138/4138 [==============================] - ETA: 0s - loss: 0.9988 - box_loss: 0.4388 - classification_loss: 0.5600 - percent_boxes_matched_with_anchor: 0.9110

---------------------------------------------------------------------------

TypeError                                 Traceback (most recent call last)

[<ipython-input-18-816f3ea473f9>](https://localhost:8080/#) in <cell line: 1>()
----> 1 history = model.fit(
      2     train_ds,
      3     validation_data=eval_ds,
      4     epochs=FLAGS.epochs,
      5     callbacks=callbacks,

1 frames

[/usr/local/lib/python3.10/dist-packages/tensorflow/python/autograph/impl/api.py](https://localhost:8080/#) in converted_call(f, args, kwargs, caller_fn_scope, options)
    437     try:
    438       if kwargs is not None:
--> 439         result = converted_f(*effective_args, **kwargs)
    440       else:
    441         result = converted_f(*effective_args)

TypeError: in user code:
https://colab.research.google.com/drive/1pYE4gv3yfQ14JaQcpbAxKwd0Gh6uNsvf?usp=sharing

    TypeError: outer_factory.<locals>.inner_factory.<locals>.tf__images_only() takes 1 positional argument but 2 were given
```
Colab link:
https://colab.research.google.com/drive/1pYE4gv3yfQ14JaQcpbAxKwd0Gh6uNsvf?usp=sharing
Code from:
https://github.com/keras-team/keras-cv/blob/master/examples/training/object_detection/pascal_voc/retinanet.py#L52

Discussion:

Patch:
diff --git a/examples/training/object_detection/pascal_voc/retinanet.py b/examples/training/object_detection/pascal_voc/retinanet.py
index a6e932e..ee787e2 100644
--- a/examples/training/object_detection/pascal_voc/retinanet.py
+++ b/examples/training/object_detection/pascal_voc/retinanet.py
@@ -149,9 +149,10 @@ train_ds = train_ds.map(augmenter, num_parallel_calls=tf.data.AUTOTUNE)
 
 
 def pad_fn(inputs):
-    return inputs["images"], keras_cv.bounding_box.to_dense(
+    inputs["bounding_boxes"] = keras_cv.bounding_box.to_dense(
         inputs["bounding_boxes"], max_boxes=32
     )
+    return inputs
 
 
 train_ds = train_ds.shuffle(8 * strategy.num_replicas_in_sync)
diff --git a/examples/visualization/plot_image_gallery.py b/examples/visualization/plot_image_gallery.py
index 297cc6d..17197ac 100644
--- a/examples/visualization/plot_image_gallery.py
+++ b/examples/visualization/plot_image_gallery.py
@@ -1,8 +1,9 @@
 """
 Title: Plot an image gallery
-Author: [lukewood](https://lukewood.xyz)
+Author: [lukewood](https://lukewood.xyz), updated by
+[Suvaditya Mukherjee](https://twitter.com/halcyonrayes)
 Date created: 2022/10/16
-Last modified: 2022/10/16
+Last modified: 2022/06/24
 Description: Visualize ground truth and predicted bounding boxes for a given
              dataset.
 """
@@ -11,7 +12,7 @@ Description: Visualize ground truth and predicted bounding boxes for a given
 Plotting images from a TensorFlow dataset is easy with KerasCV. Behold:
 """
 
-import tensorflow as tf
+import numpy as np
 import tensorflow_datasets as tfds
 
 import keras_cv
@@ -23,18 +24,20 @@ train_ds = tfds.load(
     shuffle_files=True,
 )
 
+keras_cv.visualization.plot_image_gallery(
+    train_ds,
+    value_range=(0, 255),
+    scale=3,
+)
 
-def unpackage_tfds_inputs(inputs):
-    return inputs["image"]
+"""
+If you want to use plain NumPy arrays, you can do that too:
+"""
 
+# Prepare some sample NumPy arrays from random noise
 
-train_ds = train_ds.map(unpackage_tfds_inputs)
-train_ds = train_ds.apply(tf.data.experimental.dense_to_ragged_batch(16))
+samples = np.random.randint(0, 255, (20, 224, 224, 3))
 
 keras_cv.visualization.plot_image_gallery(
-    next(iter(train_ds.take(1))),
-    value_range=(0, 255),
-    scale=3,
-    rows=2,
-    cols=2,
+    samples, value_range=(0, 255), scale=3, rows=4, cols=5
 )
diff --git a/keras_cv/metrics/coco/pycoco_wrapper.py b/keras_cv/metrics/coco/pycoco_wrapper.py
index 60c94a2..ff82b2d 100644
--- a/keras_cv/metrics/coco/pycoco_wrapper.py
+++ b/keras_cv/metrics/coco/pycoco_wrapper.py
@@ -152,7 +152,7 @@ def _convert_groundtruths_to_coco_dataset(groundtruths, label_map=None):
     gt_annotations = []
     num_batches = len(groundtruths["source_id"])
     for i in range(num_batches):
-        max_num_instances = groundtruths["classes"][i].shape[1]
+        max_num_instances = max(x.shape[0] for x in groundtruths["classes"][i])
         batch_size = groundtruths["source_id"][i].shape[0]
         for j in range(batch_size):
             num_instances = groundtruths["num_detections"][i][j]
diff --git a/keras_cv/metrics/object_detection/box_coco_metrics.py b/keras_cv/metrics/object_detection/box_coco_metrics.py
index 2e58f10..865af69 100644
--- a/keras_cv/metrics/object_detection/box_coco_metrics.py
+++ b/keras_cv/metrics/object_detection/box_coco_metrics.py
@@ -223,6 +223,13 @@ class BoxCOCOMetrics(keras.metrics.Metric):
     def update_state(self, y_true, y_pred, sample_weight=None):
         self._eval_step_count += 1
 
+        if isinstance(y_true["boxes"], tf.RaggedTensor) != isinstance(
+            y_pred["boxes"], tf.RaggedTensor
+        ):
+            # Make sure we have same ragged/dense status for y_true and y_pred
+            y_true = bounding_box.to_dense(y_true)
+            y_pred = bounding_box.to_dense(y_pred)
+
         self.ground_truths.append(y_true)
         self.predictions.append(y_pred)
 
diff --git a/keras_cv/visualization/plot_image_gallery.py b/keras_cv/visualization/plot_image_gallery.py
index bb9bb06..5926929 100644
--- a/keras_cv/visualization/plot_image_gallery.py
+++ b/keras_cv/visualization/plot_image_gallery.py
@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
+
+import numpy as np
+import tensorflow as tf
+
 import keras_cv
 from keras_cv import utils
 from keras_cv.utils import assert_matplotlib_installed
@@ -22,12 +27,43 @@ except:
     plt = None
 
 
+def _extract_image_batch(images, num_images, batch_size):
+    def unpack_images(inputs):
+        return inputs["image"]
+
+    num_batches_required = math.ceil(num_images / batch_size)
+
+    if isinstance(images, tf.data.Dataset):
+        images = images.map(unpack_images)
+
+        if batch_size == 1:
+            images = images.ragged_batch(num_batches_required)
+            sample = next(iter(images.take(1)))
+        else:
+            sample = next(iter(images.take(num_batches_required)))
+
+        return sample
+
+    else:
+        if len(images.shape) != 4:
+            raise ValueError(
+                "`plot_images_gallery()` requires you to "
+                "batch your `np.array` samples together."
+            )
+        else:
+            num_samples = (
+                num_images if num_images <= batch_size else num_batches_required
+            )
+            sample = images[:num_samples, ...]
+    return sample
+
+
 def plot_image_gallery(
     images,
     value_range,
-    rows=3,
-    cols=3,
     scale=2,
+    rows=None,
+    cols=None,
     path=None,
     show=None,
     transparent=True,
@@ -45,32 +81,27 @@ def plot_image_gallery(
         shuffle_files=True,
     )
 
-
-    def unpackage_tfds_inputs(inputs):
-        return inputs["image"]
-
-    train_ds = train_ds.map(unpackage_tfds_inputs)
-    train_ds = train_ds.apply(tf.data.experimental.dense_to_ragged_batch(16))
-
     keras_cv.visualization.plot_image_gallery(
-        next(iter(train_ds.take(1))),
+        train_ds,
         value_range=(0, 255),
         scale=3,
-        rows=2,
-        cols=2,
     )
     ```
 
     ![example gallery](https://i.imgur.com/r0ndse0.png)
 
     Args:
-        images: a Tensor or NumPy array containing images to show in the
-            gallery.
+        images: a Tensor, `tf.data.Dataset` or NumPy array containing images
+            to show in the gallery. Note: If using a `tf.data.Dataset`,
+            images should be present in the `FeaturesDict` under
+            the key `image`.
         value_range: value range of the images. Common examples include
             `(0, 255)` and `(0, 1)`.
-        rows: number of rows in the gallery to show.
-        cols: number of columns in the gallery to show.
         scale: how large to scale the images in the gallery
+        rows: (Optional) number of rows in the gallery to show.
+            Required if inputs are unbatched.
+        cols: (Optional) number of columns in the gallery to show.
+            Required if inputs are unbatched.
         path: (Optional) path to save the resulting gallery to.
         show: (Optional) whether to show the gallery of images.
         transparent: (Optional) whether to give the image a transparent
@@ -80,6 +111,7 @@ def plot_image_gallery(
         legend_handles: (Optional) matplotlib.patches List of legend handles.
             I.e. passing: `[patches.Patch(color='red', label='mylabel')]` will
             produce a legend with a single red patch and the label 'mylabel'.
+
     """
     assert_matplotlib_installed("plot_bounding_box_gallery")
 
@@ -92,27 +124,57 @@ def plot_image_gallery(
             "to be true."
         )
 
-    fig = plt.figure(figsize=(cols * scale, rows * scale))
-    fig.tight_layout()  # Or equivalently,  "plt.tight_layout()"
-    plt.subplots_adjust(wspace=0, hspace=0)
-    plt.margins(x=0, y=0)
-    plt.axis("off")
+    if isinstance(images, tf.data.Dataset):
+        sample = next(iter(images.take(1)))
+        batch_size = (
+            sample["image"].shape[0] if len(sample["image"].shape) == 4 else 1
+        )  # batch_size from within passed `tf.data.Dataset`
+    else:
+        batch_size = (
+            images.shape[0] if len(images.shape) == 4 else 1
+        )  # batch_size from np.array or single image
+
+    rows = rows or int(math.ceil(math.sqrt(batch_size)))
+    cols = cols or int(math.ceil(batch_size // rows))
+
+    num_images = rows * cols
+    images = _extract_image_batch(images, num_images, batch_size)
+
+    # Generate subplots
+    fig, axes = plt.subplots(
+        nrows=rows,
+        ncols=cols,
+        figsize=(cols * scale, rows * scale),
+        frameon=False,
+        layout="tight",
+        squeeze=True,
+        sharex="row",
+        sharey="col",
+    )
+    fig.subplots_adjust(wspace=0, hspace=0)
+
+    if isinstance(axes, np.ndarray) and len(axes.shape) == 1:
+        expand_axis = 0 if rows == 1 else -1
+        axes = np.expand_dims(axes, expand_axis)
 
     if legend_handles is not None:
         fig.legend(handles=legend_handles, loc="lower center")
 
+    # Perform image range transform
     images = keras_cv.utils.transform_value_range(
         images, original_range=value_range, target_range=(0, 255)
     )
     images = utils.to_numpy(images)
-    images = images.astype(int)
+
     for row in range(rows):
         for col in range(cols):
             index = row * cols + col
-            plt.subplot(rows, cols, index + 1)
-            plt.imshow(images[index].astype("uint8"))
-            plt.axis("off")
-            plt.margins(x=0, y=0)
+            current_axis = (
+                axes[row, col] if isinstance(axes, np.ndarray) else axes
+            )
+            current_axis.imshow(images[index].astype("uint8"))
+            current_axis.margins(x=0, y=0)
+            current_axis.axis("off")
 
     if path is None and not show:
         return


Test patch:
diff --git a/keras_cv/conftest.py b/keras_cv/conftest.py
index a91dfbe..97a31db 100644
--- a/keras_cv/conftest.py
+++ b/keras_cv/conftest.py
@@ -61,6 +61,8 @@ def pytest_collection_modifyitems(config, items):
     for item in items:
         if "keras_format" in item.name:
             item.add_marker(skip_keras_saving_test)
+        if "tf_format" in item.name:
+            item.add_marker(skip_extra_large)
         if "large" in item.keywords:
             item.add_marker(skip_large)
         if "extra_large" in item.keywords:
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 9638d9e..b4ee7e2 100644
--- a/keras_cv/metrics/object_detection/box_coco_metrics_test.py
+++ b/keras_cv/metrics/object_detection/box_coco_metrics_test.py
@@ -139,41 +139,129 @@ class BoxCOCOMetricsTest(tf.test.TestCase):
             )
 
     def test_coco_metric_suite_ragged_prediction(self):
-        bounding_boxes = {
+        suite = BoxCOCOMetrics(bounding_box_format="xyxy", evaluate_freq=1)
+        ragged_bounding_boxes = {
+            # shape: (2, (2, 1), 4)
             "boxes": tf.ragged.constant(
                 [
                     [[10, 10, 20, 20], [100, 100, 150, 150]],  # small, medium
-                    [
-                        [200, 200, 400, 400],  # large
-                    ],
+                    [[200, 200, 400, 400]],  # large
                 ],
+                ragged_rank=1,
                 dtype=tf.float32,
             ),
             "classes": tf.ragged.constant(
+                [[0, 1], [2]],
+                ragged_rank=1,
+                dtype=tf.float32,
+            ),
+            "confidence": tf.ragged.constant(
+                [[0.7, 0.8], [0.9]],
+                ragged_rank=1,
+                dtype=tf.float32,
+            ),
+        }
+        different_ragged_bounding_boxes = {
+            # shape: (2, (2, 3), 4)
+            "boxes": tf.ragged.constant(
                 [
-                    [0, 1],
-                    [2],
+                    [[10, 10, 25, 25], [100, 105, 155, 155]],
+                    [[200, 200, 450, 450], [1, 1, 5, 5], [50, 50, 300, 300]],
                 ],
+                ragged_rank=1,
+                dtype=tf.float32,
+            ),
+            "classes": tf.ragged.constant(
+                [[0, 1], [2, 3, 3]],
+                ragged_rank=1,
                 dtype=tf.float32,
             ),
             "confidence": tf.ragged.constant(
+                [[0.7, 0.8], [0.9, 0.7, 0.7]],
+                ragged_rank=1,
+                dtype=tf.float32,
+            ),
+        }
+
+        suite.update_state(
+            ragged_bounding_boxes,
+            bounding_box.to_dense(ragged_bounding_boxes),
+        )
+        metrics = suite.result()
+        for metric in metrics:
+            # The metrics will be all 1.0 because the predictions and ground
+            # truth values are identical.
+            self.assertEqual(metrics[metric], 1.0)
+
+        suite.reset_state()
+        suite.update_state(
+            ragged_bounding_boxes,
+            bounding_box.to_dense(different_ragged_bounding_boxes),
+        )
+        metrics = suite.result()
+        for metric in metrics:
+            # The metrics will not be 1.0 because the predictions and ground
+            # truth values are completely different.
+            self.assertNotEqual(metrics[metric], 1.0)
+
+    def test_coco_metric_suite_ragged_labels(self):
+        suite = BoxCOCOMetrics(bounding_box_format="xyxy", evaluate_freq=1)
+        ragged_bounding_boxes = {
+            # shape: (2, (2, 1), 4)
+            "boxes": tf.ragged.constant(
                 [
-                    [0.7, 0.8],
-                    [0.9],
+                    [[10, 10, 20, 20], [100, 100, 150, 150]],  # small, medium
+                    [[200, 200, 400, 400]],  # large
                 ],
+                ragged_rank=1,
+                dtype=tf.float32,
+            ),
+            "classes": tf.ragged.constant(
+                [[0, 1], [2]],
+                ragged_rank=1,
+                dtype=tf.float32,
+            ),
+            "confidence": tf.ragged.constant(
+                [[0.7, 0.8], [0.9]],
+                ragged_rank=1,
+                dtype=tf.float32,
+            ),
+        }
+        different_ragged_bounding_boxes = {
+            # shape: (2, (2, 3), 4)
+            "boxes": tf.ragged.constant(
+                [
+                    [[10, 10, 25, 25], [100, 105, 155, 155]],
+                    [[200, 200, 450, 450], [1, 1, 5, 5], [50, 50, 300, 300]],
+                ],
+                ragged_rank=1,
+                dtype=tf.float32,
+            ),
+            "classes": tf.ragged.constant(
+                [[0, 1], [2, 3, 3]],
+                ragged_rank=1,
+                dtype=tf.float32,
+            ),
+            "confidence": tf.ragged.constant(
+                [[0.7, 0.8], [0.9, 0.7, 0.7]],
+                ragged_rank=1,
                 dtype=tf.float32,
             ),
         }
-        dense_bounding_boxes = bounding_box.to_dense(bounding_boxes)
-        ragged_bounding_boxes = bounding_box.to_ragged(dense_bounding_boxes)
-        suite = BoxCOCOMetrics(bounding_box_format="xyxy", evaluate_freq=1)
-        y_true = dense_bounding_boxes
-        y_pred = ragged_bounding_boxes
 
-        suite.update_state(y_true, y_pred)
-        metrics = suite.result(force=True)
+        suite.update_state(ragged_bounding_boxes, ragged_bounding_boxes)
+        metrics = suite.result()
+        for metric in metrics:
+            # The metrics will be all 1.0 because the predictions and ground
+            # truth values are identical.
+            self.assertEqual(metrics[metric], 1.0)
 
+        suite.reset_state()
+        suite.update_state(
+            ragged_bounding_boxes, different_ragged_bounding_boxes
+        )
+        metrics = suite.result()
         for metric in metrics:
-            # The metrics will be all 1.0 because the prediction and ground
-            # truth is identical.
-            self.assertAllEqual(metrics[metric], 1.0)
+            # The metrics will not be 1.0 because the predictions and ground
+            # truth values are completely different.
+            self.assertNotEqual(metrics[metric], 1.0)
diff --git a/keras_cv/models/object_detection/retinanet/retinanet_test.py b/keras_cv/models/object_detection/retinanet/retinanet_test.py
index 508fb85..1297d94 100644
--- a/keras_cv/models/object_detection/retinanet/retinanet_test.py
+++ b/keras_cv/models/object_detection/retinanet/retinanet_test.py
@@ -265,6 +265,9 @@ class RetinaNetSmokeTest(tf.test.TestCase, parameterized.TestCase):
         )
         xs, _ = _create_bounding_box_dataset(bounding_box_format="xywh")
         output = model(xs)
+
+        # 4 represents number of parameters in a box
+        # 49104 is the number of anchors for a 512x512 image
         self.assertEqual(output["box"].shape, (xs.shape[0], 49104, 4))
 
     def test_full_preset_weight_loading(self):
diff --git a/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector_test.py b/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector_test.py
index 1eea74f..1d1a9b7 100644
--- a/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector_test.py
+++ b/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector_test.py
@@ -21,6 +21,9 @@ from tensorflow import keras
 
 import keras_cv
 from keras_cv import bounding_box
+from keras_cv.models.backbones.test_backbone_presets import (
+    test_backbone_presets,
+)
 from keras_cv.models.object_detection.__test_utils__ import (
     _create_bounding_box_dataset,
 )
@@ -195,8 +198,23 @@ class YOLOV8DetectorTest(tf.test.TestCase, parameterized.TestCase):
 
 
 @pytest.mark.large
-class YOLOV8DetectorSmokeTest(tf.test.TestCase):
-    # TODO(ianstenbit): Update this test to use a KerasCV-trained preset.
+class YOLOV8DetectorSmokeTest(tf.test.TestCase, parameterized.TestCase):
+    @parameterized.named_parameters(
+        *[(preset, preset) for preset in test_backbone_presets]
+    )
+    def test_backbone_preset(self, preset):
+        model = keras_cv.models.YOLOV8Detector.from_preset(
+            preset,
+            num_classes=20,
+            bounding_box_format="xywh",
+        )
+        xs, _ = _create_bounding_box_dataset(bounding_box_format="xywh")
+        output = model(xs)
+
+        # 64 represents number of parameters in a box
+        # 5376 is the number of anchors for a 512x512 image
+        self.assertEqual(output["boxes"].shape, (xs.shape[0], 5376, 64))
+
     def test_preset_with_forward_pass(self):
         model = keras_cv.models.YOLOV8Detector.from_preset(
             "yolo_v8_m_pascalvoc",