From 69364f547b8c876d5eb673f198cc039779ab83b0 Mon Sep 17 00:00:00 2001
From: SkafteNicki <skaftenicki@gmail.com>
Date: Thu, 28 Apr 2022 14:02:50 +0200
Subject: [PATCH] update code

---
 CHANGELOG.md                      |  3 +++
 tests/detection/test_map.py       | 13 +++++++++++++
 torchmetrics/detection/mean_ap.py | 14 ++++++++++----
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8064ef60063..3c67110d222 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -58,6 +58,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed `CalibrationError` to work on logit input ([#985](https://github.com/PyTorchLightning/metrics/pull/985))
 
 
+- Fixed MAP metric when using custom list of thresholds ([#995](https://github.com/PyTorchLightning/metrics/issues/995))
+
+
 ## [0.8.0] - 2022-04-14
 
 ### Added
diff --git a/tests/detection/test_map.py b/tests/detection/test_map.py
index f63e3d80814..151d9251208 100644
--- a/tests/detection/test_map.py
+++ b/tests/detection/test_map.py
@@ -312,6 +312,19 @@ def test_map_gpu(inputs):
     metric.compute()
 
 
+@pytest.mark.skipif(_pytest_condition, reason="test requires that torchvision=>0.8.0 is installed")
+@pytest.mark.skipif(_gpu_test_condition, reason="test requires CUDA availability")
+def test_map_with_custom_thresholds():
+    """Test that map works with custom iou thresholds."""
+    metric = MeanAveragePrecision(iou_thresholds=[0.1, 0.2])
+    metric = metric.to("cuda")
+    for preds, targets in zip(_inputs.preds, _inputs.target):
+        metric.update(_move_to_gpu(preds), _move_to_gpu(targets))
+    res = metric.compute()
+    assert res["map_50"].item() == -1
+    assert res["map_75"].item() == -1
+
+
 @pytest.mark.skipif(_pytest_condition, reason="test requires that pycocotools and torchvision=>0.8.0 is installed")
 def test_empty_metric():
     """Test empty metric."""
diff --git a/torchmetrics/detection/mean_ap.py b/torchmetrics/detection/mean_ap.py
index 1ef6f5ae69c..8941ecf2c8b 100644
--- a/torchmetrics/detection/mean_ap.py
+++ b/torchmetrics/detection/mean_ap.py
@@ -653,8 +653,14 @@ def _summarize_results(self, precisions: Tensor, recalls: Tensor) -> Tuple[MAPMe
         map_metrics = MAPMetricResults()
         map_metrics.map = self._summarize(results, True)
         last_max_det_thr = self.max_detection_thresholds[-1]
-        map_metrics.map_50 = self._summarize(results, True, iou_threshold=0.5, max_dets=last_max_det_thr)
-        map_metrics.map_75 = self._summarize(results, True, iou_threshold=0.75, max_dets=last_max_det_thr)
+        if 0.5 in self.iou_thresholds:
+            map_metrics.map_50 = self._summarize(results, True, iou_threshold=0.5, max_dets=last_max_det_thr)
+        else:
+            map_metrics.map_50 = torch.tensor([-1])
+        if 0.75 in self.iou_thresholds:
+            map_metrics.map_75 = self._summarize(results, True, iou_threshold=0.75, max_dets=last_max_det_thr)
+        else:
+            map_metrics.map_75 = torch.tensor([-1])
         map_metrics.map_small = self._summarize(results, True, area_range="small", max_dets=last_max_det_thr)
         map_metrics.map_medium = self._summarize(results, True, area_range="medium", max_dets=last_max_det_thr)
         map_metrics.map_large = self._summarize(results, True, area_range="large", max_dets=last_max_det_thr)
@@ -750,8 +756,6 @@ def compute(self) -> dict:
             dict containing
 
             - map: ``torch.Tensor``
-            - map_50: ``torch.Tensor``
-            - map_75: ``torch.Tensor``
             - map_small: ``torch.Tensor``
             - map_medium: ``torch.Tensor``
             - map_large: ``torch.Tensor``
@@ -761,6 +765,8 @@ def compute(self) -> dict:
             - mar_small: ``torch.Tensor``
             - mar_medium: ``torch.Tensor``
             - mar_large: ``torch.Tensor``
+            - map_50: ``torch.Tensor`` (-1 if 0.5 not in the list of iou thresholds)
+            - map_75: ``torch.Tensor`` (-1 if 0.75 not in the list of iou thresholds)
             - map_per_class: ``torch.Tensor`` (-1 if class metrics are disabled)
             - mar_100_per_class: ``torch.Tensor`` (-1 if class metrics are disabled)
         """