Repository: Project-MONAI/MONAI

Issue:
MeanDice return Nan for all zero label
**Describe the bug**
If some channel of label is all zero and prediction is also all zero, the MeanDice will div 0 and return Nan.
Need to skip zero labels when computing the MeanDice and average value.

**To Reproduce**
Steps to reproduce the behavior:
Use all zero predictions and labels to test.

**Expected behavior**
Output Nan.




Discussion:

Resolving patch:
diff --git a/monai/handlers/mean_dice.py b/monai/handlers/mean_dice.py
--- a/monai/handlers/mean_dice.py
+++ b/monai/handlers/mean_dice.py
@@ -9,42 +9,41 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Callable, Optional, Sequence, Union
+
 import torch
-from typing import Callable, Union, Optional, Sequence
 from ignite.exceptions import NotComputableError
 from ignite.metrics import Metric
-from ignite.metrics.metric import sync_all_reduce, reinit__is_reduced
+from ignite.metrics.metric import reinit__is_reduced, sync_all_reduce
+
 from monai.metrics.compute_meandice import compute_meandice
 
-__all__ = [
-    'MeanDice'
-]
+__all__ = ['MeanDice']
 
 
 class MeanDice(Metric):
     """Computes dice score metric from full size Tensor and collects average over batch, class-channels, iterations.
     """
-    def __init__(
-        self,
-        include_background=True,
-        to_onehot_y=True,
-        logit_thresh=None,
-        add_sigmoid=False,
-        mutually_exclusive=True,
-        output_transform: Callable = lambda x: x,
-        device: Optional[Union[str, torch.device]] = None
-    ):
+
+    def __init__(self,
+                 include_background=True,
+                 to_onehot_y=True,
+                 mutually_exclusive=False,
+                 add_sigmoid=False,
+                 logit_thresh=0.5,
+                 output_transform: Callable = lambda x: x,
+                 device: Optional[Union[str, torch.device]] = None):
         """
 
         Args:
             include_background (Bool): whether to include dice computation on the first channel of the predicted output.
                 Defaults to True.
             to_onehot_y (Bool): whether to convert the output prediction into the one-hot format. Defaults to True.
-            logit_thresh (Float): the threshold value to round value to 0.0 and 1.0. Defaults to None (no thresholding).
+            mutually_exclusive (Bool): if True, the output prediction will be converted into a binary matrix using
+                a combination of argmax and to_onehot. Defaults to False.
             add_sigmoid (Bool): whether to add sigmoid function to the output prediction before computing Dice.
                 Defaults to False.
-            mutually_exclusive (Bool): if True, the output prediction will be converted into a binary matrix using
-                a combination of argmax and to_onehot. Defaults to True.
+            logit_thresh (Float): the threshold value to round value to 0.0 and 1.0. Defaults to None (no thresholding).
             output_transform (Callable): transform the ignite.engine.state.output into [y_pred, y] pair.
             device (torch.device): device specification in case of distributed computation usage.
 
@@ -54,9 +53,9 @@ def __init__(
         super(MeanDice, self).__init__(output_transform, device=device)
         self.include_background = include_background
         self.to_onehot_y = to_onehot_y
-        self.logit_thresh = logit_thresh
-        self.add_sigmoid = add_sigmoid
         self.mutually_exclusive = mutually_exclusive
+        self.add_sigmoid = add_sigmoid
+        self.logit_thresh = logit_thresh
 
         self._sum = 0
         self._num_examples = 0
@@ -70,16 +69,20 @@ def reset(self):
     def update(self, output: Sequence[Union[torch.Tensor, dict]]):
         assert len(output) == 2, 'MeanDice metric can only support y_pred and y.'
         y_pred, y = output
-        average = compute_meandice(y_pred, y, self.include_background, self.to_onehot_y, self.mutually_exclusive,
-                                   self.add_sigmoid, self.logit_thresh)
+        scores = compute_meandice(y_pred, y, self.include_background, self.to_onehot_y, self.mutually_exclusive,
+                                  self.add_sigmoid, self.logit_thresh)
 
-        batch_size = len(y)
-        self._sum += average.item() * batch_size
-        self._num_examples += batch_size
+        # add all items in current batch
+        for batch in scores:
+            not_nan = ~torch.isnan(batch)
+            if not_nan.sum() == 0:
+                continue
+            class_avg = batch[not_nan].mean().item()
+            self._sum += class_avg
+            self._num_examples += 1
 
     @sync_all_reduce("_sum", "_num_examples")
     def compute(self):
         if self._num_examples == 0:
-            raise NotComputableError(
-                'MeanDice must have at least one example before it can be computed.')
+            raise NotComputableError('MeanDice must have at least one example before it can be computed.')
         return self._sum / self._num_examples
diff --git a/monai/metrics/compute_meandice.py b/monai/metrics/compute_meandice.py
--- a/monai/metrics/compute_meandice.py
+++ b/monai/metrics/compute_meandice.py
@@ -9,6 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
 import torch
 
 from monai.networks.utils import one_hot
@@ -16,48 +18,65 @@
 
 def compute_meandice(y_pred,
                      y,
-                     include_background=False,
+                     include_background=True,
                      to_onehot_y=True,
-                     mutually_exclusive=True,
+                     mutually_exclusive=False,
                      add_sigmoid=False,
-                     logit_thresh=None):
+                     logit_thresh=0.5):
     """Computes dice score metric from full size Tensor and collects average.
 
     Args:
         y_pred (torch.Tensor): input data to compute, typical segmentation model output.
-                               it must be One-Hot format and first dim is batch, example shape: [16, 3, 32, 32].
+            it must be One-Hot format and first dim is batch, example shape: [16, 3, 32, 32].
         y (torch.Tensor): ground truth to compute mean dice metric, the first dim is batch.
-                          example shape: [16, 3, 32, 32] for 3-class one-hot labels.
-                          alternative shape: [16, 1, 32, 32] and set `to_onehot_y=True` to convert it into [16, 3, 32, 32].
-        include_background (Bool): whether to skip dice computation on the first channel of the predicted output.
-        to_onehot_y (Bool): whether to convert `y` into the one-hot format.
+            example shape: [16, 1, 32, 32] will be converted into [16, 3, 32, 32].
+            alternative shape: [16, 3, 32, 32] and set `to_onehot_y=False` to use 3-class labels directly.
+        include_background (Bool): whether to skip Dice computation on the first channel of
+            the predicted output. Defaults to True.
+        to_onehot_y (Bool): whether to convert `y` into the one-hot format. Defaults to True.
         mutually_exclusive (Bool): if True, `y_pred` will be converted into a binary matrix using
-            a combination of argmax and to_onehot.
-        add_sigmoid (Bool): whether to add sigmoid function to y_pred before computation.
-        logit_thresh (Float): the threshold value used to convert `y_pred` into a binary matrix.
+            a combination of argmax and to_onehot.  Defaults to False.
+        add_sigmoid (Bool): whether to add sigmoid function to y_pred before computation. Defaults to False.
+        logit_thresh (Float): the threshold value used to convert (after sigmoid if `add_sigmoid=True`)
+            `y_pred` into a binary matrix. Defaults to 0.5.
+
+    Returns:
+        Dice scores per batch and per class (shape: [batch_size, n_classes]).
 
     Note:
         This method provide two options to convert `y_pred` into a binary matrix:
-            (1) when `mutually_exclusive` is True, it uses a combination of argmax and to_onehot
+            (1) when `mutually_exclusive` is True, it uses a combination of argmax and to_onehot,
             (2) when `mutually_exclusive` is False, it uses a threshold `logit_thresh`
                 (optionally with a sigmoid function before thresholding).
 
     """
-    n_channels_y_pred = y_pred.shape[1]
-
-    if mutually_exclusive:
-        if logit_thresh is not None or add_sigmoid:
-            raise ValueError('`logit_thresh` and `add_sigmoid` are incompatible when mutually_exclusive is True.')
-        y_pred = torch.argmax(y_pred, dim=1, keepdim=True)
-        y_pred = one_hot(y_pred, n_channels_y_pred)
-    else:  # channel-wise thresholding
-        if add_sigmoid:
-            y_pred = torch.sigmoid(y_pred)
-        if logit_thresh is not None:
-            y_pred = (y_pred >= logit_thresh).float()
+    n_classes = y_pred.shape[1]
+    n_len = len(y_pred.shape)
 
-    if to_onehot_y:
-        y = one_hot(y, n_channels_y_pred)
+    if add_sigmoid:
+        y_pred = y_pred.float().sigmoid()
+
+    if n_classes == 1:
+        if mutually_exclusive:
+            warnings.warn('y_pred has only one class, mutually_exclusive=True ignored.')
+        if to_onehot_y:
+            warnings.warn('y_pred has only one channel, to_onehot_y=True ignored.')
+        if not include_background:
+            warnings.warn('y_pred has only one channel, include_background=False ignored.')
+        # make both y and y_pred binary
+        y_pred = (y_pred >= logit_thresh).float()
+        y = (y > 0).float()
+    else:  # multi-channel y_pred
+        # make both y and y_pred binary
+        if mutually_exclusive:
+            if add_sigmoid:
+                raise ValueError('add_sigmoid=True is incompatible with mutually_exclusive=True.')
+            y_pred = torch.argmax(y_pred, dim=1, keepdim=True)
+            y_pred = one_hot(y_pred, n_classes)
+        else:
+            y_pred = (y_pred >= logit_thresh).float()
+        if to_onehot_y:
+            y = one_hot(y, n_classes)
 
     if not include_background:
         y = y[:, 1:] if y.shape[1] > 1 else y
@@ -67,13 +86,12 @@ def compute_meandice(y_pred,
                                      (y.shape, y_pred.shape))
 
     # reducing only spatial dimensions (not batch nor channels)
-    reduce_axis = list(range(2, y_pred.dim()))
+    reduce_axis = list(range(2, n_len))
     intersection = torch.sum(y * y_pred, reduce_axis)
 
     y_o = torch.sum(y, reduce_axis)
     y_pred_o = torch.sum(y_pred, reduce_axis)
     denominator = y_o + y_pred_o
 
-    f = (2.0 * intersection) / denominator
-    # final reduce_mean across batches and channels
-    return torch.mean(f)
+    f = torch.where(y_o > 0, (2.0 * intersection) / denominator, torch.tensor(float('nan')).to(y_o.float()))
+    return f  # returns array of Dice shape: [Batch, n_classes]


Test patch:
diff --git a/tests/test_compute_meandice.py b/tests/test_compute_meandice.py
--- a/tests/test_compute_meandice.py
+++ b/tests/test_compute_meandice.py
@@ -11,13 +11,14 @@
 
 import unittest
 
+import numpy as np
 import torch
 from parameterized import parameterized
 
 from monai.metrics.compute_meandice import compute_meandice
 
 # keep background
-TEST_CASE_1 = [
+TEST_CASE_1 = [  # y (1, 1, 2, 2), y_pred (1, 1, 2, 2), expected out (1, 1)
     {
         'y_pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
         'y': torch.tensor([[[[1., 0.], [1., 1.]]]]),
@@ -27,17 +28,17 @@
         'logit_thresh': 0.5,
         'add_sigmoid': True,
     },
-    0.8000,
+    [[0.8]],
 ]
 
 # remove background and not One-Hot target
-TEST_CASE_2 = [
+TEST_CASE_2 = [  # y (2, 1, 2, 2), y_pred (2, 3, 2, 2), expected out (2, 2) (no background)
     {
         'y_pred':
         torch.tensor([[[[-1., 3.], [2., -4.]], [[0., -1.], [3., 2.]], [[0., 1.], [2., -1.]]],
                       [[[-2., 0.], [3., 1.]], [[0., 2.], [1., -2.]], [[-1., 2.], [4., 0.]]]]),
         'y':
-        torch.tensor([[[[1, 2], [1, 0]]], [[[1, 1], [2, 0]]]]),
+        torch.tensor([[[[1., 2.], [1., 0.]]], [[[1., 1.], [2., 0.]]]]),
         'include_background':
         False,
         'to_onehot_y':
@@ -45,7 +46,24 @@
         'mutually_exclusive':
         True,
     },
-    0.4583,
+    [[0.5000, 0.0000], [0.6666, 0.6666]],
+]
+
+# should return Nan for all labels=0 case and skip for MeanDice
+TEST_CASE_3 = [
+    {
+        'y_pred':
+        torch.zeros(2, 3, 2, 2),
+        'y':
+        torch.tensor([[[[0., 0.], [0., 0.]]], [[[1., 0.], [0., 1.]]]]),
+        'include_background':
+        True,
+        'to_onehot_y':
+        True,
+        'mutually_exclusive':
+        True,
+    },
+    [[False, True, True], [False, False, True]],
 ]
 
 
@@ -54,7 +72,12 @@ class TestComputeMeanDice(unittest.TestCase):
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
     def test_value(self, input_data, expected_value):
         result = compute_meandice(**input_data)
-        self.assertAlmostEqual(result.item(), expected_value, places=4)
+        self.assertTrue(np.allclose(result.cpu().numpy(), expected_value, atol=1e-4))
+
+    @parameterized.expand([TEST_CASE_3])
+    def test_nans(self, input_data, expected_value):
+        result = compute_meandice(**input_data)
+        self.assertTrue(np.allclose(np.isnan(result.cpu().numpy()), expected_value))
 
 
 if __name__ == '__main__':
diff --git a/tests/test_handler_mean_dice.py b/tests/test_handler_mean_dice.py
--- a/tests/test_handler_mean_dice.py
+++ b/tests/test_handler_mean_dice.py
@@ -17,8 +17,7 @@
 from monai.handlers.mean_dice import MeanDice
 
 TEST_CASE_1 = [{'to_onehot_y': True, 'mutually_exclusive': True}, 0.75]
-TEST_CASE_2 = [{'include_background': False, 'to_onehot_y': False, 'mutually_exclusive': False}, 0.8333333]
-
+TEST_CASE_2 = [{'include_background': False, 'to_onehot_y': True, 'mutually_exclusive': False}, 0.66666]
 TEST_CASE_3 = [{'mutually_exclusive': True, 'add_sigmoid': True}]
 
 
@@ -29,16 +28,16 @@ class TestHandlerMeanDice(unittest.TestCase):
     def test_compute(self, input_params, expected_avg):
         dice_metric = MeanDice(**input_params)
 
-        y_pred = torch.Tensor([[0, 1], [1, 0]])
-        y = torch.ones((2, 1))
+        y_pred = torch.Tensor([[[0], [1]], [[1], [0]]])
+        y = torch.ones((2, 1, 1))
         dice_metric.update([y_pred, y])
 
-        y_pred = torch.Tensor([[0, 1], [1, 0]])
-        y = torch.Tensor([[1.], [0.]])
+        y_pred = torch.Tensor([[[0], [1]], [[1], [0]]])
+        y = torch.Tensor([[[1]], [[0]]])
         dice_metric.update([y_pred, y])
 
         avg_dice = dice_metric.compute()
-        self.assertAlmostEqual(avg_dice, expected_avg)
+        self.assertAlmostEqual(avg_dice, expected_avg, places=4)
 
     @parameterized.expand([TEST_CASE_3])
     def test_misconfig(self, input_params):