Implement TorchIO transforms wrapper analogous to TorchVision transforms wrapper and test case

Author: SomeUserName1

monai/transforms/__init__.py

@@ -505,6 +505,7 @@
     ToDevice,
     ToNumpy,
     ToPIL,
+    TorchIO,
     TorchVision,
     ToTensor,
     Transpose,

monai/transforms/utility/array.py

@@ -98,6 +98,7 @@
     "ConvertToMultiChannelBasedOnBratsClasses",
     "AddExtremePointsChannel",
     "TorchVision",
+    "TorchIO",
     "MapLabelValue",
     "IntensityStats",
     "ToDevice",
@@ -1163,6 +1164,42 @@ def __call__(self, img: NdarrayOrTensor):
         return out
 
 
+class TorchIO:
+    """
+    This is a wrapper transform for TorchIO transforms based on the specified transform name and args.
+    As most of the TorchIO transforms only work for PyTorch Tensor, this transform expects input
+    data to be PyTorch Tensor, users can easily call `ToTensor` transform to convert a Numpy array to Tensor.
+
+    """
+
+    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: PyTorch Tensor data for the TorchIO transform.
+
+        """
+        img_t, *_ = convert_data_type(img, torch.Tensor)
+
+        out = self.trans(img_t)
+        out, *_ = convert_to_dst_type(src=out, dst=img)
+        return out
+
+
 class MapLabelValue:
     """
     Utility to map label values to another set of values.

tests/test_torchio.py

@@ -0,0 +1,53 @@
+# 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 parameterized import parameterized
+import numpy as np
+import torch
+
+from monai.transforms import TorchIO
+from monai.utils import set_determinism
+
+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)],
+        ]
+
+
+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()