feat: add class_names support to Precision and Recall metrics

ignite/metrics/fbeta.py

@@ -17,6 +17,7 @@ def Fbeta(
     recall: Recall | None = None,
     output_transform: Callable | None = None,
     device: str | torch.device | None = None,
+    class_names: list[str] | None = None,
 ) -> MetricsLambda:
     r"""Calculates F-beta score.
 
@@ -42,6 +43,7 @@ def Fbeta(
         device: specifies which device updates are accumulated on. Setting the metric's
             device to be the same as your ``update`` arguments ensures the ``update`` method is non-blocking. By
             default, CPU.
+        class_names: list of class name strings used to label per-class output. Default: ``None``.
 
     Returns:
         MetricsLambda, F-beta metric
@@ -159,23 +161,61 @@ def thresholded_output_transform(output):
     if precision is None and recall is None and device is None:
         device = torch.device("cpu")
 
+    if class_names is not None:
+        if not isinstance(class_names, (list, tuple)) or not all(isinstance(n, str) for n in class_names):
+            raise ValueError("class_names must be a list of strings")
+        if average is not False and average is not None:
+            raise ValueError(
+                f"class_names is only applicable when average=False or average=None, got average={average!r}."
+            )
+
+    active_metrics = [m for m in (precision, recall) if m is not None]
+
+    if any(m._average for m in active_metrics):
+        raise ValueError("Input precision and recall metrics should have average=False")
+
+    if class_names is not None and any(m._class_names != class_names for m in active_metrics):
+        raise ValueError("precision and recall metric class_names must match Fbeta class_names")
+
+    if len(active_metrics) == 2 and active_metrics[0]._class_names != active_metrics[1]._class_names:
+        raise ValueError("precision and recall class_names must match")
+
+    target_class_names = class_names
+    if target_class_names is None and precision is not None:
+        target_class_names = precision._class_names
+    if target_class_names is None and recall is not None:
+        target_class_names = recall._class_names
+
+    if target_class_names is not None and average is not False and average is not None:
+        raise ValueError(f"class_names is only applicable when average=False or average=None, got average={average!r}.")
+
     if precision is None:
         precision = Precision(
             output_transform=(lambda x: x) if output_transform is None else output_transform,
             average=False,
             device=cast(str | torch.device, recall._device if recall else device),
+            class_names=target_class_names,
         )
-    elif precision._average:
-        raise ValueError("Input precision metric should have average=False")
 
     if recall is None:
         recall = Recall(
             output_transform=(lambda x: x) if output_transform is None else output_transform,
             average=False,
             device=cast(str | torch.device, precision._device if precision else device),
+            class_names=target_class_names,
         )
-    elif recall._average:
-        raise ValueError("Input recall metric should have average=False")
+
+    if target_class_names is not None:
+
+        def _fbeta_with_class_names(p: dict, r: dict) -> dict:
+            p_vals = torch.tensor(list(p.values()))
+            r_vals = torch.tensor(list(r.values()))
+            scores = (1.0 + beta**2) * p_vals * r_vals / (beta**2 * p_vals + r_vals + 1e-15)
+            if scores.ndim == 0:
+                return {target_class_names[0]: scores.item()}
+            return dict(zip(target_class_names, scores.tolist()))
+
+        return MetricsLambda(_fbeta_with_class_names, precision, recall)
 
     fbeta = (1.0 + beta**2) * precision * recall / (beta**2 * precision + recall + 1e-15)
 

ignite/metrics/precision.py

@@ -23,7 +23,17 @@ def __init__(
         is_multilabel: bool = False,
         device: str | torch.device = torch.device("cpu"),
         skip_unrolling: bool = False,
+        class_names: list[str] | None = None,
     ):
+        if class_names is not None:
+            if not isinstance(class_names, (list, tuple)) or not all(isinstance(n, str) for n in class_names):
+                raise ValueError("class_names must be a list of strings")
+            if average is not False and average is not None:
+                raise ValueError(
+                    f"class_names is only applicable when average=False or average=None, got average={average!r}."
+                )
+        self._class_names = class_names
+
         if not (average is None or isinstance(average, bool) or average in ["macro", "micro", "weighted", "samples"]):
             raise ValueError(
                 "Argument average should be None or a boolean or one of values"
@@ -125,7 +135,7 @@ def reset(self) -> None:
         super().reset()
 
     @sync_all_reduce("_numerator", "_denominator")
-    def compute(self) -> torch.Tensor | float:
+    def compute(self) -> torch.Tensor | float | dict:
         r"""
         Return value of the metric for `average` options `'weighted'` and `'macro'` is computed as follows.
 
@@ -157,6 +167,10 @@ def compute(self) -> torch.Tensor | float:
         elif self._average == "macro":
             return cast(torch.Tensor, fraction).mean().item()
         else:
+            if self._class_names is not None:
+                if isinstance(fraction, torch.Tensor) and fraction.ndim == 0:
+                    return {self._class_names[0]: fraction.item()}
+                return dict(zip(self._class_names, cast(torch.Tensor, fraction).tolist()))
             return fraction
 
 
@@ -246,6 +260,10 @@ class Precision(_BasePrecisionRecall):
         skip_unrolling: specifies whether output should be unrolled before being fed to update method. Should be
             true for multi-output model, for example, if ``y_pred`` contains multi-output as ``(y_pred_a, y_pred_b)``
             Alternatively, ``output_transform`` can be used to handle this.
+        class_names: list of class name strings used to label per-class output when ``average=False``
+            or ``average=None``. If provided, ``compute()`` returns a ``dict``  mapping each class
+            name to its metric value instead of a tensor. Must match the number of classes inferred
+            from the data. Default: ``None``.
 
     Examples:
 
@@ -428,5 +446,11 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
 
             if self._average == "weighted":
                 self._weight += y.sum(dim=0)
-
+            if self._class_names is not None:
+                num_classes = 1 if self._numerator.ndim == 0 else self._numerator.shape[0]
+                if len(self._class_names) != num_classes:
+                    raise ValueError(
+                        f"class_names has {len(self._class_names)} entries but the metric computed "
+                        f"{num_classes} classes."
+                    )
         self._updated = True

ignite/metrics/recall.py

@@ -97,7 +97,10 @@ class Recall(_BasePrecisionRecall):
         skip_unrolling: specifies whether output should be unrolled before being fed to update method. Should be
             true for multi-output model, for example, if ``y_pred`` contains multi-output as ``(y_pred_a, y_pred_b)``
             Alternatively, ``output_transform`` can be used to handle this.
-
+        class_names: list of class name strings used to label per-class output when ``average=False``
+            or ``average=None``. If provided, ``compute()`` returns a ``dict`` mapping each class
+            name to its metric value instead of a tensor. Must match the number of classes inferred
+            from the data. Default: ``None``.
     Examples:
 
         For more information on how metric works with :class:`~ignite.engine.engine.Engine`, visit :ref:`attach-engine`.
@@ -240,5 +243,12 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
 
             if self._average == "weighted":
                 self._weight += y.sum(dim=0)
+            if self._class_names is not None:
+                num_classes = 1 if self._numerator.ndim == 0 else self._numerator.shape[0]
+                if len(self._class_names) != num_classes:
+                    raise ValueError(
+                        f"class_names has {len(self._class_names)} entries but the metric computed "
+                        f"{num_classes} classes."
+                    )
 
         self._updated = True

tests/ignite/metrics/test_fbeta.py

@@ -15,11 +15,11 @@ def test_wrong_inputs():
     with pytest.raises(ValueError, match=r"Beta should be a positive integer"):
         Fbeta(0.0)
 
-    with pytest.raises(ValueError, match=r"Input precision metric should have average=False"):
+    with pytest.raises(ValueError, match=r"Input precision and recall metrics should have average=False"):
         p = Precision(average="micro")
         Fbeta(1.0, precision=p)
 
-    with pytest.raises(ValueError, match=r"Input recall metric should have average=False"):
+    with pytest.raises(ValueError, match=r"Input precision and recall metrics should have average=False"):
         r = Recall(average="samples")
         Fbeta(1.0, recall=r)
 
@@ -233,3 +233,58 @@ def test_multinode_distrib_gloo_cpu_or_gpu(distributed_context_multi_node_gloo):
 def test_multinode_distrib_nccl_gpu(distributed_context_multi_node_nccl):
     device = idist.device()
     _test_distrib_integration(device)
+
+
+def test_class_names():
+    # Invalid class_names type
+    with pytest.raises(ValueError, match="class_names must be a list of strings"):
+        Fbeta(beta=1.0, average=False, class_names=[1, 2])
+
+    # Early check for average=True when class_names is given
+    with pytest.raises(ValueError, match="class_names is only applicable when average=False or average=None"):
+        Fbeta(beta=1.0, average=True, class_names=["cat", "dog"])
+
+    # Precision metric without class_names passed to Fbeta with class_names
+    p_no_cn = Precision(average=False)
+    with pytest.raises(ValueError, match="precision and recall metric class_names must match Fbeta class_names"):
+        Fbeta(beta=1.0, average=False, class_names=["cat", "dog"], precision=p_no_cn)
+
+    # Recall metric without class_names passed to Fbeta with class_names
+    r_no_cn = Recall(average=False)
+    with pytest.raises(ValueError, match="precision and recall metric class_names must match Fbeta class_names"):
+        Fbeta(beta=1.0, average=False, class_names=["cat", "dog"], recall=r_no_cn)
+
+    # Mismatched precision and recall class_names
+    p = Precision(average=False, class_names=["cat", "dog"])
+    r = Recall(average=False, class_names=["a", "b"])
+    with pytest.raises(ValueError, match="precision and recall class_names must match"):
+        Fbeta(beta=1.0, average=False, precision=p, recall=r)
+
+    # Input precision metric with average != False
+    p_avg = Precision(average="macro")
+    with pytest.raises(ValueError, match="Input precision and recall metrics should have average=False"):
+        Fbeta(beta=1.0, average=False, precision=p_avg)
+
+    # Correct computation passing class_names directly to Fbeta
+    f1 = Fbeta(beta=1.0, average=False, class_names=["cat", "dog", "bird"])
+    y_true = torch.tensor([0, 1, 2])
+    y_pred = torch.tensor(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    f1.update((y_pred, y_true))
+    res = f1.compute()
+    assert isinstance(res, dict)
+    assert res == {"cat": 1.0, "dog": 1.0, "bird": 1.0}
+
+    # Correct computation with Precision and Recall having class_names
+    p2 = Precision(average=False, class_names=["cat", "dog", "bird"])
+    r2 = Recall(average=False, class_names=["cat", "dog", "bird"])
+    f2 = Fbeta(beta=1.0, average=False, precision=p2, recall=r2)
+    f2.update((y_pred, y_true))
+    res2 = f2.compute()
+    assert isinstance(res2, dict)
+    assert res2 == {"cat": 1.0, "dog": 1.0, "bird": 1.0}

tests/ignite/metrics/test_precision.py

@@ -325,6 +325,86 @@ def test_incorrect_y_classes(average):
     assert pr._updated is False
 
 
+@pytest.mark.parametrize(
+    "invalid_class_names",
+    [
+        [0, 1, 2],  # list of ints
+        "cat",  # string instead of list
+        [0.1, 0.2],  # list of floats
+        ["cat", 1],  # mixed
+    ],
+)
+def test_class_names_invalid_type(invalid_class_names):
+    with pytest.raises(ValueError, match="class_names must be a list of strings"):
+        Precision(average=False, class_names=invalid_class_names)
+
+
+@pytest.mark.parametrize("average", ["macro", "micro", "weighted", "samples", True])
+def test_class_names_incompatible_average(average):
+    with pytest.raises(ValueError, match="class_names is only applicable when average=False or average=None"):
+        Precision(average=average, class_names=["cat", "dog", "horse"])
+
+
+@pytest.mark.parametrize("average", [False, None])
+def test_class_names_multiclass(average):
+    class_names = ["cat", "dog", "horse"]
+    pr = Precision(average=average, class_names=class_names)
+
+    y_pred = torch.tensor(
+        [
+            [0.0266, 0.1719, 0.3055],
+            [0.6886, 0.3978, 0.8176],
+            [0.9230, 0.0197, 0.8395],
+            [0.1785, 0.2670, 0.6084],
+            [0.8448, 0.7177, 0.7288],
+        ]
+    )
+    y = torch.tensor([2, 0, 2, 1, 0])
+
+    pr.update((y_pred, y))
+    result = pr.compute()
+
+    assert isinstance(result, dict)
+    assert list(result.keys()) == class_names
+    assert result == pytest.approx({"cat": 0.5, "dog": 0.0, "horse": 0.3333333333333333})
+
+
+def test_class_names_length_mismatch():
+    pr = Precision(average=False, class_names=["cat", "dog"])
+
+    y_pred = torch.tensor(
+        [
+            [0.0266, 0.1719, 0.3055],
+            [0.6886, 0.3978, 0.8176],
+            [0.9230, 0.0197, 0.8395],
+        ]
+    )
+    y = torch.tensor([2, 0, 1])
+
+    with pytest.raises(ValueError, match="class_names has 2 entries but the metric computed 3 classes"):
+        pr.update((y_pred, y))
+
+
+def test_class_names_none_returns_tensor():
+    pr = Precision(average=False)
+
+    y_pred = torch.tensor(
+        [
+            [0.0266, 0.1719, 0.3055],
+            [0.6886, 0.3978, 0.8176],
+            [0.9230, 0.0197, 0.8395],
+            [0.1785, 0.2670, 0.6084],
+            [0.8448, 0.7177, 0.7288],
+        ]
+    )
+    y = torch.tensor([2, 0, 2, 1, 0])
+
+    pr.update((y_pred, y))
+    result = pr.compute()
+
+    assert isinstance(result, torch.Tensor)
+
+
 @pytest.mark.usefixtures("distributed")
 class TestDistributed:
     @pytest.mark.parametrize("average", [False, "macro", "weighted", "micro"])
@@ -503,3 +583,33 @@ def test_multilabel_accumulator_device(self, average):
             if average == "weighted":
                 assert pr._weight.device == metric_device, f"{type(pr._weight.device)}:{pr._weight.device} vs "
                 f"{type(metric_device)}:{metric_device}"
+
+
+def test_class_names():
+    # Invalid class_names type
+    with pytest.raises(ValueError, match="class_names must be a list of strings"):
+        Precision(average=False, class_names=[1, 2])
+
+    # Incompatible average mode
+    with pytest.raises(ValueError, match="class_names is only applicable when average=False or average=None"):
+        Precision(average="macro", class_names=["cat", "dog"])
+
+    # Correct computation returning dict
+    pr = Precision(average=False, class_names=["cat", "dog", "bird"])
+    y_true = torch.tensor([0, 1, 2])
+    y_pred = torch.tensor(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    pr.update((y_pred, y_true))
+    res = pr.compute()
+    assert isinstance(res, dict)
+    assert res == {"cat": 1.0, "dog": 1.0, "bird": 1.0}
+
+    # Class names length mismatch
+    pr_mismatch = Precision(average=False, class_names=["cat", "dog"])
+    with pytest.raises(ValueError, match="class_names has 2 entries but the metric computed 3 classes."):
+        pr_mismatch.update((y_pred, y_true))

tests/ignite/metrics/test_recall.py

@@ -505,3 +505,33 @@ def test_multilabel_accumulator_device(self, average):
             if average == "weighted":
                 assert re._weight.device == metric_device, f"{type(re._weight.device)}:{re._weight.device} vs "
                 f"{type(metric_device)}:{metric_device}"
+
+
+def test_class_names():
+    # Invalid class_names type
+    with pytest.raises(ValueError, match="class_names must be a list of strings"):
+        Recall(average=False, class_names=[1, 2])
+
+    # Incompatible average mode
+    with pytest.raises(ValueError, match="class_names is only applicable when average=False or average=None"):
+        Recall(average="macro", class_names=["cat", "dog"])
+
+    # Correct computation returning dict
+    re = Recall(average=False, class_names=["cat", "dog", "bird"])
+    y_true = torch.tensor([0, 1, 2])
+    y_pred = torch.tensor(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    re.update((y_pred, y_true))
+    res = re.compute()
+    assert isinstance(res, dict)
+    assert res == {"cat": 1.0, "dog": 1.0, "bird": 1.0}
+
+    # Class names length mismatch
+    re_mismatch = Recall(average=False, class_names=["cat", "dog"])
+    with pytest.raises(ValueError, match="class_names has 2 entries but the metric computed 3 classes."):
+        re_mismatch.update((y_pred, y_true))