Dice score as metric

tests/classification/test_dice.py

@@ -11,10 +11,60 @@
 # 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 functools import partial
+
 import pytest
-from torch import tensor
+import torch
+from torch import Tensor, tensor
+
+from tests.classification.inputs import _input_multiclass_prob as _input_mcls_prob
+from tests.helpers import seed_all
+from tests.helpers.testers import MetricTester
+from torchmetrics import Dice
+from torchmetrics.functional import dice, dice_score
+from torchmetrics.functional.classification.dice import _stat_scores
+
+seed_all(42)
+
+
+def _dice_score(
+    preds: Tensor,
+    target: Tensor,
+    background: bool = False,
+    nan_score: float = 0.0,
+) -> Tensor:
+    """
+    Compute dice score from prediction scores.
+    There is no implementation of Dice in sklearn. I used public information about
+    metric: `https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient`
+
+    Args:
+        preds: prediction tensor
+        target: target tensor
+        background: whether to also compute dice for the background
+        nan_score: the value to use for the score if denominator equals zero
+
+    Return:
+        Tensor containing dice score
 
-from torchmetrics.functional import dice_score
+    """
+    num_classes = preds.shape[1]
+    bg_inv = 1 - int(background)
+
+    tp = tensor(0, device=preds.device)
+    fp = tensor(0, device=preds.device)
+    fn = tensor(0, device=preds.device)
+
+    for i in range(bg_inv, num_classes):
+        tp_cls, fp_cls, _, fn_cls, _ = _stat_scores(preds=preds, target=target, class_index=i)
+
+        tp += tp_cls
+        fp += fp_cls
+        fn += fn_cls
+
+    denom = (2 * tp + fp + fn).to(torch.float)
+    score = (2 * tp).to(torch.float) / denom if torch.is_nonzero(denom) else nan_score
+    return score
 
 
 @pytest.mark.parametrize(
@@ -29,3 +79,46 @@
 def test_dice_score(pred, target, expected):
     score = dice_score(tensor(pred), tensor(target))
     assert score == expected
+
+
+@pytest.mark.parametrize(
+    ["pred", "target", "expected"],
+    [
+        ([[0, 0], [1, 1]], [[0, 0], [1, 1]], 1.0),
+        ([[1, 1], [0, 0]], [[0, 0], [1, 1]], 0.0),
+        ([[1, 1], [1, 1]], [[1, 1], [0, 0]], 2 / 3),
+        ([[1, 1], [0, 0]], [[1, 1], [0, 0]], 1.0),
+    ],
+)
+def test_dice(pred, target, expected):
+    score = dice(tensor(pred), tensor(target))
+    assert score == expected
+
+
+@pytest.mark.parametrize(
+    "preds, target",
+    [(_input_mcls_prob.preds, _input_mcls_prob.target)],
+)
+@pytest.mark.parametrize("background", [False, True])
+class TestDice(MetricTester):
+    @pytest.mark.parametrize("ddp", [False])
+    @pytest.mark.parametrize("dist_sync_on_step", [False])
+    def test_dice_class(self, ddp, dist_sync_on_step, preds, target, background):
+        self.run_class_metric_test(
+            ddp=ddp,
+            preds=preds,
+            target=target,
+            metric_class=Dice,
+            sk_metric=partial(_dice_score, background=background),
+            dist_sync_on_step=dist_sync_on_step,
+            metric_args={"background": background},
+        )
+
+    def test_dice_fn(self, preds, target, background):
+        self.run_functional_metric_test(
+            preds,
+            target,
+            metric_functional=dice,
+            sk_metric=partial(_dice_score, background=background),
+            metric_args={"background": background},
+        )

torchmetrics/__init__.py

@@ -33,6 +33,7 @@
     CohenKappa,
     ConfusionMatrix,
     CoverageError,
+    Dice,
     F1Score,
     FBetaScore,
     HammingDistance,

torchmetrics/classification/__init__.py

@@ -21,6 +21,7 @@
 from torchmetrics.classification.calibration_error import CalibrationError  # noqa: F401
 from torchmetrics.classification.cohen_kappa import CohenKappa  # noqa: F401
 from torchmetrics.classification.confusion_matrix import ConfusionMatrix  # noqa: F401
+from torchmetrics.classification.dice import Dice  # noqa: F401
 from torchmetrics.classification.f_beta import F1Score, FBetaScore  # noqa: F401
 from torchmetrics.classification.hamming import HammingDistance  # noqa: F401
 from torchmetrics.classification.hinge import HingeLoss  # noqa: F401

torchmetrics/classification/dice.py

@@ -0,0 +1,163 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 Any, Dict, Optional
+
+from torch import Tensor
+
+from torchmetrics.functional.classification.dice import _dice_compute
+
+from torchmetrics.classification.stat_scores import StatScores  # isort:skip
+
+
+class Dice(StatScores):
+    r"""Computes `Dice`_:
+
+    .. math:: \text{Precision} = \frac{\text{2 * TP}}{\text{2 * TP} + \text{FP} + \text{FN}}
+
+    Where :math:`\text{TP}` and :math:`\text{FP}` represent the number of true positives and
+    false positives respecitively.
+
+    The reduction method (how the precision scores are aggregated) is controlled by the
+    ``average`` parameter, and additionally by the ``mdmc_average`` parameter in the
+    multi-dimensional multi-class case. Accepts all inputs listed in :ref:`pages/classification:input types`.
+
+    Args:
+        background:
+            Whether to also compute dice for the background.
+        num_classes:
+            Number of classes. Necessary for ``'macro'``, ``'weighted'`` and ``None`` average methods.
+        threshold:
+            Threshold for transforming probability or logit predictions to binary (0,1) predictions, in the case
+            of binary or multi-label inputs. Default value of 0.5 corresponds to input being probabilities.
+        zero_division:
+            The value to use for the score if denominator equals zero.
+        average:
+            Defines the reduction that is applied. Should be one of the following:
+
+            - ``'micro'`` [default]: Calculate the metric globally, across all samples and classes.
+            - ``'macro'``: Calculate the metric for each class separately, and average the
+              metrics across classes (with equal weights for each class).
+            - ``'weighted'``: Calculate the metric for each class separately, and average the
+              metrics across classes, weighting each class by its support (``tp + fn``).
+            - ``'none'`` or ``None``: Calculate the metric for each class separately, and return
+              the metric for every class.
+            - ``'samples'``: Calculate the metric for each sample, and average the metrics
+              across samples (with equal weights for each sample).
+
+            .. note:: What is considered a sample in the multi-dimensional multi-class case
+                depends on the value of ``mdmc_average``.
+
+        mdmc_average:
+            Defines how averaging is done for multi-dimensional multi-class inputs (on top of the
+            ``average`` parameter). Should be one of the following:
+
+            - ``None`` [default]: Should be left unchanged if your data is not multi-dimensional
+              multi-class.
+
+            - ``'samplewise'``: In this case, the statistics are computed separately for each
+              sample on the ``N`` axis, and then averaged over samples.
+              The computation for each sample is done by treating the flattened extra axes ``...``
+              (see :ref:`pages/classification:input types`) as the ``N`` dimension within the sample,
+              and computing the metric for the sample based on that.
+
+            - ``'global'``: In this case the ``N`` and ``...`` dimensions of the inputs
+              (see :ref:`pages/classification:input types`) are flattened into a new ``N_X`` sample axis, i.e.
+              the inputs are treated as if they were ``(N_X, C)``.
+              From here on the ``average`` parameter applies as usual.
+
+        ignore_index:
+            Integer specifying a target class to ignore. If given, this class index does not contribute
+            to the returned score, regardless of reduction method. If an index is ignored, and ``average=None``
+            or ``'none'``, the score for the ignored class will be returned as ``nan``.
+
+        top_k:
+            Number of the highest probability or logit score predictions considered finding the correct label,
+            relevant only for (multi-dimensional) multi-class inputs. The
+            default value (``None``) will be interpreted as 1 for these inputs.
+            Should be left at default (``None``) for all other types of inputs.
+
+        multiclass:
+            Used only in certain special cases, where you want to treat inputs as a different type
+            than what they appear to be. See the parameter's
+            :ref:`documentation section <pages/classification:using the multiclass parameter>`
+            for a more detailed explanation and examples.
+
+        kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
+
+    Raises:
+        ValueError:
+            If ``average`` is none of ``"micro"``, ``"macro"``, ``"weighted"``, ``"samples"``, ``"none"``, ``None``.
+
+    Example:
+        >>> import torch
+        >>> from torchmetrics import Dice
+        >>> preds  = torch.tensor([2, 0, 2, 1])
+        >>> target = torch.tensor([1, 1, 2, 0])
+        >>> dice = Dice(average='micro')
+        >>> dice(preds, target)
+        tensor(0.3333)
+
+    """
+    is_differentiable = False
+    higher_is_better = True
+
+    def __init__(
+        self,
+        background: bool = False,
+        zero_division: int = 0,
+        num_classes: Optional[int] = None,
+        threshold: float = 0.5,
+        average: str = "micro",
+        mdmc_average: Optional[str] = None,
+        ignore_index: Optional[int] = None,
+        top_k: Optional[int] = None,
+        multiclass: Optional[bool] = None,
+        **kwargs: Dict[str, Any],
+    ) -> None:
+        allowed_average = ["micro", "macro", "weighted", "samples", "none", None]
+        if average not in allowed_average:
+            raise ValueError(f"The `average` has to be one of {allowed_average}, got {average}.")
+
+        if not background and ignore_index is None:
+            # not compute dice for the background
+            ignore_index = 0
+        elif ignore_index is not None:
+            raise ValueError("When you set `ignore_index`, you have to set background `bg` to True.")
+
+        super().__init__(
+            reduce="macro" if average in ["weighted", "none", None] else average,
+            mdmc_reduce=mdmc_average,
+            threshold=threshold,
+            top_k=top_k,
+            num_classes=num_classes,
+            multiclass=multiclass,
+            ignore_index=ignore_index,
+            **kwargs,
+        )
+
+        self.average = average
+        self.zero_division = zero_division
+
+    def compute(self) -> Tensor:
+        """Computes the dice score based on inputs passed in to ``update`` previously.
+
+        Return:
+            The shape of the returned tensor depends on the ``average`` parameter:
+
+            - If ``average in ['micro', 'macro', 'weighted', 'samples']``, a one-element tensor will be returned
+            - If ``average in ['none', None]``, the shape will be ``(C,)``, where ``C`` stands  for the number
+              of classes
+        """
+        tp, fp, _, fn = self._get_final_stats()
+        return _dice_compute(tp, fp, fn, self.average, self.mdmc_reduce, self.zero_division)

torchmetrics/functional/__init__.py

@@ -21,7 +21,7 @@
 from torchmetrics.functional.classification.calibration_error import calibration_error
 from torchmetrics.functional.classification.cohen_kappa import cohen_kappa
 from torchmetrics.functional.classification.confusion_matrix import confusion_matrix
-from torchmetrics.functional.classification.dice import dice_score
+from torchmetrics.functional.classification.dice import dice, dice_score
 from torchmetrics.functional.classification.f_beta import f1_score, fbeta_score
 from torchmetrics.functional.classification.hamming import hamming_distance
 from torchmetrics.functional.classification.hinge import hinge_loss
@@ -105,6 +105,7 @@
     "coverage_error",
     "tweedie_deviance_score",
     "dice_score",
+    "dice",
     "error_relative_global_dimensionless_synthesis",
     "explained_variance",
     "extended_edit_distance",

torchmetrics/functional/classification/__init__.py

@@ -18,7 +18,7 @@
 from torchmetrics.functional.classification.calibration_error import calibration_error  # noqa: F401
 from torchmetrics.functional.classification.cohen_kappa import cohen_kappa  # noqa: F401
 from torchmetrics.functional.classification.confusion_matrix import confusion_matrix  # noqa: F401
-from torchmetrics.functional.classification.dice import dice_score  # noqa: F401
+from torchmetrics.functional.classification.dice import dice, dice_score  # noqa: F401
 from torchmetrics.functional.classification.f_beta import f1_score, fbeta_score  # noqa: F401
 from torchmetrics.functional.classification.hamming import hamming_distance  # noqa: F401
 from torchmetrics.functional.classification.hinge import hinge_loss  # noqa: F401