diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py index 999636d0b9..8d9eb374db 100644 --- a/monai/transforms/utility/array.py +++ b/monai/transforms/utility/array.py @@ -402,18 +402,25 @@ 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): """ @@ -421,7 +428,9 @@ def __call__(self, img: NdarrayOrTensor): """ 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): diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py index b65439f067..68f556cb2b 100644 --- a/monai/transforms/utility/dictionary.py +++ b/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) diff --git a/tests/test_to_tensor.py b/tests/test_to_tensor.py index fe6c4fb0d4..cd1a814f21 100644 --- a/tests/test_to_tensor.py +++ b/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)