3142 Fix mixed data type issue in transforms

monai/transforms/intensity/array.py

@@ -521,6 +521,8 @@ class RandBiasField(RandomizableTransform):
 
     """
 
+    backend = [TransformBackends.NUMPY]
+
     def __init__(
         self,
         degree: int = 3,
@@ -560,24 +562,23 @@ def _generate_random_field(self, spatial_shape: Sequence[int], degree: int, coef
             return np.polynomial.legendre.leggrid3d(coords[0], coords[1], coords[2], coeff_mat)
         raise NotImplementedError("only supports 2D or 3D fields")
 
-    def randomize(self, data: np.ndarray) -> None:
+    def randomize(self, img_size: Sequence[int]) -> None:
         super().randomize(None)
         if not self._do_transform:
             return None
-        n_coeff = int(np.prod([(self.degree + k) / k for k in range(1, len(data.shape[1:]) + 1)]))
+        n_coeff = int(np.prod([(self.degree + k) / k for k in range(1, len(img_size) + 1)]))
         self._coeff = self.R.uniform(*self.coeff_range, n_coeff).tolist()
 
-    def __call__(self, img: np.ndarray, randomize: bool = True):
+    def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor:
         """
         Apply the transform to `img`.
         """
         if randomize:
-            self.randomize(data=img)
+            self.randomize(img_size=img.shape[1:])
 
         if not self._do_transform:
             return img
 
-        img, *_ = convert_data_type(img, np.ndarray)  # type: ignore
         num_channels, *spatial_shape = img.shape
         _bias_fields = np.stack(
             [
@@ -586,7 +587,10 @@ def __call__(self, img: np.ndarray, randomize: bool = True):
             ],
             axis=0,
         )
-        return (img * np.exp(_bias_fields)).astype(self.dtype)
+        img_np, *_ = convert_data_type(img, np.ndarray)
+        out = img_np * np.exp(_bias_fields)
+        out, *_ = convert_to_dst_type(src=out, dst=img, dtype=self.dtype)
+        return out
 
 
 class NormalizeIntensity(Transform):
@@ -1784,6 +1788,8 @@ class RandCoarseTransform(RandomizableTransform):
 
     """
 
+    backend = [TransformBackends.NUMPY]
+
     def __init__(
         self,
         holes: int,
@@ -1823,15 +1829,18 @@ def _transform_holes(self, img: np.ndarray) -> np.ndarray:
         """
         raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.")
 
-    def __call__(self, img: np.ndarray, randomize: bool = True):
+    def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor:
         if randomize:
             self.randomize(img.shape[1:])
 
         if not self._do_transform:
             return img
 
-        img, *_ = convert_data_type(img, np.ndarray)  # type: ignore
-        return self._transform_holes(img=img)
+        img_np: np.ndarray
+        img_np, *_ = convert_data_type(img, np.ndarray)  # type: ignore
+        out = self._transform_holes(img=img_np)
+        ret, *_ = convert_to_dst_type(src=out, dst=img)
+        return ret
 
 
 class RandCoarseDropout(RandCoarseTransform):

monai/transforms/intensity/dictionary.py

@@ -586,6 +586,8 @@ class RandBiasFieldd(RandomizableTransform, MapTransform):
     Dictionary-based version :py:class:`monai.transforms.RandBiasField`.
     """
 
+    backend = RandBiasField.backend
+
     def __init__(
         self,
         keys: KeysCollection,
@@ -619,14 +621,14 @@ def set_random_state(
         self.rand_bias_field.set_random_state(seed, state)
         return self
 
-    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)
         self.randomize(None)
         if not self._do_transform:
             return d
 
         # all the keys share the same random bias factor
-        self.rand_bias_field.randomize(d[self.keys[0]])
+        self.rand_bias_field.randomize(img_size=d[self.keys[0]].shape[1:])
         for key in self.key_iterator(d):
             d[key] = self.rand_bias_field(d[key], randomize=False)
         return d
@@ -1429,6 +1431,8 @@ class RandCoarseDropoutd(RandomizableTransform, MapTransform):
 
     """
 
+    backend = RandCoarseDropout.backend
+
     def __init__(
         self,
         keys: KeysCollection,
@@ -1500,6 +1504,8 @@ class RandCoarseShuffled(RandomizableTransform, MapTransform):
 
     """
 
+    backend = RandCoarseShuffle.backend
+
     def __init__(
         self,
         keys: KeysCollection,

tests/test_rand_bias_field.py

@@ -12,6 +12,7 @@
 import unittest
 
 import numpy as np
+import torch
 from parameterized import parameterized
 
 from monai.transforms import RandBiasField
@@ -28,16 +29,17 @@
 class TestRandBiasField(unittest.TestCase):
     @parameterized.expand([TEST_CASES_2D, TEST_CASES_3D])
     def test_output_shape(self, class_args, img_shape):
-        for degree in [1, 2, 3]:
-            bias_field = RandBiasField(degree=degree, **class_args)
-            img = np.random.rand(*img_shape)
-            output = bias_field(img)
-            np.testing.assert_equal(output.shape, img_shape)
-            np.testing.assert_equal(output.dtype, bias_field.dtype)
+        for fn in (np.random, torch):
+            for degree in [1, 2, 3]:
+                bias_field = RandBiasField(degree=degree, **class_args)
+                img = fn.rand(*img_shape)
+                output = bias_field(img)
+                np.testing.assert_equal(output.shape, img_shape)
+                self.assertTrue(output.dtype in (np.float32, torch.float32))
 
-            img_zero = np.zeros([*img_shape])
-            output_zero = bias_field(img_zero)
-            np.testing.assert_equal(output_zero, img_zero)
+                img_zero = np.zeros([*img_shape])
+                output_zero = bias_field(img_zero)
+                np.testing.assert_equal(output_zero, img_zero)
 
     @parameterized.expand([TEST_CASES_2D_ZERO_RANGE])
     def test_zero_range(self, class_args, img_shape):

tests/test_rand_coarse_dropout.py

@@ -12,6 +12,7 @@
 import unittest
 
 import numpy as np
+import torch
 from parameterized import parameterized
 
 from monai.transforms import RandCoarseDropout
@@ -52,12 +53,29 @@
     np.random.randint(0, 2, size=[3, 3, 3, 4]),
 ]
 
+TEST_CASE_7 = [
+    {"holes": 2, "spatial_size": [2, 2, 2], "dropout_holes": False, "fill_value": (3, 6), "prob": 1.0},
+    torch.randint(0, 2, size=[3, 3, 3, 4]),
+]
+
 
 class TestRandCoarseDropout(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6])
+    @parameterized.expand(
+        [
+            TEST_CASE_0,
+            TEST_CASE_1,
+            TEST_CASE_2,
+            TEST_CASE_3,
+            TEST_CASE_4,
+            TEST_CASE_5,
+            TEST_CASE_6,
+            TEST_CASE_7,
+        ]
+    )
     def test_value(self, input_param, input_data):
         dropout = RandCoarseDropout(**input_param)
         result = dropout(input_data)
+        self.assertEqual(type(result), type(input_data))
         holes = input_param.get("holes")
         max_holes = input_param.get("max_holes")
         spatial_size = fall_back_tuple(input_param.get("spatial_size"), input_data.shape[1:])

tests/test_rand_coarse_shuffle.py

@@ -12,6 +12,7 @@
 import unittest
 
 import numpy as np
+import torch
 from parameterized import parameterized
 
 from monai.transforms import RandCoarseShuffle
@@ -40,6 +41,11 @@
         {"img": np.arange(16).reshape((2, 2, 2, 2))},
         np.asarray([[[[6, 1], [4, 3]], [[0, 2], [7, 5]]], [[[14, 10], [9, 8]], [[12, 15], [13, 11]]]]),
     ],
+    [
+        {"holes": 2, "spatial_size": 1, "max_spatial_size": -1, "prob": 1.0},
+        {"img": torch.arange(16).reshape((2, 2, 2, 2))},
+        torch.as_tensor([[[[6, 1], [4, 3]], [[0, 2], [7, 5]]], [[[14, 10], [9, 8]], [[12, 15], [13, 11]]]]),
+    ],
 ]