MeanAveragePrecision box and segm at same time

CHANGELOG.md

@@ -17,6 +17,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `VisualInformationFidelity` to image package ([#1830](https://github.com/Lightning-AI/torchmetrics/pull/1830))
 
 
+- Added support for evaluating `"segm"` and `"bbox"` detection in `MeanAveragePrecision` at the same time ([#1830](https://github.com/Lightning-AI/torchmetrics/pull/1928))
+
+
 ### Changed
 
 -

src/torchmetrics/detection/helpers.py

@@ -11,20 +11,26 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Dict, Sequence
+from typing import Dict, List, Literal, Sequence, Union
 
 from torch import Tensor
 
 
 def _input_validator(
-    preds: Sequence[Dict[str, Tensor]], targets: Sequence[Dict[str, Tensor]], iou_type: str = "bbox"
+    preds: Sequence[Dict[str, Tensor]],
+    targets: Sequence[Dict[str, Tensor]],
+    iou_type: Union[Literal["bbox", "segm"], List[str]] = "bbox",
 ) -> None:
     """Ensure the correct input format of `preds` and `targets`."""
-    if iou_type == "bbox":
-        item_val_name = "boxes"
-    elif iou_type == "segm":
-        item_val_name = "masks"
-    else:
+    if not isinstance(iou_type, list):
+        iou_type = [iou_type]
+
+    item_val_name = []
+    if "bbox" in iou_type:
+        item_val_name.append("boxes")
+    if "segm" in iou_type:
+        item_val_name.append("masks")
+    if any(i not in ["bbox", "segm"] for i in iou_type):
         raise Exception(f"IOU type {iou_type} is not supported")
 
     if not isinstance(preds, Sequence):
@@ -36,42 +42,57 @@ def _input_validator(
             f"Expected argument `preds` and `target` to have the same length, but got {len(preds)} and {len(targets)}"
         )
 
-    for k in [item_val_name, "scores", "labels"]:
+    for k in [*item_val_name, "scores", "labels"]:
         if any(k not in p for p in preds):
             raise ValueError(f"Expected all dicts in `preds` to contain the `{k}` key")
 
-    for k in [item_val_name, "labels"]:
+    for k in [*item_val_name, "labels"]:
         if any(k not in p for p in targets):
             raise ValueError(f"Expected all dicts in `target` to contain the `{k}` key")
 
-    if any(type(pred[item_val_name]) is not Tensor for pred in preds):
-        raise ValueError(f"Expected all {item_val_name} in `preds` to be of type Tensor")
+    for ivn in item_val_name:
+        if any(type(pred[ivn]) is not Tensor for pred in preds):
+            raise ValueError(f"Expected all {ivn} in `preds` to be of type Tensor")
     if any(type(pred["scores"]) is not Tensor for pred in preds):
         raise ValueError("Expected all scores in `preds` to be of type Tensor")
     if any(type(pred["labels"]) is not Tensor for pred in preds):
         raise ValueError("Expected all labels in `preds` to be of type Tensor")
-    if any(type(target[item_val_name]) is not Tensor for target in targets):
-        raise ValueError(f"Expected all {item_val_name} in `target` to be of type Tensor")
+    for ivn in item_val_name:
+        if any(type(target[ivn]) is not Tensor for target in targets):
+            raise ValueError(f"Expected all {ivn} in `target` to be of type Tensor")
     if any(type(target["labels"]) is not Tensor for target in targets):
         raise ValueError("Expected all labels in `target` to be of type Tensor")
 
     for i, item in enumerate(targets):
-        if item[item_val_name].size(0) != item["labels"].size(0):
-            raise ValueError(
-                f"Input {item_val_name} and labels of sample {i} in targets have a"
-                f" different length (expected {item[item_val_name].size(0)} labels, got {item['labels'].size(0)})"
-            )
+        for ivn in item_val_name:
+            if item[ivn].size(0) != item["labels"].size(0):
+                raise ValueError(
+                    f"Input {ivn} and labels of sample {i} in targets have a"
+                    f" different length (expected {item[ivn].size(0)} labels, got {item['labels'].size(0)})"
+                )
     for i, item in enumerate(preds):
-        if not (item[item_val_name].size(0) == item["labels"].size(0) == item["scores"].size(0)):
-            raise ValueError(
-                f"Input {item_val_name}, labels and scores of sample {i} in predictions have a"
-                f" different length (expected {item[item_val_name].size(0)} labels and scores,"
-                f" got {item['labels'].size(0)} labels and {item['scores'].size(0)})"
-            )
+        for ivn in item_val_name:
+            if not (item[ivn].size(0) == item["labels"].size(0) == item["scores"].size(0)):
+                raise ValueError(
+                    f"Input {ivn}, labels and scores of sample {i} in predictions have a"
+                    f" different length (expected {item[ivn].size(0)} labels and scores,"
+                    f" got {item['labels'].size(0)} labels and {item['scores'].size(0)})"
+                )
 
 
 def _fix_empty_tensors(boxes: Tensor) -> Tensor:
     """Empty tensors can cause problems in DDP mode, this methods corrects them."""
     if boxes.numel() == 0 and boxes.ndim == 1:
         return boxes.unsqueeze(0)
     return boxes
+
+
+def _validate_iou_type_arg(iou_type: Union[Literal["bbox", "segm"], List[str]] = "bbox") -> List[str]:
+    allowed_iou_types = ("segm", "bbox")
+    if not isinstance(iou_type, list):
+        iou_type = [iou_type]
+    if any(i_type not in allowed_iou_types for i_type in iou_type):
+        raise ValueError(
+            f"Expected argument `iou_type` to be one of {allowed_iou_types} or a list of, but got {iou_type}"
+        )
+    return iou_type