Repository: Lightning-AI/torchmetrics

Issue:
Incorrect implementation of MAP
## 🐛 Bug

After merging #1259, some discrepancies have occurred as described here https://github.com/Lightning-AI/metrics/issues/1024#issuecomment-1281424230 and #1259.

### Additional context

I'll try to investigate this more.



Discussion:
Good finding, so could we at fist place add a test to cover this correctness? :)
#1259 replaces nested for-loop with `einsum` op, while the matching process implemented in `metrics` version 0.10.0 cannot be parallelized like this. The `gt_matches` (`n_threshold x m_gt`) matrix is used as **flags** to track previous matches, so that one `gt` is matched with one `det` at a IoU threshold level. I guess the incorrectness is caused by the different behaviour when more than one `det`s can match with a `gt`.
@aaronzs, would you be interested in exploring this issue further? :rabbit: 
@Borda Sure, I'm interested in. I'll find some test cases that can fail current implementation.
> @Borda Sure, I'm interested in. I'll find some test cases that can fail current implementation.

you can start with this one https://github.com/Lightning-AI/metrics/issues/1024#issuecomment-1281424230
@Borda @stancld I generate random bounding boxes base on https://github.com/Lightning-AI/metrics/issues/1024#issuecomment-1281424230 and find following case that gives an incorrect result.

```python
targets = [
    {
        "boxes": torch.tensor([[4., 2., 20., 21.]]),
        "labels": torch.tensor([1]),
    }
]
preds = [
    {
        "boxes": torch.tensor(
            [
                [2., 2., 18., 13.],
                [9., 2., 21., 20.],
                [5., 4., 21., 21.],
            ]
        ).round(),
        "labels": torch.tensor([1, 1, 1]),
        "scores": torch.tensor([0.75, 0.5, 0.25]),
    }
]
```

with torchmetrics v0.6.0 (pycocotools) and torchmetrics v0.10.2

```
'map': tensor(0.2500),
'map_50': tensor(0.5000),
'map_75': tensor(0.3333),
'map_small': tensor(0.2500),
'map_medium': tensor(-1.),
'map_large': tensor(-1.),
'mar_1': tensor(0.),
'mar_10': tensor(0.6000),
'mar_100': tensor(0.6000),
'mar_small': tensor(0.6000),
'mar_medium': tensor(-1.),
'mar_large': tensor(-1.),
'map_per_class': tensor(-1.),
'mar_100_per_class': tensor(-1.)
```

with torchmetrics master branch

```
'map': tensor(0.3000),
'map_50': tensor(0.6667),
'map_75': tensor(0.3333),
'map_small': tensor(0.3000),
'map_medium': tensor(-1.),
'map_large': tensor(-1.),
'mar_1': tensor(0.),
'mar_10': tensor(0.9000),
'mar_100': tensor(0.9000),
'mar_small': tensor(0.9000),
'mar_medium': tensor(-1.),
'mar_large': tensor(-1.),
'map_per_class': tensor(-1.),
'mar_100_per_class': tensor(-1.)
```

It is a 1-target and 3-detection case, with a lower IoU threshold, e.g. 0.5, two of these detections can match with the target. `pycocotools` picks the first match while the latest implementation marks both are matched. More specifically, `det_matches` in `_evaluate_image()` is different from previous implementation.

Resolving patch:
diff --git a/src/torchmetrics/detection/mean_ap.py b/src/torchmetrics/detection/mean_ap.py
--- a/src/torchmetrics/detection/mean_ap.py
+++ b/src/torchmetrics/detection/mean_ap.py
@@ -579,6 +579,10 @@ def _evaluate_image(
         det = [det[i] for i in det_label_mask]
         if len(gt) == 0 and len(det) == 0:
             return None
+        if isinstance(det, dict):
+            det = [det]
+        if isinstance(gt, dict):
+            gt = [gt]
 
         areas = compute_area(gt, iou_type=self.iou_type).to(self.device)
 
@@ -600,23 +604,23 @@ def _evaluate_image(
         # load computed ious
         ious = ious[idx, class_id][:, gtind] if len(ious[idx, class_id]) > 0 else ious[idx, class_id]
 
+        nb_iou_thrs = len(self.iou_thresholds)
         nb_gt = len(gt)
         nb_det = len(det)
         gt_matches = torch.zeros((nb_iou_thrs, nb_gt), dtype=torch.bool, device=self.device)
+        det_matches = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool, device=self.device)
         gt_ignore = ignore_area_sorted
-
-        iou_thresholds = torch.tensor(self.iou_thresholds, device=self.device)
+        det_ignore = torch.zeros((nb_iou_thrs, nb_det), dtype=torch.bool, device=self.device)
 
         if torch.numel(ious) > 0:
-            best_matches = self._find_best_gt_matches(iou_thresholds, gt_matches, gt_ignore, ious)
-            _zero_tensor = torch.tensor(0, dtype=torch.bool, device=self.device)
-            _one_tensor = torch.tensor(1, dtype=torch.bool, device=self.device)
-            det_ignore = torch.where(
-                best_matches != -1, gt_ignore[best_matches.clamp(max=gt_ignore.shape[0] - 1)], _zero_tensor
-            )
-            det_matches = torch.where(best_matches != -1, _one_tensor, _zero_tensor)
-            for idx in range(nb_iou_thrs):
-                gt_matches[idx, best_matches[idx].clamp(0, max=gt_matches.shape[1] - 1).unique()] = 1
+            for idx_iou, t in enumerate(self.iou_thresholds):
+                for idx_det, _ in enumerate(det):
+                    m = MeanAveragePrecision._find_best_gt_match(t, gt_matches, idx_iou, gt_ignore, ious, idx_det)
+                    if m == -1:
+                        continue
+                    det_ignore[idx_iou, idx_det] = gt_ignore[m]
+                    det_matches[idx_iou, idx_det] = 1
+                    gt_matches[idx_iou, m] = 1
 
         # set unmatched detections outside of area range to ignore
         det_areas = compute_area(det, iou_type=self.iou_type).to(self.device)
@@ -635,26 +639,33 @@ def _evaluate_image(
         }
 
     @staticmethod
-    def _find_best_gt_matches(thr: Tensor, gt_matches: Tensor, gt_ignore: Tensor, ious: Tensor) -> Tensor:
-        """Return matrix of indices of best ground truth match with current detection.
+    def _find_best_gt_match(
+        thr: int, gt_matches: Tensor, idx_iou: float, gt_ignore: Tensor, ious: Tensor, idx_det: int
+    ) -> int:
+        """Return id of best ground truth match with current detection.
 
         Args:
             thr:
                 Current threshold value.
             gt_matches:
                 Tensor showing if a ground truth matches for threshold ``t`` exists.
+            idx_iou:
+                Id of threshold ``t``.
             gt_ignore:
                 Tensor showing if ground truth should be ignored.
             ious:
                 IoUs for all combinations of detection and ground truth.
+            idx_det:
+                Id of current detection.
         """
+        previously_matched = gt_matches[idx_iou]
         # Remove previously matched or ignored gts
-        remove_mask = gt_matches | gt_ignore
-        gt_ious = torch.einsum("cw,dw->cdw", ~remove_mask, ious).max(-1).values
-        best_gt_matches = gt_ious.where(
-            gt_ious > thr.unsqueeze(-1), torch.tensor(-1, dtype=gt_ious.dtype, device=gt_ious.device)
-        )
-        return best_gt_matches.long()
+        remove_mask = previously_matched | gt_ignore
+        gt_ious = ious[idx_det] * ~remove_mask
+        match_idx = gt_ious.argmax().item()
+        if gt_ious[match_idx] > thr:
+            return match_idx
+        return -1
 
     def _summarize(
         self,


Test patch:
diff --git a/tests/unittests/detection/test_map.py b/tests/unittests/detection/test_map.py
--- a/tests/unittests/detection/test_map.py
+++ b/tests/unittests/detection/test_map.py
@@ -87,10 +87,22 @@
                 labels=IntTensor([4, 1, 0, 0, 0, 0, 0]),
             ),  # coco image id 74
             dict(
-                boxes=Tensor([[0.00, 2.87, 601.00, 421.52]]),
-                scores=Tensor([0.699]),
-                labels=IntTensor([5]),
-            ),  # coco image id 133
+                boxes=Tensor(
+                    [
+                        [72.92, 45.96, 91.23, 80.57],
+                        [45.17, 45.34, 66.28, 79.83],
+                        [82.28, 47.04, 99.66, 78.50],
+                        [59.96, 46.17, 80.35, 80.48],
+                        [75.29, 23.01, 91.85, 50.85],
+                        [71.14, 1.10, 96.96, 28.33],
+                        [61.34, 55.23, 77.14, 79.57],
+                        [41.17, 45.78, 60.99, 78.48],
+                        [56.18, 44.80, 64.42, 56.25],
+                    ]
+                ),
+                scores=Tensor([0.532, 0.204, 0.782, 0.202, 0.883, 0.271, 0.561, 0.204, 0.349]),
+                labels=IntTensor([49, 49, 49, 49, 49, 49, 49, 49, 49]),
+            ),  # coco image id 987 category_id 49
         ],
     ],
     target=[
@@ -125,9 +137,22 @@
                 labels=IntTensor([4, 1, 0, 0, 0, 0, 0]),
             ),  # coco image id 74
             dict(
-                boxes=Tensor([[13.99, 2.87, 640.00, 421.52]]),
-                labels=IntTensor([5]),
-            ),  # coco image id 133
+                boxes=Tensor(
+                    [
+                        [72.92, 45.96, 91.23, 80.57],
+                        [50.17, 45.34, 71.28, 79.83],
+                        [81.28, 47.04, 98.66, 78.50],
+                        [63.96, 46.17, 84.35, 80.48],
+                        [75.29, 23.01, 91.85, 50.85],
+                        [56.39, 21.65, 75.66, 45.54],
+                        [73.14, 1.10, 98.96, 28.33],
+                        [62.34, 55.23, 78.14, 79.57],
+                        [44.17, 45.78, 63.99, 78.48],
+                        [58.18, 44.80, 66.42, 56.25],
+                    ]
+                ),
+                labels=IntTensor([49, 49, 49, 49, 49, 49, 49, 49, 49, 49]),
+            ),  # coco image id 987 category_id 49
         ],
     ],
 )
@@ -205,18 +230,18 @@ def _compare_fn(preds, target) -> dict:
 
     Official pycocotools results calculated from a subset of https://github.com/cocodataset/cocoapi/tree/master/results
         All classes
-        Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.706
-        Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.901
-        Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.846
-        Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.689
+        Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.637
+        Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.859
+        Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.761
+        Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.622
         Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.800
-        Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.701
-        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.592
-        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.716
-        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.716
-        Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.767
+        Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.635
+        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.432
+        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.652
+        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.652
+        Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.673
         Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.800
-        Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.700
+        Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.633
 
         Class 0
         Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.725
@@ -224,7 +249,7 @@ def _compare_fn(preds, target) -> dict:
 
         Class 1
         Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.800
-        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.800
+        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.800
 
         Class 2
         Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.454
@@ -238,25 +263,25 @@ def _compare_fn(preds, target) -> dict:
         Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.650
         Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.650
 
-        Class 5
-        Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.900
-        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.900
+        Class 49
+        Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.556
+        Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.580
     """
     return {
-        "map": Tensor([0.706]),
-        "map_50": Tensor([0.901]),
-        "map_75": Tensor([0.846]),
-        "map_small": Tensor([0.689]),
+        "map": Tensor([0.637]),
+        "map_50": Tensor([0.859]),
+        "map_75": Tensor([0.761]),
+        "map_small": Tensor([0.622]),
         "map_medium": Tensor([0.800]),
-        "map_large": Tensor([0.701]),
-        "mar_1": Tensor([0.592]),
-        "mar_10": Tensor([0.716]),
-        "mar_100": Tensor([0.716]),
-        "mar_small": Tensor([0.767]),
+        "map_large": Tensor([0.635]),
+        "mar_1": Tensor([0.432]),
+        "mar_10": Tensor([0.652]),
+        "mar_100": Tensor([0.652]),
+        "mar_small": Tensor([0.673]),
         "mar_medium": Tensor([0.800]),
-        "mar_large": Tensor([0.700]),
-        "map_per_class": Tensor([0.725, 0.800, 0.454, -1.000, 0.650, 0.900]),
-        "mar_100_per_class": Tensor([0.780, 0.800, 0.450, -1.000, 0.650, 0.900]),
+        "mar_large": Tensor([0.633]),
+        "map_per_class": Tensor([0.725, 0.800, 0.454, -1.000, 0.650, 0.556]),
+        "mar_100_per_class": Tensor([0.780, 0.800, 0.450, -1.000, 0.650, 0.580]),
     }
 
 
@@ -279,7 +304,7 @@ def _compare_fn_segm(preds, target) -> dict:
     """
     return {
         "map": Tensor([0.352]),
-        "map_50": Tensor([0.742]),
+        "map_50": Tensor([0.752]),
         "map_75": Tensor([0.252]),
         "map_small": Tensor([-1]),
         "map_medium": Tensor([-1]),
@@ -299,7 +324,7 @@ def _compare_fn_segm(preds, target) -> dict:
 
 
 @pytest.mark.skipif(_pytest_condition, reason="test requires that torchvision=>0.8.0 is installed")
-@pytest.mark.parametrize("compute_on_cpu", [True, False])
+@pytest.mark.parametrize("compute_on_cpu", [True])
 class TestMAP(MetricTester):
     """Test the MAP metric for object detection predictions.
 
@@ -308,7 +333,7 @@ class TestMAP(MetricTester):
     https://github.com/cocodataset/cocoapi/blob/master/results/instances_val2014_fakebbox100_results.json
     """
 
-    atol = 1e-1
+    atol = 1e-2
 
     @pytest.mark.parametrize("ddp", [False, True])
     def test_map_bbox(self, compute_on_cpu, ddp):