Repository: Project-MONAI/MONAI

Issue:
Adding a loss mask? 
Are there any plans to add an easy way to include "loss mask" to `Dataset` class for segmentation task? As what I understood, it only takes `img` and `seg`?

For example I have such CT volumes with the following "windowing" that could be masked away

![image](https://user-images.githubusercontent.com/1060514/84707668-076fbb80-af57-11ea-9bba-be1183c43d68.png)

And in general the following behavior could be a nice feature

https://github.com/deepmind/deepmind-research/blob/cf0a9ca46058a61e8b66bc8708cd35ac77209bc6/hierarchical_probabilistic_unet/geco_utils.py#L100 


Discussion:
Hi @petteriTeikari ,

Thanks for your interest here.
Actually, the Dataset in MONAI accepts `dict` type input data.
So you can put a datalist like:
```py
[
    {"img": img1, "seg": seg1, "loss_mask": mask1},
    {"img": img2, "seg": seg2, "loss_mask": mask2},
    {"img": img3, "seg": seg3, "loss_mask": mask3},
    ... ...
]
```

And BTW, suggest using `CacheDataset` to accelerate your training process.

Thanks.
@petteriTeikari 
We also have the masked dice loss, for convenience, which should work in your case (if you're using dice)
https://github.com/Project-MONAI/MONAI/blob/22d975b1d8c374b2ba5a360140ccd983be193311/monai/losses/dice.py#L139

Thanks for the clarifications @Nic-Ma and @myron that indeed was something that I was trying to find from the documentation 
> @petteriTeikari
> We also have the masked dice loss, for convenience, which should work in your case (if you're using dice)
> https://github.com/Project-MONAI/MONAI/blob/22d975b1d8c374b2ba5a360140ccd983be193311/monai/losses/dice.py#L139

@myron Do you have any plans to add the "masked"-option for all types of losses? As in, would be niceto play with some own custom losses and all your off-the-shelf ones with the loss masking option 
Sorry, could I ask what is the current status of this issue? @wyli 
It seems like currently `MaskedDiceLoss` is the only type of segmentation loss with an implementation of masking. Would it be beneficial to extend masking to the other segmentation losses, such as Tversky or Focal Losses? Perhaps with a wrapper function like [MultiScaleLoss](https://docs.monai.io/en/latest/losses.html#multiscaleloss)? Thanks 😄 
sure @eddyleelin, let me look into this...

Resolving patch:
diff --git a/monai/losses/__init__.py b/monai/losses/__init__.py
--- a/monai/losses/__init__.py
+++ b/monai/losses/__init__.py
@@ -27,4 +27,5 @@
 from .focal_loss import FocalLoss
 from .image_dissimilarity import GlobalMutualInformationLoss, LocalNormalizedCrossCorrelationLoss
 from .multi_scale import MultiScaleLoss
+from .spatial_mask import MaskedLoss
 from .tversky import TverskyLoss
diff --git a/monai/losses/dice.py b/monai/losses/dice.py
--- a/monai/losses/dice.py
+++ b/monai/losses/dice.py
@@ -19,6 +19,7 @@
 from torch.nn.modules.loss import _Loss
 
 from monai.losses.focal_loss import FocalLoss
+from monai.losses.spatial_mask import MaskedLoss
 from monai.networks import one_hot
 from monai.utils import LossReduction, Weight
 
@@ -185,32 +186,21 @@ class MaskedDiceLoss(DiceLoss):
 
     """
 
-    def forward(self, input: torch.Tensor, target: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+    def __init__(self, *args, **kwargs) -> None:
+        """
+        Args follow :py:class:`monai.losses.DiceLoss`.
+        """
+        super().__init__(*args, **kwargs)
+        self.spatial_weighted = MaskedLoss(loss=super().forward)
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor, mask: Optional[torch.Tensor] = None):
         """
         Args:
             input: the shape should be BNH[WD].
             target: the shape should be BNH[WD].
             mask: the shape should B1H[WD] or 11H[WD].
         """
-        if mask is not None:
-            # checking if mask is of proper shape
-            if input.dim() != mask.dim():
-                raise AssertionError(f"dim of input ({input.shape}) is different from mask ({mask.shape})")
-            if not (input.shape[0] == mask.shape[0] or mask.shape[0] == 1):
-                raise AssertionError(f" batch size of mask ({mask.shape}) must be 1 or equal to input ({input.shape})")
-
-            if target.dim() > 1:
-                if mask.shape[1] != 1:
-                    raise AssertionError(f"mask ({mask.shape}) must have only 1 channel")
-                if input.shape[2:] != mask.shape[2:]:
-                    raise AssertionError(f"spatial size of input ({input.shape}) is different from mask ({mask.shape})")
-
-            input = input * mask
-            target = target * mask
-        else:
-            warnings.warn("no mask value specified for the MaskedDiceLoss.")
-
-        return super().forward(input=input, target=target)
+        return self.spatial_weighted(input=input, target=target, mask=mask)
 
 
 class GeneralizedDiceLoss(_Loss):
diff --git a/monai/losses/spatial_mask.py b/monai/losses/spatial_mask.py
new file mode 100644
--- /dev/null
+++ b/monai/losses/spatial_mask.py
@@ -0,0 +1,67 @@
+# Copyright 2020 - 2021 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 inspect
+import warnings
+from typing import Callable, Optional, Union
+
+import torch
+from torch.nn.modules.loss import _Loss
+
+__all__ = ["MaskedLoss"]
+
+
+class MaskedLoss(_Loss):
+    """
+    This is a wrapper class for the loss functions.  It allows for additional
+    weighting masks to be applied to both input and target.
+
+    See Also:
+        - :py:class:`monai.losses.MaskedDiceLoss`
+    """
+
+    def __init__(self, loss: Union[Callable, _Loss], *loss_args, **loss_kwargs) -> None:
+        """
+        Args:
+            loss: loss function to be wrapped, this could be a loss class or an instance of a loss class.
+            loss_args: arguments to the loss function's constructor if `loss` is a class.
+            loss_kwargs: keyword arguments to the loss function's constructor if `loss` is a class.
+        """
+        super().__init__()
+        if inspect.isclass(loss):  # call the loss function's constructor if it's a class
+            self.loss = loss(*loss_args, **loss_kwargs)
+        else:
+            self.loss = loss  # loss is a callable loss function instance.
+
+        if not callable(self.loss):
+            raise ValueError("The loss function is not callable.")
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor, mask: Optional[torch.Tensor] = None):
+        """
+        Args:
+            input: the shape should be BNH[WD].
+            target: the shape should be BNH[WD].
+            mask: the shape should be B1H[WD] or 11H[WD].
+        """
+        if mask is None:
+            warnings.warn("No mask value specified for the MaskedLoss.")
+            return self.loss(input, target)
+
+        if input.dim() != mask.dim():
+            warnings.warn(f"Dim of input ({input.shape}) is different from mask ({mask.shape}).")
+        if not (input.shape[0] == mask.shape[0] or mask.shape[0] == 1):
+            raise ValueError(f"Batch size of mask ({mask.shape}) must be one or equal to input ({input.shape}).")
+        if target.dim() > 1:
+            if mask.shape[1] != 1:
+                raise ValueError(f"Mask ({mask.shape}) must have only one channel.")
+            if input.shape[2:] != mask.shape[2:]:
+                warnings.warn(f"Spatial size of input ({input.shape}) is different from mask ({mask.shape}).")
+        return self.loss(input * mask, target * mask)


Test patch:
diff --git a/tests/test_masked_loss.py b/tests/test_masked_loss.py
new file mode 100644
--- /dev/null
+++ b/tests/test_masked_loss.py
@@ -0,0 +1,89 @@
+# Copyright 2020 - 2021 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 numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.losses.dice import DiceFocalLoss, DiceLoss
+from monai.losses.spatial_mask import MaskedLoss
+from monai.utils import set_determinism
+from tests.utils import SkipIfBeforePyTorchVersion, test_script_save
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+TEST_CASES = [
+    [
+        {
+            "loss": DiceFocalLoss,
+            "focal_weight": torch.tensor([1.0, 1.0, 2.0]),
+            "gamma": 0.1,
+            "lambda_focal": 0.5,
+            "include_background": True,
+            "to_onehot_y": True,
+            "reduction": "sum",
+        },
+        [(12.105497, 18.805185), (10.636354, 6.3138)],
+    ],
+]
+
+
+class TestMaskedLoss(unittest.TestCase):
+    def setUp(self):
+        set_determinism(0)
+
+    def tearDown(self):
+        set_determinism(None)
+
+    @parameterized.expand(TEST_CASES)
+    def test_shape(self, input_param, expected_val):
+        size = [3, 3, 5, 5]
+        label = torch.randint(low=0, high=2, size=size)
+        label = torch.argmax(label, dim=1, keepdim=True)
+        pred = torch.randn(size)
+        print(label[0, 0, 0])
+        result = MaskedLoss(**input_param)(pred, label, None)
+        out = result.detach().cpu().numpy()
+        checked = np.allclose(out, expected_val[0][0]) or np.allclose(out, expected_val[0][1])
+        self.assertTrue(checked)
+
+        mask = torch.randint(low=0, high=2, size=label.shape)
+        result = MaskedLoss(**input_param)(pred, label, mask)
+        out = result.detach().cpu().numpy()
+        checked = np.allclose(out, expected_val[1][0]) or np.allclose(out, expected_val[1][1])
+        self.assertTrue(checked)
+
+    def test_ill_opts(self):
+        with self.assertRaisesRegex(ValueError, ""):
+            MaskedLoss(loss=[])
+
+        dice_loss = DiceLoss(include_background=True, sigmoid=True, smooth_nr=1e-5, smooth_dr=1e-5)
+        with self.assertRaisesRegex(ValueError, ""):
+            masked = MaskedLoss(loss=dice_loss)
+            masked(input=torch.zeros((3, 1, 2, 2)), target=torch.zeros((3, 1, 2, 2)), mask=torch.zeros((3, 3, 2, 2)))
+        with self.assertRaisesRegex(ValueError, ""):
+            masked = MaskedLoss(loss=dice_loss)
+            masked(input=torch.zeros((3, 3, 2, 2)), target=torch.zeros((3, 2, 2, 2)), mask=torch.zeros((3, 3, 2, 2)))
+
+    @SkipIfBeforePyTorchVersion((1, 7, 0))
+    def test_script(self):
+        input_param, expected_val = TEST_CASES[0]
+        size = [3, 3, 5, 5]
+        label = torch.randint(low=0, high=2, size=size)
+        label = torch.argmax(label, dim=1, keepdim=True)
+        pred = torch.randn(size)
+        loss = MaskedLoss(**input_param)
+        test_script_save(loss, pred, label)
+
+
+if __name__ == "__main__":
+    unittest.main()