Update PyCoCo wrapper to convert GT and predictions to Numpy

keras_cv/metrics/coco/pycoco_wrapper.py

@@ -14,6 +14,7 @@
 import copy
 
 import numpy as np
+import tree
 
 try:
     from pycocotools.coco import COCO
@@ -22,6 +23,7 @@
     COCO = object
     COCOeval = None
 
+from keras_cv.backend import ops
 from keras_cv.utils.conditional_imports import assert_pycocotools_installed
 
 METRIC_NAMES = [
@@ -196,16 +198,14 @@ def _concat_numpy(groundtruths, predictions):
     """Converts tensors to numpy arrays."""
     numpy_groundtruths = {}
     for key, val in groundtruths.items():
-        if isinstance(val, tuple):
-            val = np.concatenate(val)
-        numpy_groundtruths[key] = val
-
+        numpy_groundtruths[key] = tree.map_structure(
+            lambda x: ops.convert_to_tensor(x), val
+        )
     numpy_predictions = {}
     for key, val in predictions.items():
-        if isinstance(val, tuple):
-            val = np.concatenate(val)
-        numpy_predictions[key] = val
-
+        numpy_predictions[key] = tree.map_structure(
+            lambda x: ops.convert_to_tensor(x), val
+        )
     return numpy_groundtruths, numpy_predictions
 
 

keras_cv/models/segmentation/segformer/segformer_test.py

@@ -67,10 +67,6 @@ def test_segformer_call(self):
         self.assertAllClose(mit_pred, seg_pred)
 
     @pytest.mark.large
-    @pytest.mark.skipif(
-        keras.backend.backend() == "tensorflow",
-        reason="Seg Fault with TF Backend",
-    )
     def test_weights_change(self):
         target_size = [512, 512, 2]
 
@@ -98,10 +94,6 @@ def test_weights_change(self):
             self.assertFalse(ops.any(ops.isnan(w2)))
 
     @pytest.mark.large  # Saving is slow, so mark these large.
-    @pytest.mark.skipif(
-        keras.backend.backend() == "tensorflow",
-        reason="Seg Fault with TF Backend",
-    )
     def test_saved_model(self):
         target_size = [512, 512, 3]
 
@@ -126,10 +118,6 @@ def test_saved_model(self):
         self.assertAllClose(model_output, restored_output)
 
     @pytest.mark.large  # Saving is slow, so mark these large.
-    @pytest.mark.skipif(
-        keras.backend.backend() == "tensorflow",
-        reason="Seg Fault with TF Backend",
-    )
     def test_preset_saved_model(self):
         target_size = [224, 224, 3]