AddChannel, AsChannelFirst, AsChannelLast, EnsureChannelFirst, Identity, RepeatChannel (#2840)

Author: rijobro

monai/transforms/__init__.py

@@ -518,3 +518,4 @@
     weighted_patch_samples,
     zero_margins,
 )
+from .utils_pytorch_numpy_unification import moveaxis

monai/transforms/utility/array.py

@@ -31,6 +31,7 @@
     map_binary_to_indices,
     map_classes_to_indices,
 )
+from monai.transforms.utils_pytorch_numpy_unification import moveaxis
 from monai.utils import (
     convert_to_numpy,
     convert_to_tensor,
@@ -82,17 +83,18 @@
 
 class Identity(Transform):
     """
-    Convert the input to an np.ndarray, if input data is np.ndarray or subclasses, return unchanged data.
+    Do nothing to the data.
     As the output value is same as input, it can be used as a testing tool to verify the transform chain,
     Compose or transform adaptor, etc.
-
     """
 
-    def __call__(self, img: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
+    def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
         Apply the transform to `img`.
         """
-        return np.asanyarray(img)
+        return img
 
 
 class AsChannelFirst(Transform):
@@ -111,16 +113,18 @@ class AsChannelFirst(Transform):
         channel_dim: which dimension of input image is the channel, default is the last dimension.
     """
 
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
     def __init__(self, channel_dim: int = -1) -> None:
         if not (isinstance(channel_dim, int) and channel_dim >= -1):
             raise AssertionError("invalid channel dimension.")
         self.channel_dim = channel_dim
 
-    def __call__(self, img: np.ndarray) -> np.ndarray:
+    def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
         Apply the transform to `img`.
         """
-        return np.moveaxis(img, self.channel_dim, 0)
+        return moveaxis(img, self.channel_dim, 0)
 
 
 class AsChannelLast(Transform):
@@ -138,16 +142,18 @@ class AsChannelLast(Transform):
         channel_dim: which dimension of input image is the channel, default is the first dimension.
     """
 
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
     def __init__(self, channel_dim: int = 0) -> None:
         if not (isinstance(channel_dim, int) and channel_dim >= -1):
             raise AssertionError("invalid channel dimension.")
         self.channel_dim = channel_dim
 
-    def __call__(self, img: np.ndarray) -> np.ndarray:
+    def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
         Apply the transform to `img`.
         """
-        return np.moveaxis(img, self.channel_dim, -1)
+        return moveaxis(img, self.channel_dim, -1)
 
 
 class AddChannel(Transform):
@@ -164,7 +170,9 @@ class AddChannel(Transform):
     transforms.
     """
 
-    def __call__(self, img: NdarrayTensor):
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
+    def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
         Apply the transform to `img`.
         """
@@ -179,14 +187,16 @@ class EnsureChannelFirst(Transform):
     Convert the data to `channel_first` based on the `original_channel_dim` information.
     """
 
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
     def __init__(self, strict_check: bool = True):
         """
         Args:
             strict_check: whether to raise an error when the meta information is insufficient.
         """
         self.strict_check = strict_check
 
-    def __call__(self, img: np.ndarray, meta_dict: Optional[Mapping] = None):
+    def __call__(self, img: NdarrayOrTensor, meta_dict: Optional[Mapping] = None) -> NdarrayOrTensor:
         """
         Apply the transform to `img`.
         """
@@ -220,16 +230,19 @@ class RepeatChannel(Transform):
         repeats: the number of repetitions for each element.
     """
 
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
     def __init__(self, repeats: int) -> None:
         if repeats <= 0:
             raise AssertionError("repeats count must be greater than 0.")
         self.repeats = repeats
 
-    def __call__(self, img: np.ndarray) -> np.ndarray:
+    def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
         Apply the transform to `img`, assuming `img` is a "channel-first" array.
         """
-        return np.repeat(img, self.repeats, 0)
+        repeeat_fn = torch.repeat_interleave if isinstance(img, torch.Tensor) else np.repeat
+        return repeeat_fn(img, self.repeats, 0)  # type: ignore
 
 
 class RemoveRepeatedChannel(Transform):

monai/transforms/utility/dictionary.py

@@ -169,6 +169,8 @@ class Identityd(MapTransform):
     Dictionary-based wrapper of :py:class:`monai.transforms.Identity`.
     """
 
+    backend = Identity.backend
+
     def __init__(self, keys: KeysCollection, allow_missing_keys: bool = False) -> None:
         """
         Args:
@@ -180,9 +182,7 @@ def __init__(self, keys: KeysCollection, allow_missing_keys: bool = False) -> No
         super().__init__(keys, allow_missing_keys)
         self.identity = Identity()
 
-    def __call__(
-        self, data: Mapping[Hashable, Union[np.ndarray, torch.Tensor]]
-    ) -> Dict[Hashable, Union[np.ndarray, torch.Tensor]]:
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
         for key in self.key_iterator(d):
             d[key] = self.identity(d[key])
@@ -194,6 +194,8 @@ class AsChannelFirstd(MapTransform):
     Dictionary-based wrapper of :py:class:`monai.transforms.AsChannelFirst`.
     """
 
+    backend = AsChannelFirst.backend
+
     def __init__(self, keys: KeysCollection, channel_dim: int = -1, allow_missing_keys: bool = False) -> None:
         """
         Args:
@@ -205,7 +207,7 @@ def __init__(self, keys: KeysCollection, channel_dim: int = -1, allow_missing_ke
         super().__init__(keys, allow_missing_keys)
         self.converter = AsChannelFirst(channel_dim=channel_dim)
 
-    def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
         for key in self.key_iterator(d):
             d[key] = self.converter(d[key])
@@ -217,6 +219,8 @@ class AsChannelLastd(MapTransform):
     Dictionary-based wrapper of :py:class:`monai.transforms.AsChannelLast`.
     """
 
+    backend = AsChannelLast.backend
+
     def __init__(self, keys: KeysCollection, channel_dim: int = 0, allow_missing_keys: bool = False) -> None:
         """
         Args:
@@ -228,7 +232,7 @@ def __init__(self, keys: KeysCollection, channel_dim: int = 0, allow_missing_key
         super().__init__(keys, allow_missing_keys)
         self.converter = AsChannelLast(channel_dim=channel_dim)
 
-    def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
         for key in self.key_iterator(d):
             d[key] = self.converter(d[key])
@@ -240,6 +244,8 @@ class AddChanneld(MapTransform):
     Dictionary-based wrapper of :py:class:`monai.transforms.AddChannel`.
     """
 
+    backend = AddChannel.backend
+
     def __init__(self, keys: KeysCollection, allow_missing_keys: bool = False) -> None:
         """
         Args:
@@ -250,7 +256,7 @@ def __init__(self, keys: KeysCollection, allow_missing_keys: bool = False) -> No
         super().__init__(keys, allow_missing_keys)
         self.adder = AddChannel()
 
-    def __call__(self, data: Mapping[Hashable, NdarrayTensor]) -> Dict[Hashable, NdarrayTensor]:
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
         for key in self.key_iterator(d):
             d[key] = self.adder(d[key])
@@ -262,6 +268,8 @@ class EnsureChannelFirstd(MapTransform):
     Dictionary-based wrapper of :py:class:`monai.transforms.EnsureChannelFirst`.
     """
 
+    backend = EnsureChannelFirst.backend
+
     def __init__(
         self,
         keys: KeysCollection,
@@ -289,7 +297,7 @@ def __init__(
         self.meta_keys = ensure_tuple_rep(meta_keys, len(self.keys))
         self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys))
 
-    def __call__(self, data) -> Dict[Hashable, np.ndarray]:
+    def __call__(self, data) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
         for key, meta_key, meta_key_postfix in zip(self.keys, self.meta_keys, self.meta_key_postfix):
             d[key] = self.adjuster(d[key], d[meta_key or f"{key}_{meta_key_postfix}"])
@@ -301,6 +309,8 @@ class RepeatChanneld(MapTransform):
     Dictionary-based wrapper of :py:class:`monai.transforms.RepeatChannel`.
     """
 
+    backend = RepeatChannel.backend
+
     def __init__(self, keys: KeysCollection, repeats: int, allow_missing_keys: bool = False) -> None:
         """
         Args:
@@ -312,7 +322,7 @@ def __init__(self, keys: KeysCollection, repeats: int, allow_missing_keys: bool
         super().__init__(keys, allow_missing_keys)
         self.repeater = RepeatChannel(repeats)
 
-    def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
         for key in self.key_iterator(d):
             d[key] = self.repeater(d[key])

monai/transforms/utils_pytorch_numpy_unification.py

@@ -0,0 +1,41 @@
+# 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 numpy as np
+import torch
+
+from monai.config.type_definitions import NdarrayOrTensor
+
+__all__ = [
+    "moveaxis",
+]
+
+
+def moveaxis(x: NdarrayOrTensor, src: int, dst: int) -> NdarrayOrTensor:
+    if isinstance(x, torch.Tensor):
+        if hasattr(torch, "moveaxis"):
+            return torch.moveaxis(x, src, dst)
+        # moveaxis only available in pytorch as of 1.8.0
+        else:
+            # get original indices
+            indices = list(range(x.ndim))
+            # make src and dst positive
+            if src < 0:
+                src = len(indices) + src
+            if dst < 0:
+                dst = len(indices) + dst
+            # remove desired index and insert it in new position
+            indices.pop(src)
+            indices.insert(dst, src)
+            return x.permute(indices)
+    elif isinstance(x, np.ndarray):
+        return np.moveaxis(x, src, dst)
+    raise RuntimeError()

tests/test_add_channeld.py

@@ -15,16 +15,21 @@
 from parameterized import parameterized
 
 from monai.transforms import AddChanneld
+from tests.utils import TEST_NDARRAYS
 
-TEST_CASE_1 = [
-    {"keys": ["img", "seg"]},
-    {"img": np.array([[0, 1], [1, 2]]), "seg": np.array([[0, 1], [1, 2]])},
-    (1, 2, 2),
-]
+TESTS = []
+for p in TEST_NDARRAYS:
+    TESTS.append(
+        [
+            {"keys": ["img", "seg"]},
+            {"img": p(np.array([[0, 1], [1, 2]])), "seg": p(np.array([[0, 1], [1, 2]]))},
+            (1, 2, 2),
+        ]
+    )
 
 
 class TestAddChanneld(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1])
+    @parameterized.expand(TESTS)
     def test_shape(self, input_param, input_data, expected_shape):
         result = AddChanneld(**input_param)(input_data)
         self.assertEqual(result["img"].shape, expected_shape)

tests/test_as_channel_first.py

@@ -12,23 +12,29 @@
 import unittest
 
 import numpy as np
+import torch
 from parameterized import parameterized
 
 from monai.transforms import AsChannelFirst
+from tests.utils import TEST_NDARRAYS, assert_allclose
 
-TEST_CASE_1 = [{"channel_dim": -1}, (4, 1, 2, 3)]
-
-TEST_CASE_2 = [{"channel_dim": 3}, (4, 1, 2, 3)]
-
-TEST_CASE_3 = [{"channel_dim": 2}, (3, 1, 2, 4)]
+TESTS = []
+for p in TEST_NDARRAYS:
+    TESTS.append([p, {"channel_dim": -1}, (4, 1, 2, 3)])
+    TESTS.append([p, {"channel_dim": 3}, (4, 1, 2, 3)])
+    TESTS.append([p, {"channel_dim": 2}, (3, 1, 2, 4)])
 
 
 class TestAsChannelFirst(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
-    def test_shape(self, input_param, expected_shape):
-        test_data = np.random.randint(0, 2, size=[1, 2, 3, 4])
+    @parameterized.expand(TESTS)
+    def test_value(self, in_type, input_param, expected_shape):
+        test_data = in_type(np.random.randint(0, 2, size=[1, 2, 3, 4]))
         result = AsChannelFirst(**input_param)(test_data)
         self.assertTupleEqual(result.shape, expected_shape)
+        if isinstance(test_data, torch.Tensor):
+            test_data = test_data.cpu().numpy()
+        expected = np.moveaxis(test_data, input_param["channel_dim"], 0)
+        assert_allclose(expected, result)
 
 
 if __name__ == "__main__":

tests/test_as_channel_firstd.py

@@ -15,21 +15,22 @@
 from parameterized import parameterized
 
 from monai.transforms import AsChannelFirstd
+from tests.utils import TEST_NDARRAYS
 
-TEST_CASE_1 = [{"keys": ["image", "label", "extra"], "channel_dim": -1}, (4, 1, 2, 3)]
-
-TEST_CASE_2 = [{"keys": ["image", "label", "extra"], "channel_dim": 3}, (4, 1, 2, 3)]
-
-TEST_CASE_3 = [{"keys": ["image", "label", "extra"], "channel_dim": 2}, (3, 1, 2, 4)]
+TESTS = []
+for p in TEST_NDARRAYS:
+    TESTS.append([p, {"keys": ["image", "label", "extra"], "channel_dim": -1}, (4, 1, 2, 3)])
+    TESTS.append([p, {"keys": ["image", "label", "extra"], "channel_dim": 3}, (4, 1, 2, 3)])
+    TESTS.append([p, {"keys": ["image", "label", "extra"], "channel_dim": 2}, (3, 1, 2, 4)])
 
 
 class TestAsChannelFirstd(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
-    def test_shape(self, input_param, expected_shape):
+    @parameterized.expand(TESTS)
+    def test_shape(self, in_type, input_param, expected_shape):
         test_data = {
-            "image": np.random.randint(0, 2, size=[1, 2, 3, 4]),
-            "label": np.random.randint(0, 2, size=[1, 2, 3, 4]),
-            "extra": np.random.randint(0, 2, size=[1, 2, 3, 4]),
+            "image": in_type(np.random.randint(0, 2, size=[1, 2, 3, 4])),
+            "label": in_type(np.random.randint(0, 2, size=[1, 2, 3, 4])),
+            "extra": in_type(np.random.randint(0, 2, size=[1, 2, 3, 4])),
         }
         result = AsChannelFirstd(**input_param)(test_data)
         self.assertTupleEqual(result["image"].shape, expected_shape)