Repository: Project-MONAI/MONAI

Issue:
Add backends to `ClassesToIndices`
**Is your feature request related to a problem? Please describe.**
As we already added backends to `FgBgToIndices`, also need to enhance `ClassesToIndices`.




Discussion:

Patch:
diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py
--- a/monai/transforms/utility/array.py
+++ b/monai/transforms/utility/array.py
@@ -819,6 +819,9 @@ def __call__(
 
 
 class ClassesToIndices(Transform):
+
+    backend = [TransformBackends.NUMPY, TransformBackends.TORCH]
+
     def __init__(
         self,
         num_classes: Optional[int] = None,
@@ -845,10 +848,10 @@ def __init__(
 
     def __call__(
         self,
-        label: np.ndarray,
-        image: Optional[np.ndarray] = None,
+        label: NdarrayOrTensor,
+        image: Optional[NdarrayOrTensor] = None,
         output_shape: Optional[Sequence[int]] = None,
-    ) -> List[np.ndarray]:
+    ) -> List[NdarrayOrTensor]:
         """
         Args:
             label: input data to compute the indices of every class.
@@ -857,16 +860,13 @@ def __call__(
             output_shape: expected shape of output indices. if None, use `self.output_shape` instead.
 
         """
-        label, *_ = convert_data_type(label, np.ndarray)  # type: ignore
-        if image is not None:
-            image, *_ = convert_data_type(image, np.ndarray)  # type: ignore
 
         if output_shape is None:
             output_shape = self.output_shape
-        indices: List[np.ndarray]
-        indices = map_classes_to_indices(label, self.num_classes, image, self.image_threshold)  # type: ignore
+        indices: List[NdarrayOrTensor]
+        indices = map_classes_to_indices(label, self.num_classes, image, self.image_threshold)
         if output_shape is not None:
-            indices = [np.stack([np.unravel_index(i, output_shape) for i in array]) for array in indices]
+            indices = [unravel_indices(cls_indices, output_shape) for cls_indices in indices]
 
         return indices
 
diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py
--- a/monai/transforms/utility/dictionary.py
+++ b/monai/transforms/utility/dictionary.py
@@ -1180,7 +1180,7 @@ def __init__(
         self.image_key = image_key
         self.converter = ClassesToIndices(num_classes, image_threshold, output_shape)
 
-    def __call__(self, data: Mapping[Hashable, Any]) -> Dict[Hashable, np.ndarray]:
+    def __call__(self, data: Mapping[Hashable, Any]):
         d = dict(data)
         image = d[self.image_key] if self.image_key else None
         for key in self.key_iterator(d):
diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py
--- a/monai/transforms/utils.py
+++ b/monai/transforms/utils.py
@@ -339,7 +339,9 @@ def map_classes_to_indices(
     for c in range(num_classes_):
         label_flat = ravel(any_np_pt(label[c : c + 1] if channels > 1 else label == c, 0))
         label_flat = img_flat & label_flat if img_flat is not None else label_flat
-        indices.append(nonzero(label_flat))
+        # no need to save the indices in GPU, otherwise, still need to move to CPU at runtime when crop by indices
+        cls_indices, *_ = convert_data_type(nonzero(label_flat), device=torch.device("cpu"))
+        indices.append(cls_indices)
 
     return indices
 


Test patch:
diff --git a/tests/test_classes_to_indices.py b/tests/test_classes_to_indices.py
--- a/tests/test_classes_to_indices.py
+++ b/tests/test_classes_to_indices.py
@@ -11,68 +11,80 @@
 
 import unittest
 
-import numpy as np
 from parameterized import parameterized
 
 from monai.transforms import ClassesToIndices
+from tests.utils import TEST_NDARRAYS, assert_allclose
 
-TEST_CASE_1 = [
-    # test Argmax data
-    {"num_classes": 3, "image_threshold": 0.0},
-    np.array([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]]),
-    None,
-    [np.array([0, 4, 8]), np.array([1, 5, 6]), np.array([2, 3, 7])],
-]
+TESTS_CASES = []
+for p in TEST_NDARRAYS:
+    TESTS_CASES.append(
+        [
+            # test Argmax data
+            {"num_classes": 3, "image_threshold": 0.0},
+            p([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]]),
+            None,
+            [p([0, 4, 8]), p([1, 5, 6]), p([2, 3, 7])],
+        ]
+    )
 
-TEST_CASE_2 = [
-    {"num_classes": 3, "image_threshold": 60},
-    np.array([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]]),
-    np.array([[[132, 1434, 51], [61, 0, 133], [523, 44, 232]]]),
-    [np.array([0, 8]), np.array([1, 5, 6]), np.array([3])],
-]
+    TESTS_CASES.append(
+        [
+            {"num_classes": 3, "image_threshold": 60},
+            p([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]]),
+            p([[[132, 1434, 51], [61, 0, 133], [523, 44, 232]]]),
+            [p([0, 8]), p([1, 5, 6]), p([3])],
+        ]
+    )
 
-TEST_CASE_3 = [
-    # test One-Hot data
-    {"image_threshold": 0.0},
-    np.array(
+    TESTS_CASES.append(
         [
-            [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
-            [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
-            [[0, 0, 1], [1, 0, 0], [0, 1, 0]],
+            # test One-Hot data
+            {"image_threshold": 0.0},
+            p(
+                [
+                    [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                    [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
+                    [[0, 0, 1], [1, 0, 0], [0, 1, 0]],
+                ]
+            ),
+            None,
+            [p([0, 4, 8]), p([1, 5, 6]), p([2, 3, 7])],
         ]
-    ),
-    None,
-    [np.array([0, 4, 8]), np.array([1, 5, 6]), np.array([2, 3, 7])],
-]
+    )
 
-TEST_CASE_4 = [
-    {"num_classes": None, "image_threshold": 60},
-    np.array(
+    TESTS_CASES.append(
         [
-            [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
-            [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
-            [[0, 0, 1], [1, 0, 0], [0, 1, 0]],
+            {"num_classes": None, "image_threshold": 60},
+            p(
+                [
+                    [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                    [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
+                    [[0, 0, 1], [1, 0, 0], [0, 1, 0]],
+                ]
+            ),
+            p([[[132, 1434, 51], [61, 0, 133], [523, 44, 232]]]),
+            [p([0, 8]), p([1, 5, 6]), p([3])],
         ]
-    ),
-    np.array([[[132, 1434, 51], [61, 0, 133], [523, 44, 232]]]),
-    [np.array([0, 8]), np.array([1, 5, 6]), np.array([3])],
-]
+    )
 
-TEST_CASE_5 = [
-    # test output_shape
-    {"num_classes": 3, "image_threshold": 0.0, "output_shape": [3, 3]},
-    np.array([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]]),
-    None,
-    [np.array([[0, 0], [1, 1], [2, 2]]), np.array([[0, 1], [1, 2], [2, 0]]), np.array([[0, 2], [1, 0], [2, 1]])],
-]
+    TESTS_CASES.append(
+        [
+            # test output_shape
+            {"num_classes": 3, "image_threshold": 0.0, "output_shape": [3, 3]},
+            p([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]]),
+            None,
+            [p([[0, 0], [1, 1], [2, 2]]), p([[0, 1], [1, 2], [2, 0]]), p([[0, 2], [1, 0], [2, 1]])],
+        ]
+    )
 
 
 class TestClassesToIndices(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
+    @parameterized.expand(TESTS_CASES)
     def test_value(self, input_args, label, image, expected_indices):
         indices = ClassesToIndices(**input_args)(label, image)
         for i, e in zip(indices, expected_indices):
-            np.testing.assert_allclose(i, e)
+            assert_allclose(i, e)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_classes_to_indicesd.py b/tests/test_classes_to_indicesd.py
--- a/tests/test_classes_to_indicesd.py
+++ b/tests/test_classes_to_indicesd.py
@@ -11,73 +11,91 @@
 
 import unittest
 
-import numpy as np
 from parameterized import parameterized
 
 from monai.transforms import ClassesToIndicesd
+from tests.utils import TEST_NDARRAYS, assert_allclose
 
-TEST_CASE_1 = [
-    # test Argmax data
-    {"keys": "label", "num_classes": 3, "image_threshold": 0.0},
-    {"label": np.array([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]])},
-    [np.array([0, 4, 8]), np.array([1, 5, 6]), np.array([2, 3, 7])],
-]
+TESTS_CASES = []
+for p in TEST_NDARRAYS:
+    TESTS_CASES.append(
+        [
+            # test Argmax data
+            {"keys": "label", "num_classes": 3, "image_threshold": 0.0},
+            {"label": p([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]])},
+            [p([0, 4, 8]), p([1, 5, 6]), p([2, 3, 7])],
+        ]
+    )
 
-TEST_CASE_2 = [
-    {"keys": "label", "image_key": "image", "num_classes": 3, "image_threshold": 60},
-    {
-        "label": np.array([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]]),
-        "image": np.array([[[132, 1434, 51], [61, 0, 133], [523, 44, 232]]]),
-    },
-    [np.array([0, 8]), np.array([1, 5, 6]), np.array([3])],
-]
+    TESTS_CASES.append(
+        [
+            {"keys": "label", "image_key": "image", "num_classes": 3, "image_threshold": 60},
+            {
+                "label": p([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]]),
+                "image": p([[[132, 1434, 51], [61, 0, 133], [523, 44, 232]]]),
+            },
+            [p([0, 8]), p([1, 5, 6]), p([3])],
+        ]
+    )
 
-TEST_CASE_3 = [
-    # test One-Hot data
-    {"keys": "label", "image_threshold": 0.0},
-    {
-        "label": np.array(
-            [
-                [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
-                [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
-                [[0, 0, 1], [1, 0, 0], [0, 1, 0]],
-            ]
-        )
-    },
-    [np.array([0, 4, 8]), np.array([1, 5, 6]), np.array([2, 3, 7])],
-]
+    TESTS_CASES.append(
+        [
+            # test One-Hot data
+            {"keys": "label", "image_threshold": 0.0},
+            {
+                "label": p(
+                    [
+                        [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                        [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
+                        [[0, 0, 1], [1, 0, 0], [0, 1, 0]],
+                    ]
+                )
+            },
+            [p([0, 4, 8]), p([1, 5, 6]), p([2, 3, 7])],
+        ]
+    )
 
-TEST_CASE_4 = [
-    {"keys": "label", "image_key": "image", "num_classes": None, "image_threshold": 60},
-    {
-        "label": np.array(
-            [
-                [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
-                [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
-                [[0, 0, 1], [1, 0, 0], [0, 1, 0]],
-            ]
-        ),
-        "image": np.array([[[132, 1434, 51], [61, 0, 133], [523, 44, 232]]]),
-    },
-    [np.array([0, 8]), np.array([1, 5, 6]), np.array([3])],
-]
+    TESTS_CASES.append(
+        [
+            {"keys": "label", "image_key": "image", "num_classes": None, "image_threshold": 60},
+            {
+                "label": p(
+                    [
+                        [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
+                        [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
+                        [[0, 0, 1], [1, 0, 0], [0, 1, 0]],
+                    ]
+                ),
+                "image": p([[[132, 1434, 51], [61, 0, 133], [523, 44, 232]]]),
+            },
+            [p([0, 8]), p([1, 5, 6]), p([3])],
+        ]
+    )
 
-TEST_CASE_5 = [
-    # test output_shape
-    {"keys": "label", "indices_postfix": "cls", "num_classes": 3, "image_threshold": 0.0, "output_shape": [3, 3]},
-    {"label": np.array([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]])},
-    [np.array([[0, 0], [1, 1], [2, 2]]), np.array([[0, 1], [1, 2], [2, 0]]), np.array([[0, 2], [1, 0], [2, 1]])],
-]
+    TESTS_CASES.append(
+        [
+            # test output_shape
+            {
+                "keys": "label",
+                "indices_postfix": "cls",
+                "num_classes": 3,
+                "image_threshold": 0.0,
+                "output_shape": [3, 3],
+            },
+            {"label": p([[[0, 1, 2], [2, 0, 1], [1, 2, 0]]])},
+            [p([[0, 0], [1, 1], [2, 2]]), p([[0, 1], [1, 2], [2, 0]]), p([[0, 2], [1, 0], [2, 1]])],
+        ]
+    )
 
 
 class TestClassesToIndicesd(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
+    @parameterized.expand(TESTS_CASES)
     def test_value(self, input_args, input_data, expected_indices):
         result = ClassesToIndicesd(**input_args)(input_data)
         key_postfix = input_args.get("indices_postfix")
         key_postfix = "_cls_indices" if key_postfix is None else key_postfix
         for i, e in zip(result["label" + key_postfix], expected_indices):
-            np.testing.assert_allclose(i, e)
+            assert_allclose(i, e)
 
 
 if __name__ == "__main__":