4741 Support MetaTensor data type in Concatitemd (#4745)

* add Metatenosr support in ConcatItemd

Signed-off-by: KumoLiu <[email protected]>

Author: KumoLiu

Diff for: monai/transforms/utility/dictionary.py

@@ -929,7 +929,8 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
 class ConcatItemsd(MapTransform):
     """
     Concatenate specified items from data dictionary together on the first dim to construct a big array.
-    Expect all the items are numpy array or PyTorch Tensor.
+    Expect all the items are numpy array or PyTorch Tensor or MetaTensor.
+    Return the first input's meta information when items are MetaTensor.
     """
 
     backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
@@ -951,7 +952,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
         """
         Raises:
             TypeError: When items in ``data`` differ in type.
-            TypeError: When the item type is not in ``Union[numpy.ndarray, torch.Tensor]``.
+            TypeError: When the item type is not in ``Union[numpy.ndarray, torch.Tensor, MetaTensor]``.
 
         """
         d = dict(data)
@@ -969,10 +970,12 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
 
         if data_type is np.ndarray:
             d[self.name] = np.concatenate(output, axis=self.dim)
-        elif data_type is torch.Tensor:
+        elif issubclass(data_type, torch.Tensor):  # type: ignore
             d[self.name] = torch.cat(output, dim=self.dim)  # type: ignore
         else:
-            raise TypeError(f"Unsupported data type: {data_type}, available options are (numpy.ndarray, torch.Tensor).")
+            raise TypeError(
+                f"Unsupported data type: {data_type}, available options are (numpy.ndarray, torch.Tensor, MetaTensor)."
+            )
         return d
 
 

Diff for: tests/test_concat_itemsd.py

@@ -14,6 +14,7 @@
 import numpy as np
 import torch
 
+from monai.data import MetaTensor
 from monai.transforms import ConcatItemsd
 
 
@@ -30,6 +31,20 @@ def test_tensor_values(self):
         torch.testing.assert_allclose(result["img1"], torch.tensor([[0, 1], [1, 2]], device=device))
         torch.testing.assert_allclose(result["cat_img"], torch.tensor([[1, 2], [2, 3], [1, 2], [2, 3]], device=device))
 
+    def test_metatensor_values(self):
+        device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu:0")
+        input_data = {
+            "img1": MetaTensor([[0, 1], [1, 2]], device=device),
+            "img2": MetaTensor([[0, 1], [1, 2]], device=device),
+        }
+        result = ConcatItemsd(keys=["img1", "img2"], name="cat_img")(input_data)
+        self.assertTrue("cat_img" in result)
+        self.assertTrue(isinstance(result["cat_img"], MetaTensor))
+        self.assertEqual(result["img1"].meta, result["cat_img"].meta)
+        result["cat_img"] += 1
+        torch.testing.assert_allclose(result["img1"], torch.tensor([[0, 1], [1, 2]], device=device))
+        torch.testing.assert_allclose(result["cat_img"], torch.tensor([[1, 2], [2, 3], [1, 2], [2, 3]], device=device))
+
     def test_numpy_values(self):
         input_data = {"img1": np.array([[0, 1], [1, 2]]), "img2": np.array([[0, 1], [1, 2]])}
         result = ConcatItemsd(keys=["img1", "img2"], name="cat_img")(input_data)
@@ -52,6 +67,13 @@ def test_single_tensor(self):
         torch.testing.assert_allclose(result["img"], torch.tensor([[0, 1], [1, 2]]))
         torch.testing.assert_allclose(result["cat_img"], torch.tensor([[1, 2], [2, 3]]))
 
+    def test_single_metatensor(self):
+        input_data = {"img": MetaTensor([[0, 1], [1, 2]])}
+        result = ConcatItemsd(keys="img", name="cat_img")(input_data)
+        result["cat_img"] += 1
+        torch.testing.assert_allclose(result["img"], torch.tensor([[0, 1], [1, 2]]))
+        torch.testing.assert_allclose(result["cat_img"], torch.tensor([[1, 2], [2, 3]]))
+
 
 if __name__ == "__main__":
     unittest.main()