Repository: Project-MONAI/MONAI

Issue:
Port the Tversky loss
a subtask of #84, port tversky loss function into Monai

https://arxiv.org/abs/1706.05721
~https://weisslab.cs.ucl.ac.uk/FernandoPerezGarcia/episurg/blob/master/episurg/loss.py~
https://niftynet.readthedocs.io/en/dev/niftynet.layer.loss_segmentation.html#niftynet.layer.loss_segmentation.tversky


Discussion:

Resolving patch:
diff --git a/monai/losses/__init__.py b/monai/losses/__init__.py
--- a/monai/losses/__init__.py
+++ b/monai/losses/__init__.py
@@ -10,3 +10,4 @@
 # limitations under the License.
 
 from .dice import *
+from .tversky import *
diff --git a/monai/losses/tversky.py b/monai/losses/tversky.py
new file mode 100644
--- /dev/null
+++ b/monai/losses/tversky.py
@@ -0,0 +1,110 @@
+# 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 warnings
+
+import torch
+from torch.nn.modules.loss import _Loss
+
+from monai.networks.utils import one_hot
+
+
+class TverskyLoss(_Loss):
+
+    """
+    Compute the Tversky loss defined in:        
+
+        Sadegh et al. (2017) Tversky loss function for image segmentation
+        using 3D fully convolutional deep networks. (https://arxiv.org/abs/1706.05721)
+
+    Adapted from:
+        https://github.com/NifTK/NiftyNet/blob/v0.6.0/niftynet/layer/loss_segmentation.py#L631
+
+    """
+
+    def __init__(
+        self,
+        include_background=True,
+        to_onehot_y=False,
+        do_sigmoid=False,
+        do_softmax=False,
+        alpha=0.5,
+        beta=0.5
+    ):
+
+        """
+        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.
+            alpha (float): weight of false positives
+            beta  (float): weight of false negatives
+        """
+
+        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=True are not compatible.')
+        self.do_sigmoid = do_sigmoid
+        self.do_softmax = do_softmax
+        self.alpha = alpha
+        self.beta = beta
+
+    def forward(self, pred, ground, smooth=1e-5):
+        """
+        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:
+            pred = torch.sigmoid(pred)
+        n_pred_ch = pred.shape[1]
+        if n_pred_ch == 1:
+            if self.do_softmax:
+                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.')
+        else:
+            if self.do_softmax:
+                pred = torch.softmax(pred, 1)
+            if self.to_onehot_y:
+                ground = one_hot(ground, n_pred_ch)
+            if not self.include_background:
+                # 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))
+
+        p0 = pred
+        p1 = 1 - p0
+        g0 = ground
+        g1 = 1 - g0
+
+        # reducing only spatial dimensions (not batch nor channels)
+        reduce_axis = list(range(2, len(pred.shape)))
+
+        tp = torch.sum(p0 * g0 , reduce_axis)
+        fp = self.alpha * torch.sum(p0 * g1 , reduce_axis)
+        fn = self.beta * torch.sum(p1 * g0 , reduce_axis)
+
+        numerator = tp + smooth
+        denominator = tp + fp + fn + smooth
+
+        score = numerator / denominator 
+
+        return 1.0 - score.mean()  


Test patch:
diff --git a/tests/test_tversky_loss.py b/tests/test_tversky_loss.py
new file mode 100644
--- /dev/null
+++ b/tests/test_tversky_loss.py
@@ -0,0 +1,128 @@
+# 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 import TverskyLoss
+
+TEST_CASE_1 = [  # 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_2 = [  # 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.416657,
+]
+
+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.]]]),
+        'ground': torch.tensor([[[0., 0., 1.]], [[0., 1., 0.]]]),
+        'smooth': 0.0,
+    },
+    0.0,
+]
+
+TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
+    {
+        'include_background': True,
+        'to_onehot_y': 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.435050,
+]
+
+TEST_CASE_5 = [  # shape: (2, 2, 3), (2, 1, 3)
+    {
+        'include_background': True,
+        'to_onehot_y': 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.383713,
+]
+
+TEST_CASE_6 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
+    {
+        'include_background': True,
+        'do_sigmoid': True,
+        'alpha': 0.3,
+        'beta': 0.7,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 0.], [1., 1.]]]]),
+        'smooth': 1e-6,
+    },
+    0.3589,
+]
+
+TEST_CASE_7 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
+    {
+        'include_background': True,
+        'do_sigmoid': True,
+        'alpha': 0.7,
+        'beta': 0.3,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 0.], [1., 1.]]]]),
+        'smooth': 1e-6,
+    },
+    0.2474,
+]
+
+
+
+
+
+class TestTverskyLoss(unittest.TestCase):
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7])
+    def test_shape(self, input_param, input_data, expected_val):
+        result = TverskyLoss(**input_param).forward(**input_data)
+        self.assertAlmostEqual(result.item(), expected_val, places=4)
+
+if __name__ == '__main__':
+    unittest.main()