fix confidence handling for keypoint evaluation

voxel51 · Jan 6, 2025 · e37e4fd · e37e4fd
1 parent 199872a
commit e37e4fd
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 8 deletions.
diff --git a/fiftyone/utils/eval/coco.py b/fiftyone/utils/eval/coco.py
@@ -13,6 +13,7 @@
 
 import eta.core.utils as etau
 
+import fiftyone.core.labels as fol
 import fiftyone.core.plots as fop
 import fiftyone.utils.iou as foui
 
@@ -552,14 +553,19 @@ def _coco_evaluation_setup(
             label = obj.label if classwise else "all"
             cats[label]["preds"].append(obj)
 
+    if isinstance(preds, fol.Keypoints):
+        sort_key = lambda p: np.nanmean(p.confidence) if p.confidence else -1
+    else:
+        sort_key = lambda p: p.confidence or -1
+
     # Compute IoUs within each category
     pred_ious = {}
     for objects in cats.values():
         gts = objects["gts"]
         preds = objects["preds"]
 
         # Highest confidence predictions first
-        preds = sorted(preds, key=lambda p: p.confidence or -1, reverse=True)
+        preds = sorted(preds, key=sort_key, reverse=True)
 
         if max_preds is not None:
             preds = preds[:max_preds]
@@ -587,6 +593,13 @@ def _compute_matches(
 
         # Match each prediction to the highest available IoU ground truth
         for pred in objects["preds"]:
+            if isinstance(pred, fol.Keypoint):
+                pred_conf = (
+                    np.nanmean(pred.confidence) if pred.confidence else None
+                )
+            else:
+                pred_conf = pred.confidence
+
             if pred.id in pred_ious:
                 best_match = None
                 best_match_iou = iou_thresh
@@ -638,7 +651,7 @@ def _compute_matches(
                             gt.label,
                             pred.label,
                             best_match_iou,
-                            pred.confidence,
+                            pred_conf,
                             gt.id,
                             pred.id,
                             iscrowd(gt),
@@ -651,7 +664,7 @@ def _compute_matches(
                             None,
                             pred.label,
                             None,
-                            pred.confidence,
+                            pred_conf,
                             None,
                             pred.id,
                             None,
@@ -665,7 +678,7 @@ def _compute_matches(
                         None,
                         pred.label,
                         None,
-                        pred.confidence,
+                        pred_conf,
                         None,
                         pred.id,
                         None,

diff --git a/fiftyone/utils/eval/openimages.py b/fiftyone/utils/eval/openimages.py
@@ -10,6 +10,7 @@
 
 import numpy as np
 
+import fiftyone.core.labels as fol
 import fiftyone.core.plots as fop
 import fiftyone.utils.iou as foui
 
@@ -545,14 +546,19 @@ def _open_images_evaluation_setup(
                 if config.expand_pred_hierarchy and label != "all":
                     _expand_detection_hierarchy(cats, obj, config, "preds")
 
+    if isinstance(preds, fol.Keypoints):
+        sort_key = lambda p: np.nanmean(p.confidence) if p.confidence else -1
+    else:
+        sort_key = lambda p: p.confidence or -1
+
     # Compute IoUs within each category
     pred_ious = {}
     for objects in cats.values():
         gts = objects["gts"]
         preds = objects["preds"]
 
         # Highest confidence predictions first
-        preds = sorted(preds, key=lambda p: p.confidence or -1, reverse=True)
+        preds = sorted(preds, key=sort_key, reverse=True)
 
         if max_preds is not None:
             preds = preds[:max_preds]
@@ -583,6 +589,13 @@ def _compute_matches(
 
         # Match each prediction to the highest available IoU ground truth
         for pred in objects["preds"]:
+            if isinstance(pred, fol.Keypoint):
+                pred_conf = (
+                    np.nanmean(pred.confidence) if pred.confidence else None
+                )
+            else:
+                pred_conf = pred.confidence
+
             if pred.id in pred_ious:
                 best_match = None
                 best_match_iou = iou_thresh
@@ -667,7 +680,7 @@ def _compute_matches(
                                 gt.label,
                                 pred.label,
                                 best_match_iou,
-                                pred.confidence,
+                                pred_conf,
                                 gt.id,
                                 pred.id,
                             )
@@ -679,15 +692,15 @@ def _compute_matches(
                             None,
                             pred.label,
                             None,
-                            pred.confidence,
+                            pred_conf,
                             None,
                             pred.id,
                         )
                     )
             elif pred.label == cat:
                 pred[eval_key] = "fp"
                 matches.append(
-                    (None, pred.label, None, pred.confidence, None, pred.id)
+                    (None, pred.label, None, pred_conf, None, pred.id)
                 )
 
         # Leftover GTs are false negatives