3142 Fix mixed output data type in transforms

monai/transforms/intensity/array.py

@@ -224,9 +224,9 @@ def __call__(self, img: NdarrayOrTensor, offset: Optional[float] = None) -> Ndar
 
         offset = self.offset if offset is None else offset
         out = img + offset
-        if isinstance(out, torch.Tensor):
-            return out.type(img.dtype)
-        return out.astype(img.dtype)  # type: ignore
+        out, *_ = convert_data_type(data=out, dtype=img.dtype)
+
+        return out
 
 
 class RandShiftIntensity(RandomizableTransform):
@@ -941,11 +941,13 @@ class MaskIntensity(Transform):
 
     """
 
-    def __init__(self, mask_data: Optional[np.ndarray] = None, select_fn: Callable = is_positive) -> None:
+    backend = [TransformBackends.NUMPY]
+
+    def __init__(self, mask_data: Optional[NdarrayOrTensor] = None, select_fn: Callable = is_positive) -> None:
         self.mask_data = mask_data
         self.select_fn = select_fn
 
-    def __call__(self, img: np.ndarray, mask_data: Optional[np.ndarray] = None) -> np.ndarray:
+    def __call__(self, img: NdarrayOrTensor, mask_data: Optional[NdarrayOrTensor] = None) -> NdarrayOrTensor:
         """
         Args:
             mask_data: if mask data is single channel, apply to every channel
@@ -958,21 +960,20 @@ def __call__(self, img: np.ndarray, mask_data: Optional[np.ndarray] = None) -> n
             - ValueError: When ``mask_data`` and ``img`` channels differ and ``mask_data`` is not single channel.
 
         """
-        img, *_ = convert_data_type(img, np.ndarray)  # type: ignore
         mask_data = self.mask_data if mask_data is None else mask_data
         if mask_data is None:
             raise ValueError("must provide the mask_data when initializing the transform or at runtime.")
 
-        mask_data, *_ = convert_data_type(mask_data, np.ndarray)  # type: ignore
+        mask_data_, *_ = convert_to_dst_type(src=mask_data, dst=img)
 
-        mask_data = np.asarray(self.select_fn(mask_data))
-        if mask_data.shape[0] != 1 and mask_data.shape[0] != img.shape[0]:
+        mask_data_ = self.select_fn(mask_data_)
+        if mask_data_.shape[0] != 1 and mask_data_.shape[0] != img.shape[0]:
             raise ValueError(
                 "When mask_data is not single channel, mask_data channels must match img, "
-                f"got img channels={img.shape[0]} mask_data channels={mask_data.shape[0]}."
+                f"got img channels={img.shape[0]} mask_data channels={mask_data_.shape[0]}."
             )
 
-        return np.asarray(img * mask_data)
+        return img * mask_data_
 
 
 class SavitzkyGolaySmooth(Transform):
@@ -1032,6 +1033,8 @@ class DetectEnvelope(Transform):
 
     """
 
+    backend = [TransformBackends.TORCH]
+
     def __init__(self, axis: int = 1, n: Union[int, None] = None) -> None:
 
         if PT_BEFORE_1_7:
@@ -1043,7 +1046,7 @@ def __init__(self, axis: int = 1, n: Union[int, None] = None) -> None:
         self.axis = axis
         self.n = n
 
-    def __call__(self, img: np.ndarray):
+    def __call__(self, img: NdarrayOrTensor):
         """
 
         Args:
@@ -1053,12 +1056,15 @@ def __call__(self, img: np.ndarray):
             np.ndarray containing envelope of data in img along the specified axis.
 
         """
-        img, *_ = convert_data_type(img, np.ndarray)  # type: ignore
+        img_t: torch.Tensor
+        img_t, *_ = convert_data_type(img, torch.Tensor)  # type: ignore
         # add one to transform axis because a batch axis will be added at dimension 0
         hilbert_transform = HilbertTransform(self.axis + 1, self.n)
         # convert to Tensor and add Batch axis expected by HilbertTransform
-        input_data = torch.as_tensor(np.ascontiguousarray(img)).unsqueeze(0)
-        return np.abs(hilbert_transform(input_data).squeeze(0).numpy())
+        out = hilbert_transform(img_t.unsqueeze(0)).squeeze(0).abs()
+        out, *_ = convert_to_dst_type(src=out, dst=img)
+
+        return out
 
 
 class GaussianSmooth(Transform):

monai/transforms/intensity/dictionary.py

@@ -871,10 +871,12 @@ class MaskIntensityd(MapTransform):
 
     """
 
+    backend = MaskIntensity.backend
+
     def __init__(
         self,
         keys: KeysCollection,
-        mask_data: Optional[np.ndarray] = None,
+        mask_data: Optional[NdarrayOrTensor] = None,
         mask_key: Optional[str] = None,
         select_fn: Callable = is_positive,
         allow_missing_keys: bool = False,
@@ -883,7 +885,7 @@ def __init__(
         self.converter = MaskIntensity(mask_data=mask_data, select_fn=select_fn)
         self.mask_key = mask_key if mask_data is None else None
 
-    def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
         for key in self.key_iterator(d):
             d[key] = self.converter(d[key], d[self.mask_key]) if self.mask_key is not None else self.converter(d[key])

monai/transforms/spatial/array.py

@@ -185,11 +185,11 @@ def __call__(
             raise ValueError("data_array must have at least one spatial dimension.")
         if affine is None:
             # default to identity
-            affine = np.eye(sr + 1, dtype=np.float64)
+            affine_np = affine = np.eye(sr + 1, dtype=np.float64)
             affine_ = np.eye(sr + 1, dtype=np.float64)
         else:
-            affine, *_ = convert_data_type(affine, np.ndarray)
-            affine_ = to_affine_nd(sr, affine)  # type: ignore
+            affine_np, *_ = convert_data_type(affine, np.ndarray)  # type: ignore
+            affine_ = to_affine_nd(sr, affine_np)
 
         out_d = self.pixdim[:sr]
         if out_d.size < sr:
@@ -205,8 +205,7 @@ def __call__(
 
         # no resampling if it's identity transform
         if np.allclose(transform, np.diag(np.ones(len(transform))), atol=1e-3):
-            output_data, *_ = convert_data_type(data_array, dtype=torch.float32)
-            new_affine = to_affine_nd(affine, new_affine)  # type: ignore
+            output_data = data_array
         else:
             # resample
             affine_xform = AffineTransform(
@@ -224,8 +223,10 @@ def __call__(
                 convert_data_type(transform, torch.Tensor, data_array_t.device, dtype=_dtype)[0],
                 spatial_size=output_shape if output_spatial_shape is None else output_spatial_shape,
             ).squeeze(0)
-            output_data, *_ = convert_to_dst_type(output_data, data_array, dtype=torch.float32)
-            new_affine = to_affine_nd(affine, new_affine)  # type: ignore
+
+        output_data, *_ = convert_to_dst_type(output_data, data_array, dtype=torch.float32)
+        new_affine = to_affine_nd(affine_np, new_affine)  # type: ignore
+        new_affine, *_ = convert_to_dst_type(src=new_affine, dst=affine, dtype=torch.float32)
 
         if self.image_only:
             return output_data
@@ -237,6 +238,8 @@ class Orientation(Transform):
     Change the input image's orientation into the specified based on `axcodes`.
     """
 
+    backend = [TransformBackends.NUMPY]
+
     def __init__(
         self,
         axcodes: Optional[str] = None,
@@ -273,8 +276,8 @@ def __init__(
         self.image_only = image_only
 
     def __call__(
-        self, data_array: np.ndarray, affine: Optional[np.ndarray] = None
-    ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray, np.ndarray]]:
+        self, data_array: NdarrayOrTensor, affine: Optional[NdarrayOrTensor] = None
+    ) -> Union[NdarrayOrTensor, Tuple[NdarrayOrTensor, NdarrayOrTensor, NdarrayOrTensor]]:
         """
         original orientation of `data_array` is defined by `affine`.
 
@@ -291,15 +294,18 @@ def __call__(
             (data_array [reoriented in `self.axcodes`], original axcodes, current axcodes).
 
         """
-        data_array, *_ = convert_data_type(data_array, np.ndarray)  # type: ignore
-        sr = data_array.ndim - 1
+        data_array_np, *_ = convert_data_type(data_array, np.ndarray)  # type: ignore
+        sr = data_array_np.ndim - 1
         if sr <= 0:
             raise ValueError("data_array must have at least one spatial dimension.")
         if affine is None:
-            affine = np.eye(sr + 1, dtype=np.float64)
+            # default to identity
+            affine_np = affine = np.eye(sr + 1, dtype=np.float64)
             affine_ = np.eye(sr + 1, dtype=np.float64)
         else:
-            affine_ = to_affine_nd(sr, affine)
+            affine_np, *_ = convert_data_type(affine, np.ndarray)  # type: ignore
+            affine_ = to_affine_nd(sr, affine_np)
+
         src = nib.io_orientation(affine_)
         if self.as_closest_canonical:
             spatial_ornt = src
@@ -315,14 +321,16 @@ def __call__(
         ornt = spatial_ornt.copy()
         ornt[:, 0] += 1  # skip channel dim
         ornt = np.concatenate([np.array([[0, 1]]), ornt])
-        shape = data_array.shape[1:]
-        data_array = np.ascontiguousarray(nib.orientations.apply_orientation(data_array, ornt))
+        shape = data_array_np.shape[1:]
+        data_array_np = np.ascontiguousarray(nib.orientations.apply_orientation(data_array_np, ornt))
         new_affine = affine_ @ nib.orientations.inv_ornt_aff(spatial_ornt, shape)
-        new_affine = to_affine_nd(affine, new_affine)
+        new_affine = to_affine_nd(affine_np, new_affine)
+        out, *_ = convert_to_dst_type(src=data_array_np, dst=data_array)
+        new_affine, *_ = convert_to_dst_type(src=new_affine, dst=affine, dtype=torch.float32)
 
         if self.image_only:
-            return data_array
-        return data_array, affine, new_affine
+            return out
+        return out, affine, new_affine
 
 
 class Flip(Transform):

monai/transforms/spatial/dictionary.py

@@ -226,7 +226,7 @@ def __init__(
 
     def __call__(
         self, data: Mapping[Union[Hashable, str], Dict[str, NdarrayOrTensor]]
-    ) -> Dict[Union[Hashable, str], Union[NdarrayOrTensor, Dict[str, NdarrayOrTensor]]]:
+    ) -> Dict[Hashable, NdarrayOrTensor]:
         d: Dict = dict(data)
         for key, mode, padding_mode, align_corners, dtype, meta_key, meta_key_postfix in self.key_iterator(
             d, self.mode, self.padding_mode, self.align_corners, self.dtype, self.meta_keys, self.meta_key_postfix
@@ -309,6 +309,8 @@ class Orientationd(MapTransform, InvertibleTransform):
     to the `affine` field of metadata which is formed by ``key_{meta_key_postfix}``.
     """
 
+    backend = Orientation.backend
+
     def __init__(
         self,
         keys: KeysCollection,
@@ -358,8 +360,8 @@ def __init__(
         self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys))
 
     def __call__(
-        self, data: Mapping[Union[Hashable, str], Dict[str, np.ndarray]]
-    ) -> Dict[Union[Hashable, str], Union[np.ndarray, Dict[str, np.ndarray]]]:
+        self, data: Mapping[Union[Hashable, str], Dict[str, NdarrayOrTensor]]
+    ) -> Dict[Hashable, NdarrayOrTensor]:
         d: Dict = dict(data)
         for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix):
             meta_key = meta_key or f"{key}_{meta_key_postfix}"
@@ -372,12 +374,12 @@ def __call__(
             d[meta_key]["affine"] = new_affine
         return d
 
-    def inverse(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
+    def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = deepcopy(dict(data))
         for key in self.key_iterator(d):
             transform = self.get_most_recent_transform(d, key)
             # Create inverse transform
-            meta_data = d[transform[InverseKeys.EXTRA_INFO]["meta_key"]]
+            meta_data: Dict = d[transform[InverseKeys.EXTRA_INFO]["meta_key"]]  # type: ignore
             orig_affine = transform[InverseKeys.EXTRA_INFO]["old_affine"]
             orig_axcodes = nib.orientations.aff2axcodes(orig_affine)
             inverse_transform = Orientation(
@@ -708,7 +710,7 @@ class RandAffined(RandomizableTransform, MapTransform, InvertibleTransform):
     Dictionary-based wrapper of :py:class:`monai.transforms.RandAffine`.
     """
 
-    backend = Affine.backend
+    backend = RandAffine.backend
 
     @deprecated_arg(name="as_tensor_output", since="0.6")
     def __init__(
@@ -1371,8 +1373,9 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd
             transform_t: torch.Tensor
             transform_t, *_ = convert_to_dst_type(inv_rot_mat, img_t)  # type: ignore
 
-            output = xform(img_t.unsqueeze(0), transform_t, spatial_size=transform[InverseKeys.ORIG_SIZE])
-            d[key] = output.squeeze(0).detach().float()
+            out = xform(img_t.unsqueeze(0), transform_t, spatial_size=transform[InverseKeys.ORIG_SIZE]).squeeze(0)
+            out, *_ = convert_to_dst_type(out, dst=d[key], dtype=out.dtype)
+            d[key] = out
             # Remove the applied transform
             self.pop_transform(d, key)
 
@@ -1504,8 +1507,9 @@ def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, Nd
                 transform_t: torch.Tensor
                 transform_t, *_ = convert_to_dst_type(inv_rot_mat, img_t)  # type: ignore
                 output: torch.Tensor
-                output = xform(img_t.unsqueeze(0), transform_t, spatial_size=transform[InverseKeys.ORIG_SIZE])
-                d[key] = output.squeeze(0).detach().float()
+                out = xform(img_t.unsqueeze(0), transform_t, spatial_size=transform[InverseKeys.ORIG_SIZE]).squeeze(0)
+                out, *_ = convert_to_dst_type(out, dst=d[key], dtype=out.dtype)
+                d[key] = out
             # Remove the applied transform
             self.pop_transform(d, key)
 

tests/test_detect_envelope.py

@@ -12,6 +12,7 @@
 import unittest
 
 import numpy as np
+import torch
 from parameterized import parameterized
 
 from monai.transforms import DetectEnvelope
@@ -70,9 +71,9 @@
 TEST_CASE_2_CHAN_3D_SINE = [
     {},  # args (empty, so use default (i.e. process along first spatial dimension, axis=1)
     # Create 100 identical windowed sine waves as a (n_samples x 10 x 10) 3D numpy array, twice (2 channels)
-    np.stack([np.stack([np.stack([hann_windowed_sine] * 10, axis=1)] * 10, axis=2)] * 2, axis=0),
+    torch.as_tensor(np.stack([np.stack([np.stack([hann_windowed_sine] * 10, axis=1)] * 10, axis=2)] * 2, axis=0)),
     # Expected output: Set of 100 identical Hann windows in (n_samples x 10 x 10) 3D numpy array, twice (2 channels)
-    np.stack([np.stack([np.stack([np.hanning(n_samples)] * 10, axis=1)] * 10, axis=2)] * 2, axis=0),
+    torch.as_tensor(np.stack([np.stack([np.stack([np.hanning(n_samples)] * 10, axis=1)] * 10, axis=2)] * 2, axis=0)),
     1e-4,  # absolute tolerance
 ]
 

tests/test_invertd.py

@@ -36,7 +36,7 @@
     Spacingd,
     ToTensord,
 )
-from monai.utils.misc import set_determinism
+from monai.utils import set_determinism
 from tests.utils import make_nifti_image
 
 KEYS = ["image", "label"]
@@ -134,7 +134,7 @@ def test_invert(self):
                 torch.testing.assert_allclose(i.to(torch.uint8).to(torch.float), i.to(torch.float))
                 self.assertTupleEqual(i.shape[1:], (100, 101, 107))
                 i = item["label_inverted"]
-                np.testing.assert_allclose(i.astype(np.uint8).astype(np.float32), i.astype(np.float32))
+                torch.testing.assert_allclose(i.to(torch.uint8).to(torch.float), i.to(torch.float))
                 self.assertTupleEqual(i.shape[1:], (100, 101, 107))
                 # test inverted test_dict
                 self.assertTrue(isinstance(item["test_dict"]["affine"], np.ndarray))
@@ -152,7 +152,7 @@ def test_invert(self):
                 self.assertTupleEqual(d.shape, (1, 100, 101, 107))
 
         # check labels match
-        reverted = item["label_inverted"].astype(np.int32)
+        reverted = item["label_inverted"].detach().cpu().numpy().astype(np.int32)
         original = LoadImaged(KEYS)(data[-1])["label"]
         n_good = np.sum(np.isclose(reverted, original, atol=1e-3))
         reverted_name = item["label_inverted_meta_dict"]["filename_or_obj"]

tests/test_mask_intensity.py

@@ -12,6 +12,7 @@
 import unittest
 
 import numpy as np
+import torch
 from parameterized import parameterized
 
 from monai.transforms import MaskIntensity
@@ -43,9 +44,15 @@
     np.array([[[0, 0, 0], [2, 2, 2], [0, 0, 0]], [[0, 0, 0], [5, 5, 5], [0, 0, 0]]]),
 ]
 
+TEST_CASE_5 = [
+    {"mask_data": np.array([[[0, 0, 0], [0, 1, 0], [0, 0, 0]], [[0, 1, 0], [0, 1, 0], [0, 1, 0]]])},
+    torch.as_tensor([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], [[4, 4, 4], [5, 5, 5], [6, 6, 6]]]),
+    torch.as_tensor([[[0, 0, 0], [0, 2, 0], [0, 0, 0]], [[0, 4, 0], [0, 5, 0], [0, 6, 0]]]),
+]
+
 
 class TestMaskIntensity(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
     def test_value(self, argments, image, expected_data):
         result = MaskIntensity(**argments)(image)
         np.testing.assert_allclose(result, expected_data)

tests/test_spacingd.py

@@ -17,7 +17,7 @@
 from parameterized import parameterized
 
 from monai.transforms import Spacingd
-from tests.utils import TEST_NDARRAYS
+from tests.utils import TEST_NDARRAYS, assert_allclose
 
 TESTS: List[Tuple] = []
 for p in TEST_NDARRAYS:
@@ -28,7 +28,7 @@
             dict(keys="image", pixdim=(1, 2, 1.4)),
             ("image", "image_meta_dict", "image_transforms"),
             (2, 10, 8, 15),
-            np.diag([1, 2, 1.4, 1.0]),
+            p(np.diag([1, 2, 1.4, 1.0])),
         )
     )
     TESTS.append(
@@ -91,7 +91,7 @@ def test_spacingd(self, _, data, kw_args, expected_keys, expected_shape, expecte
             self.assertEqual(data["image"].device, res["image"].device)
         self.assertEqual(expected_keys, tuple(sorted(res)))
         np.testing.assert_allclose(res["image"].shape, expected_shape)
-        np.testing.assert_allclose(res["image_meta_dict"]["affine"], expected_affine)
+        assert_allclose(res["image_meta_dict"]["affine"], expected_affine)
 
 
 if __name__ == "__main__":