Repository: Project-MONAI/MONAI

Issue:
Pre-segment images using a mask
Hi, is there a way to mask images given a mask by e.g. using a "monai.transforms" function? For example: Use a lung mask Nifti's for "pre-segmenting" Chest CT Nifti's before feeding it into the [classification example](https://github.com/Project-MONAI/MONAI/blob/master/examples/classification_3d/densenet_training_dict.py)

Thank you


Discussion:
Hi @bartmch ,

Thanks for your great suggestion!
Actually, this transform is already in our internal to-do list, I will submit a task for it.

Thanks.
if you are looking for some quick prototype, This could be done by creating a simple `masked` function (the transforms are callables, can be complicated, but in this case it's quite simple provided that the mask is already aligned with the image)
```py
def masked(data_dict):
     data_dict["image"] = data_dict["image"] * data_dict["label"]
     return data_dict

```
use it in a composed:
```python

val_transforms = Compose([
    LoadNiftid(keys=['image', 'label']),
    AddChanneld(keys=['image', 'label']),
    Spacingd(keys=['image', 'label'], pixdim=(1.5, 1.5, 2.), mode=('bilinear', 'nearest')),
    Orientationd(keys=['image', 'label'], axcodes='RAS'),
    masked,
    ScaleIntensityRanged(keys=['image'], a_min=-57, a_max=164, b_min=0.0, b_max=1.0, clip=True),
    ToTensord(keys=['image', 'label'])
])
```
![Screenshot 2020-06-30 at 20 20 01](https://user-images.githubusercontent.com/831580/86167864-31a4aa00-bb0f-11ea-8978-3f41370a6348.png)

 if you comment out the `masked` in the transform chain, you get the original:
```py
val_transforms = Compose([
    LoadNiftid(keys=['image', 'label']),
    AddChanneld(keys=['image', 'label']),
    Spacingd(keys=['image', 'label'], pixdim=(1.5, 1.5, 2.), mode=('bilinear', 'nearest')),
    Orientationd(keys=['image', 'label'], axcodes='RAS'),
    #  masked,
    ScaleIntensityRanged(keys=['image'], a_min=-57, a_max=164, b_min=0.0, b_max=1.0, clip=True),
    ToTensord(keys=['image', 'label'])
])
```
![Screenshot 2020-06-30 at 20 20 11](https://user-images.githubusercontent.com/831580/86167936-4ed97880-bb0f-11ea-9955-1ed4a496aaa9.png)


Resolving patch:
diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py
--- a/monai/transforms/intensity/array.py
+++ b/monai/transforms/intensity/array.py
@@ -431,3 +431,29 @@ def __call__(self, img):
             img = np.clip(img, self.b_min, self.b_max)
 
         return img
+
+
+class MaskIntensity(Transform):
+    """
+    Mask the intensity values of input image with the specified mask data.
+    Mask data must have the same spatial size as the input image, and all
+    the intensity values of input image corresponding to `0` in the mask
+    data will be set to `0`, others will keep the original value.
+
+    Args:
+        mask_data: if mask data is single channel, apply to evey channel
+            of input image. if multiple channels, the channel number must
+            match input data. mask_data will be converted to `bool` values
+            by `mask_data > 0` before applying transform to input image.
+
+    """
+
+    def __init__(self, mask_data: np.ndarray):
+        self.mask_data = mask_data
+
+    def __call__(self, img, mask_data: Optional[np.ndarray] = None):
+        mask_data_ = self.mask_data > 0 if mask_data is None else mask_data > 0
+        if mask_data_.shape[0] != 1 and mask_data_.shape[0] != img.shape[0]:
+            raise RuntimeError("mask data has more than 1 channel and do not match channels of input data.")
+
+        return img * mask_data_
diff --git a/monai/transforms/intensity/dictionary.py b/monai/transforms/intensity/dictionary.py
--- a/monai/transforms/intensity/dictionary.py
+++ b/monai/transforms/intensity/dictionary.py
@@ -29,6 +29,7 @@
     ShiftIntensity,
     ScaleIntensity,
     ScaleIntensityRangePercentiles,
+    MaskIntensity,
 )
 
 
@@ -380,6 +381,31 @@ def __call__(self, data):
         return d
 
 
+class MaskIntensityd(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.MaskIntensity`.
+
+    Args:
+        keys: keys of the corresponding items to be transformed.
+            See also: :py:class:`monai.transforms.compose.MapTransform`
+        mask_data: if mask data is single channel, apply to evey channel
+            of input image. if multiple channels, the channel number must
+            match input data. mask_data will be converted to `bool` values
+            by `mask_data > 0` before applying transform to input image.
+
+    """
+
+    def __init__(self, keys: KeysCollection, mask_data: np.ndarray):
+        super().__init__(keys)
+        self.converter = MaskIntensity(mask_data)
+
+    def __call__(self, data):
+        d = dict(data)
+        for key in self.keys:
+            d[key] = self.converter(d[key])
+        return d
+
+
 RandGaussianNoiseD = RandGaussianNoiseDict = RandGaussianNoised
 ShiftIntensityD = ShiftIntensityDict = ShiftIntensityd
 RandShiftIntensityD = RandShiftIntensityDict = RandShiftIntensityd
@@ -391,3 +417,4 @@ def __call__(self, data):
 AdjustContrastD = AdjustContrastDict = AdjustContrastd
 RandAdjustContrastD = RandAdjustContrastDict = RandAdjustContrastd
 ScaleIntensityRangePercentilesD = ScaleIntensityRangePercentilesDict = ScaleIntensityRangePercentilesd
+MaskIntensityD = MaskIntensityDict = MaskIntensityd


Test patch:
diff --git a/tests/test_mask_intensity.py b/tests/test_mask_intensity.py
new file mode 100644
--- /dev/null
+++ b/tests/test_mask_intensity.py
@@ -0,0 +1,45 @@
+# 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 MaskIntensity
+
+TEST_CASE_1 = [
+    {"mask_data": np.array([[[0, 0, 0], [0, 1, 0], [0, 0, 0]]])},
+    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], [0, 2, 0], [0, 0, 0]], [[0, 0, 0], [0, 5, 0], [0, 0, 0]]]),
+]
+
+TEST_CASE_2 = [
+    {"mask_data": np.array([[[0, 0, 0], [0, 5, 0], [0, 0, 0]]])},
+    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], [0, 2, 0], [0, 0, 0]], [[0, 0, 0], [0, 5, 0], [0, 0, 0]]]),
+]
+
+TEST_CASE_3 = [
+    {"mask_data": np.array([[[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 1, 0], [0, 1, 0], [0, 1, 0]]])},
+    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], [0, 2, 0], [0, 0, 0]], [[0, 4, 0], [0, 5, 0], [0, 6, 0]]]),
+]
+
+
+class TestMaskIntensity(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
+    def test_value(self, argments, image, expected_data):
+        result = MaskIntensity(**argments)(image)
+        np.testing.assert_allclose(result, expected_data)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_mask_intensityd.py b/tests/test_mask_intensityd.py
new file mode 100644
--- /dev/null
+++ b/tests/test_mask_intensityd.py
@@ -0,0 +1,45 @@
+# 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 MaskIntensityd
+
+TEST_CASE_1 = [
+    {"keys": "img", "mask_data": np.array([[[0, 0, 0], [0, 1, 0], [0, 0, 0]]])},
+    {"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], [0, 2, 0], [0, 0, 0]], [[0, 0, 0], [0, 5, 0], [0, 0, 0]]]),
+]
+
+TEST_CASE_2 = [
+    {"keys": "img", "mask_data": np.array([[[0, 0, 0], [0, 5, 0], [0, 0, 0]]])},
+    {"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], [0, 2, 0], [0, 0, 0]], [[0, 0, 0], [0, 5, 0], [0, 0, 0]]]),
+]
+
+TEST_CASE_3 = [
+    {"keys": "img", "mask_data": np.array([[[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 1, 0], [0, 1, 0], [0, 1, 0]]])},
+    {"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], [0, 2, 0], [0, 0, 0]], [[0, 4, 0], [0, 5, 0], [0, 6, 0]]]),
+]
+
+
+class TestMaskIntensityd(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
+    def test_value(self, argments, image, expected_data):
+        result = MaskIntensityd(**argments)(image)
+        np.testing.assert_allclose(result["img"], expected_data)
+
+
+if __name__ == "__main__":
+    unittest.main()