Add validation for et_label in ConvertToMultiChannelBasedOnBratsClasses and extend tests

Signed-off-by: Abdessamad <abdoomis6@gmail.com>

Author: AbdoMis

monai/transforms/utility/array.py

@@ -1061,6 +1061,8 @@ class ConvertToMultiChannelBasedOnBratsClasses(Transform):
     backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
 
     def __init__(self, et_label: int = 4) -> None:
+        if et_label in (1, 2):
+            raise ValueError(f"et_label cannot be 1 or 2, these are reserved.Got {et_label}.")
         self.et_label = et_label
 
     def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:

tests/transforms/test_convert_to_multi_channel.py

@@ -20,6 +20,9 @@
 from tests.test_utils import TEST_NDARRAYS, assert_allclose
 
 TESTS = []
+TESTS_ET_LABEL_3 = []
+
+# Tests for default et_label = 4
 for p in TEST_NDARRAYS:
     TESTS.extend(
         [
@@ -46,6 +49,22 @@
         ]
     )
 
+# Tests for et_label = 3
+for p in TEST_NDARRAYS:
+    TESTS_ET_LABEL_3.extend(
+        [
+            [
+                p([[0, 1, 2], [1, 2, 3], [0, 1, 3]]),
+                p(
+                    [
+                        [[0, 1, 0], [1, 0, 1], [0, 1, 1]],
+                        [[0, 1, 1], [1, 1, 1], [0, 1, 1]],
+                        [[0, 0, 0], [0, 0, 1], [0, 0, 1]],
+                    ]
+                ),
+            ],
+        ]
+    )
 
 class TestConvertToMultiChannel(unittest.TestCase):
     @parameterized.expand(TESTS)
@@ -54,6 +73,17 @@ def test_type_shape(self, data, expected_result):
         assert_allclose(result, expected_result)
         self.assertTrue(result.dtype in (bool, torch.bool))
 
+    @parameterized.expand(TESTS_ET_LABEL_3)
+    def test_type_shape_et_label_3(self, data, expected_result):
+        result = ConvertToMultiChannelBasedOnBratsClasses(et_label=3)(data)
+        assert_allclose(result, expected_result)
+        self.assertTrue(result.dtype in (bool, torch.bool))
+
+    def test_invalid_et_label(self):
+        with self.assertRaises(ValueError):
+            ConvertToMultiChannelBasedOnBratsClasses(et_label=1)
+        with self.assertRaises(ValueError):
+            ConvertToMultiChannelBasedOnBratsClasses(et_label=2)
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()

tests/transforms/test_convert_to_multi_channeld.py

@@ -24,6 +24,12 @@
     np.array([[[0, 1, 0], [1, 0, 1], [0, 1, 1]], [[0, 1, 1], [1, 1, 1], [0, 1, 1]], [[0, 0, 0], [0, 0, 1], [0, 0, 1]]]),
 ]
 
+TEST_CASE_ET_LABEL_3 = [
+    {"keys": "label", "et_label": 3},
+    {"label": np.array([[0, 1, 2], [1, 2, 3], [0, 1, 3]])},
+    np.array([[[0, 1, 0], [1, 0, 1], [0, 1, 1]], [[0, 1, 1], [1, 1, 1], [0, 1, 1]], [[0, 0, 0], [0, 0, 1], [0, 0, 1]]]),
+]
+
 
 class TestConvertToMultiChanneld(unittest.TestCase):
 
@@ -32,6 +38,11 @@ def test_type_shape(self, keys, data, expected_result):
         result = ConvertToMultiChannelBasedOnBratsClassesd(**keys)(data)
         np.testing.assert_equal(result["label"], expected_result)
 
+    @parameterized.expand([TEST_CASE_ET_LABEL_3])
+    def test_et_label_3(self, keys, data, expected_result):
+        result = ConvertToMultiChannelBasedOnBratsClassesd(**keys)(data)
+        np.testing.assert_equal(result["label"], expected_result)
+
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()