Implement TorchIO transforms wrapper analogous to TorchVision transfo…

environment-dev.yml

@@ -7,6 +7,7 @@ channels:
 dependencies:
   - numpy>=1.20
   - pytorch>=1.9
+  - torchio
   - torchvision
   - pytorch-cuda=11.6
   - pip

monai/transforms/__init__.py

@@ -517,6 +517,7 @@
     ToDevice,
     ToNumpy,
     ToPIL,
+    TorchIO,
     TorchVision,
     ToTensor,
     Transpose,
@@ -627,6 +628,9 @@
     ToPILd,
     ToPILD,
     ToPILDict,
+    TorchIOd,
+    TorchIOD,
+    TorchIODict,
     TorchVisiond,
     TorchVisionD,
     TorchVisionDict,

monai/transforms/utility/array.py

@@ -98,6 +98,7 @@
     "ConvertToMultiChannelBasedOnBratsClasses",
     "AddExtremePointsChannel",
     "TorchVision",
+    "TorchIO",
     "MapLabelValue",
     "IntensityStats",
     "ToDevice",
@@ -1127,12 +1128,10 @@ def __call__(
         return concatenate((img, points_image), axis=0)
 
 
-class TorchVision:
+class TorchVision(Transform):
     """
     This is a wrapper transform for PyTorch TorchVision transform based on the specified transform name and args.
-    As most of the TorchVision transforms only work for PIL image and PyTorch Tensor, this transform expects input
-    data to be PyTorch Tensor, users can easily call `ToTensor` transform to convert a Numpy array to Tensor.
-
+    Data is converted to a torch.tensor before applying the transform and then converted back to the original data type.
     """
 
     backend = [TransformBackends.TORCH]
@@ -1163,6 +1162,36 @@ def __call__(self, img: NdarrayOrTensor):
         return out
 
 
+class TorchIO(Transform, RandomizableTrait):
+    """
+    This is a wrapper transform for TorchIO transforms based on the specified transform name and args.
+    """
+
+    backend = [TransformBackends.TORCH]
+
+    def __init__(self, name: str, *args, **kwargs) -> None:
+        """
+        Args:
+            name: The transform name in TorchIO package.
+            args: parameters for the TorchIO transform.
+            kwargs: parameters for the TorchIO transform.
+
+        """
+        super().__init__()
+        self.name = name
+        transform, _ = optional_import("torchio.transforms", "0.18.0", min_version, name=name)
+        self.trans = transform(*args, **kwargs)
+
+    def __call__(self, img: NdarrayOrTensor):
+        """
+        Args:
+            img: an instance of torchio.Subject, torchio.Image, numpy.ndarray, torch.Tensor, SimpleITK.Image,
+                 or dict containing 4D tensors as values
+
+        """
+        return self.trans(img)
+
+
 class MapLabelValue:
     """
     Utility to map label values to another set of values.

monai/transforms/utility/dictionary.py

@@ -59,6 +59,7 @@
     ToDevice,
     ToNumpy,
     ToPIL,
+    TorchIO,
     TorchVision,
     ToTensor,
     Transpose,
@@ -171,6 +172,9 @@
     "ToTensorD",
     "ToTensorDict",
     "ToTensord",
+    "TorchIOD",
+    "TorchIODict",
+    "TorchIOd",
     "TorchVisionD",
     "TorchVisionDict",
     "TorchVisiond",
@@ -1419,6 +1423,36 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
         return d
 
 
+class TorchIOd(MapTransform, RandomizableTrait):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.TorchIO` for transforms.
+    All transforms in TorchIO can be applied randomly with probability p by specifying the `p=` argument.
+    """
+
+    backend = TorchIO.backend
+
+    def __init__(self, keys: KeysCollection, name: str, allow_missing_keys: bool = False, *args, **kwargs) -> None:
+        """
+        Args:
+            keys: keys of the corresponding items to be transformed.
+                See also: :py:class:`monai.transforms.compose.MapTransform`
+            name: The transform name in TorchIO package.
+            allow_missing_keys: don't raise exception if key is missing.
+            args: parameters for the TorchIO transform.
+            kwargs: parameters for the TorchIO transform.
+
+        """
+        super().__init__(keys, allow_missing_keys)
+        self.name = name
+        self.trans = TorchIO(name, *args, **kwargs)
+
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
+        d = dict(data)
+        for key in self.key_iterator(d):
+            d[key] = self.trans(d[key])
+        return d
+
+
 class MapLabelValued(MapTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.MapLabelValue`.
@@ -1771,6 +1805,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
 )
 AddExtremePointsChannelD = AddExtremePointsChannelDict = AddExtremePointsChanneld
 TorchVisionD = TorchVisionDict = TorchVisiond
+TorchIOD = TorchIODict = TorchIOd
 RandTorchVisionD = RandTorchVisionDict = RandTorchVisiond
 RandLambdaD = RandLambdaDict = RandLambdad
 MapLabelValueD = MapLabelValueDict = MapLabelValued

requirements-dev.txt

@@ -24,6 +24,7 @@ pytype>=2020.6.1; platform_system != "Windows"
 types-pkg_resources
 mypy>=1.5.0
 ninja
+torchio
 torchvision
 psutil
 cucim-cu12; platform_system == "Linux" and python_version >= "3.9" and python_version <= "3.10"

setup.cfg

@@ -54,6 +54,7 @@ all =
     tensorboard
     gdown>=4.7.3
     pytorch-ignite==0.4.11
+    torchio
     torchvision
     itk>=5.2
     tqdm>=4.47.0
@@ -100,6 +101,8 @@ gdown =
     gdown>=4.7.3
 ignite =
     pytorch-ignite==0.4.11
+torchio =
+    torchio
 torchvision =
     torchvision
 itk =

tests/test_torchio.py

@@ -0,0 +1,56 @@
+# Copyright (c) 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.
+
+from __future__ import annotations
+
+import unittest
+from unittest import skipUnless
+
+import numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.transforms import TorchIO
+from monai.utils import optional_import, set_determinism
+
+_, has_torchio = optional_import("torchio")
+
+TEST_DIMS = [3, 128, 160, 160]
+TESTS = [
+    [{"name": "RescaleIntensity"}, torch.rand(TEST_DIMS)],
+    [{"name": "ZNormalization"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomAffine"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomElasticDeformation"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomAnisotropy"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomMotion"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomGhosting"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomSpike"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomBiasField"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomBlur"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomNoise"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomSwap"}, torch.rand(TEST_DIMS)],
+    [{"name": "RandomGamma"}, torch.rand(TEST_DIMS)],
+]
+
+
+@skipUnless(has_torchio, "Requires torchio")
+class TestTorchIO(unittest.TestCase):
+
+    @parameterized.expand(TESTS)
+    def test_value(self, input_param, input_data):
+        set_determinism(seed=0)
+        result = TorchIO(**input_param)(input_data)
+        self.assertIsNotNone(result)
+        self.assertFalse(np.array_equal(result.numpy(), input_data.numpy()), f"{input_param} failed")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_torchiod.py

@@ -0,0 +1,48 @@
+# Copyright (c) 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.
+
+from __future__ import annotations
+
+import unittest
+from unittest import skipUnless
+
+import torch
+from parameterized import parameterized
+
+from monai.transforms import TorchIOd
+from monai.utils import optional_import, set_determinism
+from tests.utils import assert_allclose
+
+_, has_torchio = optional_import("torchio")
+
+TEST_DIMS = [3, 128, 160, 160]
+TEST_TENSOR = torch.rand(TEST_DIMS)
+TESTS = [
+    [
+        {"keys": "img", "name": "RescaleIntensity", "out_min_max": (0, 42)},
+        {"img": TEST_TENSOR},
+        ((TEST_TENSOR - TEST_TENSOR.min()) / (TEST_TENSOR.max() - TEST_TENSOR.min())) * 42,
+    ]
+]
+
+
+@skipUnless(has_torchio, "Requires torchio")
+class TestTorchVisiond(unittest.TestCase):
+
+    @parameterized.expand(TESTS)
+    def test_value(self, input_param, input_data, expected_value):
+        set_determinism(seed=0)
+        result = TorchIOd(**input_param)(input_data)
+        assert_allclose(result["img"], expected_value, atol=1e-4, rtol=1e-4, type_test=False)
+
+
+if __name__ == "__main__":
+    unittest.main()