4720 Add track_meta option to the ToTensor transform

monai/transforms/utility/array.py

@@ -402,26 +402,35 @@ class ToTensor(Transform):
         device: target device to put the converted Tensor data.
         wrap_sequence: if `False`, then lists will recursively call this function, default to `True`.
             E.g., if `False`, `[1, 2]` -> `[tensor(1), tensor(2)]`, if `True`, then `[1, 2]` -> `tensor([1, 2])`.
+        track_meta: whether to convert to `MetaTensor`, default to `False`, output type will be `torch.Tensor`.
+            if `None`, use the return value of ``get_track_meta``.
 
     """
 
     backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
 
     def __init__(
-        self, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, wrap_sequence: bool = True
+        self,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        wrap_sequence: bool = True,
+        track_meta: Optional[bool] = False,
     ) -> None:
         super().__init__()
         self.dtype = dtype
         self.device = device
         self.wrap_sequence = wrap_sequence
+        self.track_meta = get_track_meta() if track_meta is None else bool(track_meta)
 
     def __call__(self, img: NdarrayOrTensor):
         """
         Apply the transform to `img` and make it contiguous.
         """
         if isinstance(img, MetaTensor):
             img.applied_operations = []  # drops tracking info
-        return convert_to_tensor(img, dtype=self.dtype, device=self.device, wrap_sequence=self.wrap_sequence)
+        return convert_to_tensor(
+            img, dtype=self.dtype, device=self.device, wrap_sequence=self.wrap_sequence, track_meta=self.track_meta
+        )
 
 
 class EnsureType(Transform):

monai/transforms/utility/dictionary.py

@@ -485,6 +485,7 @@ def __init__(
         dtype: Optional[torch.dtype] = None,
         device: Optional[torch.device] = None,
         wrap_sequence: bool = True,
+        track_meta: Optional[bool] = False,
         allow_missing_keys: bool = False,
     ) -> None:
         """
@@ -495,10 +496,13 @@ def __init__(
             device: specify the target device to put the Tensor data.
             wrap_sequence: if `False`, then lists will recursively call this function, default to `True`.
                 E.g., if `False`, `[1, 2]` -> `[tensor(1), tensor(2)]`, if `True`, then `[1, 2]` -> `tensor([1, 2])`.
+            track_meta: whether to convert to `MetaTensor`, default to `False`, output type will be `torch.Tensor`.
+                if `None`, use the return value of ``get_track_meta``.
             allow_missing_keys: don't raise exception if key is missing.
+
         """
         super().__init__(keys, allow_missing_keys)
-        self.converter = ToTensor(dtype=dtype, device=device, wrap_sequence=wrap_sequence)
+        self.converter = ToTensor(dtype=dtype, device=device, wrap_sequence=wrap_sequence, track_meta=track_meta)
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)

tests/test_to_tensor.py

@@ -40,7 +40,7 @@ def test_array_input(self, test_data, expected_shape):
 
     @parameterized.expand(TESTS_SINGLE)
     def test_single_input(self, test_data):
-        result = ToTensor()(test_data)
+        result = ToTensor(track_meta=True)(test_data)
         self.assertTrue(isinstance(result, torch.Tensor))
         assert_allclose(result, test_data, type_test=False)
         self.assertEqual(result.ndim, 0)