Repository: keras-team/keras-cv

Issue:
BoxCOCOMetrics don't support Ragged inputs
Hi KerasCV ! 

I am using your RetinaNet Model and i would like to use coco metrics. When I train the model without specifying the metrics no problems. However when I precise the coco metric the following error is yield after the first epoch : 

Epoch 1: 
  `19/Unknown - 54s 384ms/step - loss: 1.3833 - box_loss: 0.4004 - classification_loss: 0.9828 - MaP: 0.0000e+00 - MaP@[IoU=50]: 0.0000e+00 - MaP@[IoU=75]: 0.0000e+00 - MaP@[area=small]: 0.0000e+00 - MaP@[area=medium]: 0.0000e+00 - MaP@[area=large]: 0.0000e+00 - Recall@[max_detections=1]: 0.0000e+00 - Recall@[max_detections=10]: 0.0000e+00 - Recall@[max_detections=100]: 0.0000e+00 - Recall@[area=small]: 0.0000e+00 - Recall@[area=medium]: 0.0000e+00 - Recall@[area=large]: 0.0000e+00 - percent_boxes_matched_with_anchor: 0.8336`


```
Traceback (most recent call last):
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/tensorflow/python/ops/script_ops.py", line 269, in __call__
    return func(device, token, args)
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/tensorflow/python/ops/script_ops.py", line 147, in __call__
    outputs = self._call(device, args)
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/tensorflow/python/ops/script_ops.py", line 154, in _call
    ret = self._func(*args)
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/tensorflow/python/ops/script_ops.py", line 813, in wrapped_func
    out = func(*structured_inp)
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py", line 642, in wrapper
    return func(*args, **kwargs)
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/keras_cv/metrics/object_detection/box_coco_metrics.py", line 167, in update_state_on_cpu
    return obj_update_state(y_true, y_pred, sample_weight)
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/keras_cv/metrics/object_detection/box_coco_metrics.py", line 233, in update_state
    self._cached_result = self._compute_result()
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/keras_cv/metrics/object_detection/box_coco_metrics.py", line 250, in _compute_result
    metrics = compute_pycocotools_metric(
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/keras_cv/metrics/object_detection/box_coco_metrics.py", line 301, in compute_pycocotools_metric
    return keras_cv.metrics.coco.compute_pycoco_metrics(
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/keras_cv/metrics/coco/pycoco_wrapper.py", line 225, in compute_pycoco_metrics
    gt_dataset = _convert_groundtruths_to_coco_dataset(groundtruths)
  File "/home/hugo/miniconda3/envs/sinbad/lib/python3.10/site-packages/keras_cv/metrics/coco/pycoco_wrapper.py", line 155, in _convert_groundtruths_to_coco_dataset
    max_num_instances = groundtruths["classes"][i].shape[1]
IndexError: tuple index out of range
	 [[{{node EagerPyFunc_1}}]] [Op:__inference_train_function_17573]
 ```
 
 I give you a link to a gist,  I tried to reproduce the error(maybe few modifications to do)  :
 
 https://gist.github.com/hugopi/a3d168046df760b601fdae3af9d12607
 
 Thank you for your support ! 
 
 Regards
 
 Hugo PIRES
 
 


Discussion:
IMvision12: @ianstenbit @james77777778  I am getting this error : ValueError: Expected `boxes` to be a Tensor with a final dimension of `4`. Instead, got `boxes.shape=(20, None, None)`.

while using BoxCOCOMetrics, have a look!

colab : https://colab.research.google.com/drive/1Z7e75LwOwzVufE_flWJ3kxHHnNMtcaYZ?usp=sharing
james77777778: Sorry for the delay.
I've been busy lately, but I'll take a look when I have free time.
IMvision12: @james77777778 can you please rebase this branch so that I can do some experiments?
james77777778: > @james77777778 can you please rebase this branch so that I can do some experiments?

Sure. I have rebased the PR.
IMvision12: I'm still uncertain about the cause of the error, may @ianstenbit could help us!
james77777778: Hi @ianstenbit @IMvision12 

I have successfully conducted a toy training of YOLOV8 with fully ragged bounding boxes:
https://colab.research.google.com/drive/16so9QPb6WkHkXDgMoZbjKofEYHwkBzgi?usp=sharing

The root cause lies in the usage of `keras_cv.layers.Resizing`, which modifies the shape of the ragged bounding boxes to `[batch_size, None, None]` instead of the desired `[batch_size, None, 4]`.

I fixed it by
```python
def ensure_shape(inputs):
    inputs["bounding_boxes"] = keras_cv.bounding_box.to_dense(
        inputs["bounding_boxes"]
    )
    inputs["bounding_boxes"] = keras_cv.bounding_box.to_ragged(
        inputs["bounding_boxes"]
    )
    return inputs

train_ds = train_ds.map(ensure_shape, num_parallel_calls=tf.data.AUTOTUNE)
eval_ds = eval_ds.map(ensure_shape, num_parallel_calls=tf.data.AUTOTUNE)
```

Therefore, the bug mentioned by @IMvision12 appears to be unrelated to this PR.
james77777778: I changed the PR to utilize dense tensors for the conversion.
The only drawback is that padding could result in increased memory usage, although this should be generally acceptable.
ianstenbit: /gcbrun

Patch:
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/training/contrastive/contrastive_trainer.py b/keras_cv/training/contrastive/contrastive_trainer.py
index a70e194..e97435a 100644
--- a/keras_cv/training/contrastive/contrastive_trainer.py
+++ b/keras_cv/training/contrastive/contrastive_trainer.py
@@ -109,6 +109,10 @@ class ContrastiveTrainer(keras.Model):
             augmenter if type(augmenter) is tuple else (augmenter, augmenter)
         )
         self.encoder = encoder
+        # Check to see if the projector is being shared or are distinct.
+        self._is_shared_projector = (
+            True if not isinstance(projector, tuple) else False
+        )
         self.projectors = (
             projector if type(projector) is tuple else (projector, projector)
         )
@@ -230,19 +234,23 @@ class ContrastiveTrainer(keras.Model):
                 regularization_losses=self.encoder.losses,
             )
 
+        # If the projector is shared, then take the trainable weights of just
+        # one of the projectors in the tuple. If not, use both the projectors.
+        projector_weights = (
+            self.projectors[0].trainable_weights
+            if self._is_shared_projector
+            else self.projectors[0].trainable_weights
+            + self.projectors[1].trainable_weights
+        )
         gradients = tape.gradient(
             loss,
-            self.encoder.trainable_weights
-            + self.projectors[0].trainable_weights
-            + self.projectors[1].trainable_weights,
+            self.encoder.trainable_weights + projector_weights,
         )
 
         self.optimizer.apply_gradients(
             zip(
                 gradients,
-                self.encoder.trainable_weights
-                + self.projectors[0].trainable_weights
-                + self.projectors[1].trainable_weights,
+                self.encoder.trainable_weights + projector_weights,
             )
         )
         self.loss_metric.update_state(loss)


Test patch:
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)