Repository: Project-MONAI/MONAI

Issue:
dice loss support multi-labels task
**Is your feature request related to a problem? Please describe.**
The current implementation of DiceLoss can only support Binary and Multi-classes tasks.
Need to also support Multi-labels task.


Discussion:
Nic-Ma: Hi @wyli ,

Thanks for your review.
I update the code according to your suggestions.
Could you please help review it again?
Thanks.


Patch:
diff --git a/monai/losses/dice.py b/monai/losses/dice.py
index 10dc15ad..8e8ff20f 100644
--- a/monai/losses/dice.py
+++ b/monai/losses/dice.py
@@ -23,63 +23,93 @@ from monai.utils.aliases import alias
 @alias("dice", "Dice")
 class DiceLoss(_Loss):
     """
-    Multiclass dice loss. Input logits `pred` (BNHW[D] where N is number of classes) is compared with ground truth
-    `ground' (B1HW[D]). Axis N of `pred` is expected to have logit predictions for each class rather than being image
-    channels, while the same axis of `ground` should be 1. If the N channel of `pred` is 1 binary dice loss will be
-    calculated. The `smooth` parameter is a value added to the intersection and union components of the inter-over-union
-    calculation to smooth results and prevent divide-by-0, this value should be small. The `include_background` class
-    attribute can be set to False for an instance of DiceLoss to exclude the first category (channel index 0) which is
-    by convention assumed to be background. If the non-background segmentations are small compared to the total image
-    size they can get overwhelmed by the signal from the background so excluding it in such cases helps convergence.
+    Compute average Dice loss between two tensors. It can support both multi-classes and multi-labels tasks.
+    Input logits `pred` (BNHW[D] where N is number of classes) is compared with ground truth `ground' (BNHW[D]).
+    Axis N of `pred` is expected to have logit predictions for each class rather than being image channels,
+    while the same axis of `ground` can be 1 or N(one-hot format). The `smooth` parameter is a value added to the
+    intersection and union components of the inter-over-union calculation to smooth results and prevent divide by 0,
+    this value should be small. The `include_background` class attribute can be set to False for an instance of
+    DiceLoss to exclude the first category (channel index 0) which is by convention assumed to be background.
+    If the non-background segmentations are small compared to the total image size they can get overwhelmed by
+    the signal from the background so excluding it in such cases helps convergence.
     """
 
-    def __init__(self, include_background=True, do_sigmoid=False, do_softmax=False):
+    def __init__(
+        self,
+        include_background=True,
+        to_onehot_y=False,
+        do_sigmoid=False,
+        do_softmax=False,
+        squared_pred=False,
+        jaccard=False
+    ):
         """
         Args:
             include_background (bool): If False channel index 0 (background category) is excluded from the calculation.
+            to_onehot_y (bool): whether to convert `y` into the one-hot format. Defaults to False.
             do_sigmoid (bool): If True, apply a sigmoid function to the prediction.
             do_softmax (bool): If True, apply a softmax function to the prediction.
+            squared_pred (bool): use squared versions of targets and predictions in the denominator or not.
+            jaccard (bool): compute Jaccard Index (soft IoU) instead of dice or not.
+
         """
         super().__init__()
         self.include_background = include_background
+        self.to_onehot_y = to_onehot_y
         if do_sigmoid and do_softmax:
             raise ValueError('do_sigmoid=True and do_softmax=Ture are not compatible.')
         self.do_sigmoid = do_sigmoid
         self.do_softmax = do_softmax
+        self.squared_pred = squared_pred
+        self.jaccard = jaccard
 
     def forward(self, pred, ground, smooth=1e-5):
-        if ground.shape[1] != 1:
-            raise ValueError("Ground truth should have only a single channel, shape is " + str(ground.shape))
-
-        psum = pred.float()
+        """
+        Args:
+            pred (tensor): the shape should be BNH[WD].
+            ground (tensor): the shape should be BNH[WD].
+            smooth (float): a small constant to avoid nan.
+        """
         if self.do_sigmoid:
-            psum = psum.sigmoid()  # use sigmoid activation
-        if pred.shape[1] == 1:
+            pred = torch.sigmoid(pred)
+        n_pred_ch = pred.shape[1]
+        if n_pred_ch == 1:
             if self.do_softmax:
-                raise ValueError('do_softmax is not compatible with single channel prediction.')
+                warnings.warn('single channel prediction, `do_softmax=True` ignored.')
+            if self.to_onehot_y:
+                warnings.warn('single channel prediction, `to_onehot_y=True` ignored.')
             if not self.include_background:
                 warnings.warn('single channel prediction, `include_background=False` ignored.')
-            tsum = ground
-        else:  # multiclass dice loss
+        else:
             if self.do_softmax:
-                psum = torch.softmax(pred, 1)
-            tsum = one_hot(ground, pred.shape[1])  # B1HW(D) -> BNHW(D)
-            # exclude background category so that it doesn't overwhelm the other segmentations if they are small
+                pred = torch.softmax(pred, 1)
+            if self.to_onehot_y:
+                ground = one_hot(ground, n_pred_ch)
             if not self.include_background:
-                tsum = tsum[:, 1:]
-                psum = psum[:, 1:]
-        assert tsum.shape == psum.shape, ("Ground truth one-hot has differing shape (%r) from source (%r)" %
-                                          (tsum.shape, psum.shape))
+                # if skipping background, removing first channel
+                ground = ground[:, 1:]
+                pred = pred[:, 1:]
+                assert ground.shape == pred.shape, ('ground truth one-hot has differing shape (%r) from pred (%r)' %
+                                                    (ground.shape, pred.shape))
 
-        batchsize, n_classes = tsum.shape[:2]
-        tsum = tsum.float().view(batchsize, n_classes, -1)
-        psum = psum.view(batchsize, n_classes, -1)
+        # reducing only spatial dimensions (not batch nor channels)
+        reduce_axis = list(range(2, len(pred.shape)))
+        intersection = torch.sum(ground * pred, reduce_axis)
 
-        intersection = psum * tsum
-        sums = psum + tsum
+        if self.squared_pred:
+            ground = torch.pow(ground, 2)
+            pred = torch.pow(pred, 2)
 
-        score = (2.0 * intersection.sum(2) + smooth) / (sums.sum(2) + smooth)
-        return 1 - score.mean()
+        ground_o = torch.sum(ground, reduce_axis)
+        pred_o = torch.sum(pred, reduce_axis)
+
+        denominator = ground_o + pred_o
+
+        if self.jaccard:
+            denominator -= intersection
+
+        f = (2.0 * intersection + smooth) / (denominator + smooth)
+        return 1.0 - f.mean()  # final reduce_mean across batches and channels
 
 
 @alias("generalized_dice", "generalised_dice")
@@ -94,24 +124,33 @@ class GeneralizedDiceLoss(_Loss):
         https://github.com/NifTK/NiftyNet/blob/v0.6.0/niftynet/layer/loss_segmentation.py#L279
     """
 
-    def __init__(self, include_background=True, do_sigmoid=False, do_softmax=False, w_type='square'):
+    def __init__(
+        self,
+        include_background=True,
+        to_onehot_y=False,
+        do_sigmoid=False,
+        do_softmax=False,
+        w_type='square'
+    ):
         """
         Args:
             include_background (bool): If False channel index 0 (background category) is excluded from the calculation.
+            to_onehot_y (bool): whether to convert `y` into the one-hot format. Defaults to False.
             do_sigmoid (bool): If True, apply a sigmoid function to the prediction.
             do_softmax (bool): If True, apply a softmax function to the prediction.
             w_type ('square'|'simple'|'uniform'): type of function to transform ground truth volume to a weight factor.
+
         """
         super().__init__()
         self.include_background = include_background
+        self.to_onehot_y = to_onehot_y
         if do_sigmoid and do_softmax:
             raise ValueError('do_sigmoid=True and do_softmax=Ture are not compatible.')
         self.do_sigmoid = do_sigmoid
         self.do_softmax = do_softmax
-
         self.w_func = torch.ones_like
         if w_type == 'simple':
-            self.w_func = lambda x: torch.reciprocal(x)
+            self.w_func = torch.reciprocal
         elif w_type == 'square':
             self.w_func = lambda x: torch.reciprocal(x * x)
         else:
@@ -121,44 +160,45 @@ class GeneralizedDiceLoss(_Loss):
         """
         Args:
             pred (tensor): the shape should be BNH[WD].
-            ground (tensor): the shape should be B1H[WD].
+            ground (tensor): the shape should be BNH[WD].
             smooth (float): a small constant to avoid nan.
         """
-        if ground.shape[1] != 1:
-            raise ValueError("Ground truth should have only a single channel, shape is " + str(ground.shape))
-
-        psum = pred.float()
         if self.do_sigmoid:
-            psum = psum.sigmoid()  # use sigmoid activation
-        if pred.shape[1] == 1:
+            pred = torch.sigmoid(pred)
+        n_pred_ch = pred.shape[1]
+        if n_pred_ch == 1:
             if self.do_softmax:
-                raise ValueError('do_softmax is not compatible with single channel prediction.')
+                warnings.warn('single channel prediction, `do_softmax=True` ignored.')
+            if self.to_onehot_y:
+                warnings.warn('single channel prediction, `to_onehot_y=True` ignored.')
             if not self.include_background:
                 warnings.warn('single channel prediction, `include_background=False` ignored.')
-            tsum = ground
-        else:  # multiclass dice loss
+        else:
             if self.do_softmax:
-                psum = torch.softmax(pred, 1)
-            tsum = one_hot(ground, pred.shape[1])  # B1HW(D) -> BNHW(D)
-            # exclude background category so that it doesn't overwhelm the other segmentations if they are small
+                pred = torch.softmax(pred, 1)
+            if self.to_onehot_y:
+                ground = one_hot(ground, n_pred_ch)
             if not self.include_background:
-                tsum = tsum[:, 1:]
-                psum = psum[:, 1:]
-        assert tsum.shape == psum.shape, ("Ground truth one-hot has differing shape (%r) from source (%r)" %
-                                          (tsum.shape, psum.shape))
+                # if skipping background, removing first channel
+                ground = ground[:, 1:]
+                pred = pred[:, 1:]
+                assert ground.shape == pred.shape, ('ground truth one-hot has differing shape (%r) from pred (%r)' %
+                                                    (ground.shape, pred.shape))
+
+        # reducing only spatial dimensions (not batch nor channels)
+        reduce_axis = list(range(2, len(pred.shape)))
+        intersection = torch.sum(ground * pred, reduce_axis)
 
-        batchsize, n_classes = tsum.shape[:2]
-        tsum = tsum.float().view(batchsize, n_classes, -1)
-        psum = psum.view(batchsize, n_classes, -1)
+        ground_o = torch.sum(ground, reduce_axis)
+        pred_o = torch.sum(pred, reduce_axis)
 
-        intersection = psum * tsum
-        sums = psum + tsum
+        denominator = ground_o + pred_o
 
-        w = self.w_func(tsum.sum(2))
+        w = self.w_func(ground_o.float())
         for b in w:
             infs = torch.isinf(b)
             b[infs] = 0.0
             b[infs] = torch.max(b)
 
-        score = (2.0 * intersection.sum(2) * w + smooth) / (sums.sum(2) * w + smooth)
-        return 1 - score.mean()
+        f = (2.0 * intersection * w + smooth) / (denominator * w + smooth)
+        return 1.0 - f.mean()  # final reduce_mean across batches and channels
diff --git a/monai/networks/utils.py b/monai/networks/utils.py
index 628e4ea7..9d60cfc1 100644
--- a/monai/networks/utils.py
+++ b/monai/networks/utils.py
@@ -46,12 +46,22 @@ def slice_channels(tensor, *slicevals):
     return tensor[slices]
 
 
-def predict_segmentation(logits):
+def predict_segmentation(logits, mutually_exclusive=False, threshold=0):
     """
-    Given the logits from a network, computing the segmentation by thresholding all values above 0 if `logits` has one
-    channel, or computing the `argmax` along the channel axis otherwise, logits has shape `BCHW[D]`
+    Given the logits from a network, computing the segmentation by thresholding all values above 0
+    if multi-labels task, computing the `argmax` along the channel axis if multi-classes task,
+    logits has shape `BCHW[D]`.
+
+    Args:
+        logits (Tensor): raw data of model output.
+        mutually_exclusive (bool): if True, `logits` will be converted into a binary matrix using
+            a combination of argmax, which is suitable for multi-classes task. Defaults to False.
+        threshold (float): thresholding the prediction values if multi-labels task.
     """
-    if logits.shape[1] == 1:
-        return (logits >= 0).int()  # for binary segmentation threshold on channel 0
+    if not mutually_exclusive:
+        return (logits >= threshold).int()
     else:
-        return logits.argmax(1).unsqueeze(1)  # take the index of the max value along dimension 1
+        if logits.shape[1] == 1:
+            warnings.warn('single channel prediction, `mutually_exclusive=True` ignored, use threshold instead.')
+            return (logits >= threshold).int()
+        return logits.argmax(1, keepdim=True)


Test patch:
diff --git a/tests/test_dice_loss.py b/tests/test_dice_loss.py
index e937185f..c0f12332 100644
--- a/tests/test_dice_loss.py
+++ b/tests/test_dice_loss.py
@@ -45,6 +45,7 @@ TEST_CASE_2 = [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
 TEST_CASE_3 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': False,
+        'to_onehot_y': True,
     },
     {
         'pred': torch.tensor([[[1., 1., 0.], [0., 0., 1.]], [[1., 0., 1.], [0., 1., 0.]]]),
@@ -57,6 +58,7 @@ TEST_CASE_3 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': True,
+        'to_onehot_y': True,
         'do_sigmoid': True,
     },
     {
@@ -70,6 +72,7 @@ TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_5 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': True,
+        'to_onehot_y': True,
         'do_softmax': True,
     },
     {
@@ -93,10 +96,39 @@ TEST_CASE_6 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
     0.307576,
 ]
 
+TEST_CASE_7 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
+    {
+        'include_background': True,
+        'do_sigmoid': True,
+        'squared_pred': True,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 0.], [1., 1.]]]]),
+        'smooth': 1e-5,
+    },
+    0.178337,
+]
+
+TEST_CASE_8 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
+    {
+        'include_background': True,
+        'do_sigmoid': True,
+        'jaccard': True,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 0.], [1., 1.]]]]),
+        'smooth': 1e-5,
+    },
+    -0.059094,
+]
+
 
 class TestDiceLoss(unittest.TestCase):
 
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6])
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4,
+                           TEST_CASE_5, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8])
     def test_shape(self, input_param, input_data, expected_val):
         result = DiceLoss(**input_param).forward(**input_data)
         self.assertAlmostEqual(result.item(), expected_val, places=5)
diff --git a/tests/test_generalized_dice_loss.py b/tests/test_generalized_dice_loss.py
index b2ce9616..a3b328d1 100644
--- a/tests/test_generalized_dice_loss.py
+++ b/tests/test_generalized_dice_loss.py
@@ -45,6 +45,7 @@ TEST_CASE_1 = [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
 TEST_CASE_2 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': False,
+        'to_onehot_y': True,
     },
     {
         'pred': torch.tensor([[[1., 1., 0.], [0., 0., 1.]], [[1., 0., 1.], [0., 1., 0.]]]),
@@ -57,6 +58,7 @@ TEST_CASE_2 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_3 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': True,
+        'to_onehot_y': True,
         'do_sigmoid': True,
     },
     {
@@ -70,6 +72,7 @@ TEST_CASE_3 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': True,
+        'to_onehot_y': True,
         'do_softmax': True,
     },
     {
@@ -83,6 +86,7 @@ TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_5 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': False,
+        'to_onehot_y': True,
     },
     {
         'pred': torch.tensor([[[1., 1., 0.], [0., 0., 1.]], [[1., 0., 1.], [0., 1., 0.]]]),