Repository: Project-MONAI/MONAI

Issue:
Add LabelToMask transform
**Is your feature request related to a problem? Please describe.**
Need a transform to convert different kinds of segmentation labels to binary masks to pre-process image for other tasks.
labels can be 1 channel or multi-channels in One-Hot format.



Discussion:

Resolving patch:
diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py
--- a/monai/transforms/utility/array.py
+++ b/monai/transforms/utility/array.py
@@ -13,7 +13,7 @@
 https://github.com/Project-MONAI/MONAI/wiki/MONAI_Design
 """
 
-from typing import Callable, Optional
+from typing import Callable, Optional, Union, Sequence
 
 import time
 import logging
@@ -22,6 +22,7 @@
 import torch
 
 from monai.transforms.compose import Transform
+from monai.utils import ensure_tuple
 
 
 class Identity(Transform):
@@ -372,3 +373,44 @@ def __call__(self, img, func: Optional[Callable] = None):
             return self.func(img)
         else:
             raise RuntimeError("neither func or self.func is callable.")
+
+
+class LabelToMask(Transform):
+    """
+    Convert labels to mask for other tasks. A typical usage is to convert segmentation labels
+    to mask data to pre-process images and then feed the images into classification network.
+    It can support single channel labels or One-Hot labels with specified `select_labels`.
+    For example, users can select `label value = [2, 3]` to construct mask data, or select the
+    second and the third channels of labels to construct mask data.
+    The output mask data can be a multiple channels binary data or a single channel binary
+    data that merges all the channels.
+
+    Args:
+        select_labels: labels to generate mask from. for 1 channel label, the `select_labels`
+            is the expected label values, like: [1, 2, 3]. for One-Hot format label, the
+            `select_labels` is the expected channel indexes.
+        merge_channels: whether to use `np.any()` to merge the result on channel dim. if yes,
+            will return a single channel mask with binary data.
+
+    """
+
+    def __init__(self, select_labels: Union[Sequence[int], int], merge_channels: bool = False):
+        self.select_labels = ensure_tuple(select_labels)
+        self.merge_channels = merge_channels
+
+    def __call__(
+        self, img, select_labels: Optional[Union[Sequence[int], int]] = None, merge_channels: Optional[bool] = None
+    ):
+        if select_labels is None:
+            select_labels = self.select_labels
+        else:
+            select_labels = ensure_tuple(select_labels)
+        if merge_channels is None:
+            merge_channels = self.merge_channels
+
+        if img.shape[0] > 1:
+            data = img[[*(select_labels)]]
+        else:
+            data = np.where(np.in1d(img, select_labels), True, False).reshape(img.shape)
+
+        return np.any(data, axis=0, keepdims=True) if merge_channels else data
diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py
--- a/monai/transforms/utility/dictionary.py
+++ b/monai/transforms/utility/dictionary.py
@@ -39,6 +39,7 @@
     SimulateDelay,
     Identity,
     Lambda,
+    LabelToMask,
 )
 
 
@@ -461,6 +462,35 @@ def __call__(self, data):
         return d
 
 
+class LabelToMaskd(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.LabelToMask`.
+
+    Args:
+        keys: keys of the corresponding items to be transformed.
+            See also: :py:class:`monai.transforms.compose.MapTransform`
+        select_labels: labels to generate mask from. for 1 channel label, the `select_labels`
+            is the expected label values, like: [1, 2, 3]. for One-Hot format label, the
+            `select_labels` is the expected channel indexes.
+        merge_channels: whether to use `np.any()` to merge the result on channel dim.
+            if yes, will return a single channel mask with binary data.
+
+    """
+
+    def __init__(
+        self, keys: KeysCollection, select_labels: Union[Sequence[int], int], merge_channels: bool = False
+    ) -> None:
+        super().__init__(keys)
+        self.converter = LabelToMask(select_labels, merge_channels)
+
+    def __call__(self, data):
+        d = dict(data)
+        for key in self.keys:
+            d[key] = self.converter(d[key])
+
+        return d
+
+
 IdentityD = IdentityDict = Identityd
 AsChannelFirstD = AsChannelFirstDict = AsChannelFirstd
 AsChannelLastD = AsChannelLastDict = AsChannelLastd
@@ -475,3 +505,4 @@ def __call__(self, data):
 CopyItemsD = CopyItemsDict = CopyItemsd
 ConcatItemsD = ConcatItemsDict = ConcatItemsd
 LambdaD = LambdaDict = Lambdad
+LabelToMaskD = LabelToMaskDict = LabelToMaskd


Test patch:
diff --git a/tests/test_label_to_mask.py b/tests/test_label_to_mask.py
new file mode 100644
--- /dev/null
+++ b/tests/test_label_to_mask.py
@@ -0,0 +1,57 @@
+# 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 numpy as np
+from parameterized import parameterized
+
+from monai.transforms import LabelToMask
+
+TEST_CASE_1 = [
+    {"select_labels": [2, 3], "merge_channels": False},
+    np.array([[[1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4], [5, 5, 5], [6, 6, 6]]]),
+    np.array([[[0, 0, 0], [1, 1, 1], [1, 1, 1], [0, 0, 0], [0, 0, 0], [0, 0, 0]]]),
+]
+
+TEST_CASE_2 = [
+    {"select_labels": 2, "merge_channels": False},
+    np.array([[[1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4], [5, 5, 5], [6, 6, 6]]]),
+    np.array([[[0, 0, 0], [1, 1, 1], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]]),
+]
+
+TEST_CASE_3 = [
+    {"select_labels": [1, 2], "merge_channels": False},
+    np.array([[[0, 0, 1], [0, 1, 0]], [[1, 0, 0], [0, 1, 1]], [[1, 0, 1], [1, 1, 0]]]),
+    np.array([[[1, 0, 0], [0, 1, 1]], [[1, 0, 1], [1, 1, 0]]]),
+]
+
+TEST_CASE_4 = [
+    {"select_labels": 2, "merge_channels": False},
+    np.array([[[0, 0, 1], [0, 1, 0]], [[1, 0, 0], [0, 1, 1]], [[1, 0, 1], [1, 1, 0]]]),
+    np.array([[[1, 0, 1], [1, 1, 0]]]),
+]
+
+TEST_CASE_5 = [
+    {"select_labels": [1, 2], "merge_channels": True},
+    np.array([[[0, 0, 1], [0, 1, 0]], [[1, 0, 0], [0, 1, 1]], [[1, 0, 1], [1, 1, 0]]]),
+    np.array([[[1, 0, 1], [1, 1, 1]]]),
+]
+
+
+class TestLabelToMask(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
+    def test_value(self, argments, image, expected_data):
+        result = LabelToMask(**argments)(image)
+        np.testing.assert_allclose(result, expected_data)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_label_to_maskd.py b/tests/test_label_to_maskd.py
new file mode 100644
--- /dev/null
+++ b/tests/test_label_to_maskd.py
@@ -0,0 +1,57 @@
+# 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 numpy as np
+from parameterized import parameterized
+
+from monai.transforms import LabelToMaskd
+
+TEST_CASE_1 = [
+    {"keys": "img", "select_labels": [2, 3], "merge_channels": False},
+    {"img": np.array([[[1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4], [5, 5, 5], [6, 6, 6]]])},
+    np.array([[[0, 0, 0], [1, 1, 1], [1, 1, 1], [0, 0, 0], [0, 0, 0], [0, 0, 0]]]),
+]
+
+TEST_CASE_2 = [
+    {"keys": "img", "select_labels": 2, "merge_channels": False},
+    {"img": np.array([[[1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4], [5, 5, 5], [6, 6, 6]]])},
+    np.array([[[0, 0, 0], [1, 1, 1], [0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]]),
+]
+
+TEST_CASE_3 = [
+    {"keys": "img", "select_labels": [1, 2], "merge_channels": False},
+    {"img": np.array([[[0, 0, 1], [0, 1, 0]], [[1, 0, 0], [0, 1, 1]], [[1, 0, 1], [1, 1, 0]]])},
+    np.array([[[1, 0, 0], [0, 1, 1]], [[1, 0, 1], [1, 1, 0]]]),
+]
+
+TEST_CASE_4 = [
+    {"keys": "img", "select_labels": 2, "merge_channels": False},
+    {"img": np.array([[[0, 0, 1], [0, 1, 0]], [[1, 0, 0], [0, 1, 1]], [[1, 0, 1], [1, 1, 0]]])},
+    np.array([[[1, 0, 1], [1, 1, 0]]]),
+]
+
+TEST_CASE_5 = [
+    {"keys": "img", "select_labels": [1, 2], "merge_channels": True},
+    {"img": np.array([[[0, 0, 1], [0, 1, 0]], [[1, 0, 0], [0, 1, 1]], [[1, 0, 1], [1, 1, 0]]])},
+    np.array([[[1, 0, 1], [1, 1, 1]]]),
+]
+
+
+class TestLabelToMaskd(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
+    def test_value(self, argments, image, expected_data):
+        result = LabelToMaskd(**argments)(image)
+        np.testing.assert_allclose(result["img"], expected_data)
+
+
+if __name__ == "__main__":
+    unittest.main()