Repository: Project-MONAI/MONAI

Issue:
add SplitChannel transform for model output
**Is your feature request related to a problem? Please describe.**
`SplitChannel` transform can be used to split a multiple class label or model outputs.
For example, split multiple classes and compute the `mean_dice` metric separately.


Discussion:
Nic-Ma: /black
Nic-Ma: Hi @ericspod ,

Thanks very much for your review.
I updated the PR according to your comments.
And I got 2 small questions here:
1. Why you think "assert" is not good for parameters? not easy to "try-catch" outside?
2. Do you think it's OK to put post transforms at `/transforms/post/` or move to some other places like `/networks/`?

Thanks.
ericspod: Hi @Nic-Ma,
The thing with assert is that it is intended to be an internal sanity check rather than a mechanism to easily check some condition. Asserts can be turned off at runtime through a command line argument for efficiency, but we don't ever want input validation to be disabled. The assert statement is more convenient than having a `if ...: raise ...` kind of construct everywhere but it is better practice. Even better still would be some sort of design-by-contract facility so that we could define preconditions for inputs which are automatically turned into these checks for us. Something like [PyContracts](http://andreacensi.github.io/contracts/) or [icontract](https://github.com/Parquery/icontract) though I'm not recommending we start using these in MONAI without discussing it further. 

I think it's good to keep all the transforms we define in one place so yes under the `/transforms/post` directory. My feeling is that `/networks` should be strictly for the network definitions themselves, the layers they use, and a few utilities. Transforms are ancillary to the networks themselves even if defined in Pytorch. 

Patch:
diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst
index 5b552988..e505d962 100644
--- a/docs/source/transforms.rst
+++ b/docs/source/transforms.rst
@@ -315,6 +315,12 @@ Vanilla Transforms
     :members:
     :special-members: __call__
 
+`SplitChannel`
+~~~~~~~~~~~~~~
+.. autoclass:: SplitChannel
+    :members:
+    :special-members: __call__
+
 Dictionary-based Transforms
 ----------------------------
 
@@ -570,6 +576,12 @@ Dictionary-based Transforms
     :members:
     :special-members: __call__
 
+`SplitChanneld`
+~~~~~~~~~~~~~~~
+.. autoclass:: SplitChanneld
+    :members:
+    :special-members: __call__
+
 Transform Adaptors
 ------------------
 
diff --git a/monai/transforms/__init__.py b/monai/transforms/__init__.py
index b5671265..5cc032b6 100644
--- a/monai/transforms/__init__.py
+++ b/monai/transforms/__init__.py
@@ -20,3 +20,5 @@ from .spatial.array import *
 from .spatial.dictionary import *
 from .utility.array import *
 from .utility.dictionary import *
+from .post.array import *
+from .post.dictionary import *
diff --git a/monai/transforms/post/__init__.py b/monai/transforms/post/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py
new file mode 100644
index 00000000..92d171fd
--- /dev/null
+++ b/monai/transforms/post/array.py
@@ -0,0 +1,48 @@
+# 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.
+"""
+A collection of "vanilla" transforms for the model output tensors
+https://github.com/Project-MONAI/MONAI/wiki/MONAI_Design
+"""
+
+from monai.transforms.compose import Transform
+from monai.networks.utils import one_hot
+
+
+class SplitChannel(Transform):
+    """
+    Split PyTorch Tensor data according to the channel dim, if only 1 channel, convert to One-Hot
+    format first based on the class number. Users can use this transform to compute metrics on every
+    single class to get more details of validation/evaluation. Expected input shape:
+    (batch_size, num_channels, spatial_dim_1[, spatial_dim_2, ...])
+
+    Args:
+        to_onehot (bool): whether to convert the data to One-Hot format, default is False.
+        num_classes (int): the class number used to convert to One-Hot format if `to_onehot` is True.
+
+    """
+
+    def __init__(self, to_onehot=False, num_classes=None):
+        self.to_onehot = to_onehot
+        self.num_classes = num_classes
+
+    def __call__(self, img, to_onehot=None, num_classes=None):
+        if self.to_onehot if to_onehot is None else to_onehot:
+            if num_classes is None:
+                num_classes = self.num_classes
+            assert isinstance(num_classes, int), "must specify class number for One-Hot."
+            img = one_hot(img, num_classes)
+        n_classes = img.shape[1]
+        outputs = list()
+        for i in range(n_classes):
+            outputs.append(img[:, i : i + 1])
+
+        return outputs
diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py
new file mode 100644
index 00000000..097e8b51
--- /dev/null
+++ b/monai/transforms/post/dictionary.py
@@ -0,0 +1,59 @@
+# 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.
+"""
+A collection of dictionary-based wrappers around the "vanilla" transforms for model output tensors
+defined in :py:class:`monai.transforms.utility.array`.
+
+Class names are ended with 'd' to denote dictionary-based transforms.
+"""
+
+from monai.utils.misc import ensure_tuple_rep
+from monai.transforms.compose import MapTransform
+from monai.transforms.post.array import SplitChannel
+
+
+class SplitChanneld(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.SplitChannel`.
+    All the input specified by `keys` should be splitted into same count of data.
+
+    """
+
+    def __init__(self, keys, output_postfixes, to_onehot=False, num_classes=None):
+        """
+        Args:
+            keys (hashable items): keys of the corresponding items to be transformed.
+                See also: :py:class:`monai.transforms.compose.MapTransform`
+            output_postfixes (list, tuple): the postfixes to construct keys to store splitted data.
+                for example: if the key of input data is `pred` and split 2 classes, the output
+                data keys will be: pred_(output_postfixes[0]), pred_(output_postfixes[1])
+            to_onehot (bool or list of bool): whether to convert the data to One-Hot format, default is False.
+        num_classes (int or list of int): the class number used to convert to One-Hot format if `to_onehot` is True.
+        """
+        super().__init__(keys)
+        if not isinstance(output_postfixes, (list, tuple)):
+            raise ValueError("must specify key postfixes to store splitted data.")
+        self.output_postfixes = output_postfixes
+        self.to_onehot = ensure_tuple_rep(to_onehot, len(self.keys))
+        self.num_classes = ensure_tuple_rep(num_classes, len(self.keys))
+        self.splitter = SplitChannel()
+
+    def __call__(self, data):
+        d = dict(data)
+        for idx, key in enumerate(self.keys):
+            rets = self.splitter(d[key], self.to_onehot[idx], self.num_classes[idx])
+            assert len(self.output_postfixes) == len(rets), "count of splitted results must match output_postfixes."
+            for i, r in enumerate(rets):
+                d[f"{key}_{self.output_postfixes[i]}"] = r
+        return d
+
+
+SplitChannelD = SplitChannelDict = SplitChanneld


Test patch:
diff --git a/tests/test_split_channel.py b/tests/test_split_channel.py
new file mode 100644
index 00000000..76d1b3ea
--- /dev/null
+++ b/tests/test_split_channel.py
@@ -0,0 +1,31 @@
+# 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.transforms import SplitChannel
+
+TEST_CASE_1 = [{"to_onehot": False}, torch.randint(0, 2, size=(4, 3, 3, 4)), (4, 1, 3, 4)]
+
+TEST_CASE_2 = [{"to_onehot": True, "num_classes": 3}, torch.randint(0, 3, size=(4, 1, 3, 4)), (4, 1, 3, 4)]
+
+
+class TestSplitChannel(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    def test_shape(self, input_param, test_data, expected_shape):
+        result = SplitChannel(**input_param)(test_data)
+        for data in result:
+            self.assertTupleEqual(data.shape, expected_shape)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_split_channeld.py b/tests/test_split_channeld.py
new file mode 100644
index 00000000..c2821c91
--- /dev/null
+++ b/tests/test_split_channeld.py
@@ -0,0 +1,46 @@
+# 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.transforms import SplitChanneld
+
+TEST_CASE_1 = [
+    {"keys": ["pred"], "output_postfixes": ["cls1", "cls2", "cls3"], "to_onehot": False},
+    {"pred": torch.randint(0, 2, size=(4, 3, 3, 4))},
+    (4, 1, 3, 4),
+]
+
+TEST_CASE_2 = [
+    {"keys": ["pred"], "output_postfixes": ["cls1", "cls2", "cls3"], "to_onehot": True, "num_classes": 3},
+    {"pred": torch.randint(0, 3, size=(4, 1, 3, 4))},
+    (4, 1, 3, 4),
+]
+
+TEST_CASE_3 = [
+    {"keys": ["pred", "label"], "output_postfixes": ["cls1", "cls2", "cls3"], "to_onehot": True, "num_classes": 3},
+    {"pred": torch.randint(0, 3, size=(4, 1, 3, 4)), "label": torch.randint(0, 3, size=(4, 1, 3, 4))},
+    (4, 1, 3, 4),
+]
+
+
+class TestSplitChanneld(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
+    def test_shape(self, input_param, test_data, expected_shape):
+        result = SplitChanneld(**input_param)(test_data)
+        for k, v in result.items():
+            if "cls" in k:
+                self.assertTupleEqual(v.shape, expected_shape)
+
+
+if __name__ == "__main__":
+    unittest.main()