Repository: Project-MONAI/MONAI

Issue:
Port generalized Dice loss
a subtask of #84, port a generalized dice loss function into Monai


Discussion:

Resolving patch:
diff --git a/monai/losses/dice.py b/monai/losses/dice.py
--- a/monai/losses/dice.py
+++ b/monai/losses/dice.py
@@ -9,6 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
 import torch
 from torch.nn.modules.loss import _Loss
 
@@ -40,10 +42,87 @@ def __init__(self, include_background=True, do_sigmoid=False, do_softmax=False):
         """
         super().__init__()
         self.include_background = include_background
+        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
+
+    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()
+        if self.do_sigmoid:
+            psum = psum.sigmoid()  # use sigmoid activation
+        if pred.shape[1] == 1:
+            if self.do_softmax:
+                raise ValueError('do_softmax is not compatible with single channel prediction.')
+            if not self.include_background:
+                warnings.warn('single channel prediction, `include_background=False` ignored.')
+            tsum = ground
+        else:  # multiclass dice loss
+            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
+            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))
+
+        batchsize, n_classes = tsum.shape[:2]
+        tsum = tsum.float().view(batchsize, n_classes, -1)
+        psum = psum.view(batchsize, n_classes, -1)
+
+        intersection = psum * tsum
+        sums = psum + tsum
+
+        score = 2.0 * (intersection.sum(2) + smooth) / (sums.sum(2) + smooth)
+        return 1 - score.mean()
+
+
+@alias("generalized_dice", "generalised_dice")
+class GeneralizedDiceLoss(_Loss):
+    """
+    Compute the generalised Dice loss defined in:
+        Sudre, C. et. al. (2017) Generalised Dice overlap as a deep learning
+        loss function for highly unbalanced segmentations. DLMIA 2017.
+
+    Adapted from:
+        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'):
+        """
+        Args:
+            include_background (bool): If False channel index 0 (background category) is excluded from the calculation.
+            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
+        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)
+        elif w_type == 'square':
+            self.w_func = lambda x: torch.reciprocal(x * x)
+        else:
+            raise ValueError('unknown option for `w_type`: {}'.format(w_type))
+
     def forward(self, pred, ground, smooth=1e-5):
+        """
+        Args:
+            pred (tensor): the shape should be BNH[WD].
+            ground (tensor): the shape should be B1H[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))
 
@@ -54,12 +133,10 @@ def forward(self, pred, ground, smooth=1e-5):
             if self.do_softmax:
                 raise ValueError('do_softmax is not compatible with single channel prediction.')
             if not self.include_background:
-                raise RuntimeWarning('single channel prediction, `include_background=False` ignored.')
+                warnings.warn('single channel prediction, `include_background=False` ignored.')
             tsum = ground
         else:  # multiclass dice loss
             if self.do_softmax:
-                if self.do_sigmoid:
-                    raise ValueError('do_sigmoid=True and do_softmax=Ture are not compatible.')
                 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
@@ -69,12 +146,18 @@ def forward(self, pred, ground, smooth=1e-5):
         assert tsum.shape == psum.shape, ("Ground truth one-hot has differing shape (%r) from source (%r)" %
                                           (tsum.shape, psum.shape))
 
-        batchsize = ground.size(0)
-        tsum = tsum.float().view(batchsize, -1)
-        psum = psum.view(batchsize, -1)
+        batchsize, n_classes = tsum.shape[:2]
+        tsum = tsum.float().view(batchsize, n_classes, -1)
+        psum = psum.view(batchsize, n_classes, -1)
 
         intersection = psum * tsum
         sums = psum + tsum
 
-        score = 2.0 * (intersection.sum(1) + smooth) / (sums.sum(1) + smooth)
-        return 1 - score.sum() / batchsize
+        w = self.w_func(tsum.sum(2))
+        for b in w:
+            infs = torch.isinf(b)
+            b[infs] = 0.0
+            b[infs] = torch.max(b)
+
+        score = 2.0 * (intersection.sum(2) * w) / (sums.sum(2) * w + smooth)
+        return 1 - score.mean()


Test patch:
diff --git a/tests/test_dice_loss.py b/tests/test_dice_loss.py
--- a/tests/test_dice_loss.py
+++ b/tests/test_dice_loss.py
@@ -64,7 +64,7 @@
         'ground': torch.tensor([[[1., 0., 0.]], [[1., 1., 0.]]]),
         'smooth': 1e-4,
     },
-    0.422957,
+    0.435015,
 ]
 
 TEST_CASE_5 = [  # shape: (2, 2, 3), (2, 1, 3)
@@ -77,13 +77,26 @@
         'ground': torch.tensor([[[1., 0., 0.]], [[1., 1., 0.]]]),
         'smooth': 1e-4,
     },
-    0.373045,
+    0.383678,
+]
+
+TEST_CASE_6 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
+    {
+        'include_background': False,
+        'do_sigmoid': True,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 0.], [1., 1.]]]]),
+        'smooth': 1e-6,
+    },
+    0.307576,
 ]
 
 
 class TestDiceLoss(unittest.TestCase):
 
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6])
     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
new file mode 100644
--- /dev/null
+++ b/tests/test_generalized_dice_loss.py
@@ -0,0 +1,118 @@
+# Copyright 2020 MONAI Consortium
+# 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.
+
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.losses.dice import GeneralizedDiceLoss
+
+TEST_CASE_0 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
+    {
+        'include_background': True,
+        'do_sigmoid': True,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 0.], [1., 1.]]]]),
+        'smooth': 1e-6,
+    },
+    0.307576,
+]
+
+TEST_CASE_1 = [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
+    {
+        'include_background': True,
+        'do_sigmoid': True,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]], [[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 1.], [1., 1.]]], [[[1., 0.], [1., 0.]]]]),
+        'smooth': 1e-4,
+    },
+    0.41678,
+]
+
+TEST_CASE_2 = [  # shape: (2, 2, 3), (2, 1, 3)
+    {
+        'include_background': False,
+    },
+    {
+        'pred': torch.tensor([[[1., 1., 0.], [0., 0., 1.]], [[1., 0., 1.], [0., 1., 0.]]]),
+        'ground': torch.tensor([[[0., 0., 1.]], [[0., 1., 0.]]]),
+        'smooth': 0.0,
+    },
+    0.0,
+]
+
+TEST_CASE_3 = [  # shape: (2, 2, 3), (2, 1, 3)
+    {
+        'include_background': True,
+        'do_sigmoid': True,
+    },
+    {
+        'pred': torch.tensor([[[-1., 0., 1.], [1., 0., -1.]], [[0., 0., 0.], [0., 0., 0.]]]),
+        'ground': torch.tensor([[[1., 0., 0.]], [[1., 1., 0.]]]),
+        'smooth': 1e-4,
+    },
+    0.435111,
+]
+
+TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
+    {
+        'include_background': True,
+        'do_softmax': True,
+    },
+    {
+        'pred': torch.tensor([[[-1., 0., 1.], [1., 0., -1.]], [[0., 0., 0.], [0., 0., 0.]]]),
+        'ground': torch.tensor([[[1., 0., 0.]], [[1., 1., 0.]]]),
+        'smooth': 1e-4,
+    },
+    0.383776,
+]
+
+TEST_CASE_5 = [  # shape: (2, 2, 3), (2, 1, 3)
+    {
+        'include_background': False,
+    },
+    {
+        'pred': torch.tensor([[[1., 1., 0.], [0., 0., 1.]], [[1., 0., 1.], [0., 1., 0.]]]),
+        'ground': torch.tensor([[[0., 0., 0.]], [[0., 0., 0.]]]),
+        'smooth': 1e-8,
+    },
+    1.0,
+]
+
+TEST_CASE_6 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
+    {
+        'include_background': False,
+        'do_sigmoid': True,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 0.], [1., 1.]]]]),
+        'smooth': 1e-6,
+    },
+    0.307576,
+]
+
+
+class TestGeneralizedDiceLoss(unittest.TestCase):
+
+    @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6])
+    def test_shape(self, input_param, input_data, expected_val):
+        result = GeneralizedDiceLoss(**input_param).forward(**input_data)
+        self.assertAlmostEqual(result.item(), expected_val, places=5)
+
+
+if __name__ == '__main__':
+    unittest.main()