:%s/.shape/tf.shape()/g

Replace some shape calls in COCOBase with tf.shape calls to support
running in graph mode.

Author: (Ian Stenbit)

keras_cv/metrics/coco/base.py

@@ -100,10 +100,10 @@ def update_state(self, y_true, y_pred, sample_weight=None):
             raise NotImplementedError(
                 "sample_weight is not yet supported in keras_cv COCO metrics."
             )
-        num_images = y_true.shape[0]
+        num_images = tf.shape(y_true)[0]
 
-        num_thresholds = self.iou_thresholds.shape[0]
-        num_categories = self.category_ids.shape[0]
+        num_thresholds = tf.shape(self.iou_thresholds)[0]
+        num_categories = tf.shape(self.category_ids)[0]
 
         # Sort by bbox.CONFIDENCE to make maxDetections easy to compute.
         y_pred = utils.sort_bboxes(y_pred, axis=bbox.CONFIDENCE)
@@ -151,9 +151,7 @@ def update_state(self, y_true, y_pred, sample_weight=None):
                     false_positives = tf.cast(pred_matches == -1, tf.float32)
 
                     true_positives_sum = tf.math.reduce_sum(true_positives, axis=-1)
-                    false_positives_sum = tf.math.reduce_sum(
-                        false_positives, axis=-1
-                    )
+                    false_positives_sum = tf.math.reduce_sum(false_positives, axis=-1)
 
                     true_positives_update = tf.tensor_scatter_nd_add(
                         true_positives_update, [indices], [true_positives_sum]

keras_cv/metrics/coco/recall_test.py

@@ -15,22 +15,43 @@
 
 import numpy as np
 import tensorflow as tf
+from tensorflow import keras
 
 from keras_cv.metrics.coco.recall import COCORecall
 
 
 class COCORecallTest(tf.test.TestCase):
-    def test_recall_area_range_filtering(self):
+
+    def test_runs_inside_model(self):
+        i = keras.layers.Input((None, None, 6)) 
+        model = keras.Model(i, i)
+
         recall = COCORecall(
             max_detections=100,
             category_ids=[1],
-            area_range=(32 ** 2, 64 ** 2),
+            area_range=(0, 64 ** 2),
         )
 
         # These would match if they were in the area range
-        y_true = np.array([[[0, 0, 10, 10, 1], [5, 5, 10, 10, 1]]]).astype(
+        y_true = np.array([[[0, 0, 10, 10, 1], [5, 5, 10, 10, 1]]]).astype(np.float32)
+        y_pred = np.array([[[0, 0, 10, 10, 1, 1.0], [5, 5, 10, 10, 1, 0.9]]]).astype(
             np.float32
         )
+
+        model.compile(metrics=[recall])
+        model.evaluate(y_pred, y_true)
+
+        self.assertAllEqual(recall.result(), 1.0)
+
+    def test_recall_area_range_filtering(self):
+        recall = COCORecall(
+            max_detections=100,
+            category_ids=[1],
+            area_range=(32 ** 2, 64 ** 2),
+        )
+
+        # These would match if they were in the area range
+        y_true = np.array([[[0, 0, 10, 10, 1], [5, 5, 10, 10, 1]]]).astype(np.float32)
         y_pred = np.array([[[0, 0, 10, 10, 1, 1.0], [5, 5, 10, 10, 1, 0.9]]]).astype(
             np.float32
         )