Merge branch 'dev' into docs/dints-shape-constraints

Author: Cado87

MANIFEST.in

@@ -3,3 +3,5 @@ include monai/_version.py
 
 include README.md
 include LICENSE
+
+prune tests

monai/networks/nets/swin_unetr.py

@@ -47,6 +47,19 @@ class SwinUNETR(nn.Module):
     Swin UNETR based on: "Hatamizadeh et al.,
     Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images
     <https://arxiv.org/abs/2201.01266>"
+
+    Spatial Shape Constraints:
+        Each spatial dimension of the input must be divisible by ``patch_size ** 5``.
+        With the default ``patch_size=2``, this means each spatial dimension must be divisible by **32**
+        (i.e., 2^5 = 32). This requirement comes from the patch embedding step followed by 4 stages
+        of PatchMerging downsampling, each halving the spatial resolution.
+
+        For a custom ``patch_size``, the divisibility requirement is ``patch_size ** 5``.
+
+        Examples of valid 3D input sizes (with default ``patch_size=2``):
+        ``(32, 32, 32)``, ``(64, 64, 64)``, ``(96, 96, 96)``, ``(128, 128, 128)``, ``(64, 32, 192)``.
+
+        A ``ValueError`` is raised in ``forward()`` if the input spatial shape violates this constraint.
     """
 
     def __init__(
@@ -76,7 +89,8 @@ def __init__(
         Args:
             in_channels: dimension of input channels.
             out_channels: dimension of output channels.
-            patch_size: size of the patch token.
+            patch_size: size of the patch token. Input spatial dimensions must be divisible by
+                ``patch_size ** 5`` (e.g., divisible by 32 when ``patch_size=2``).
             feature_size: dimension of network feature size.
             depths: number of layers in each stage.
             num_heads: number of attention heads.
@@ -108,6 +122,10 @@ def __init__(
             # for 2D single channel input with size (96,96), 2-channel output and gradient checkpointing.
             >>> net = SwinUNETR(in_channels=3, out_channels=2, use_checkpoint=True, spatial_dims=2)
 
+        Raises:
+            ValueError: When a spatial dimension of the input is not divisible by ``patch_size ** 5``.
+                Use ``net._check_input_size(spatial_shape)`` to validate a shape before inference.
+
         """
 
         super().__init__()

monai/transforms/signal/array.py

@@ -273,7 +273,7 @@ def __call__(self, signal: NdarrayOrTensor) -> NdarrayOrTensor:
         data = convert_to_tensor(self.freqs * time_partial)
         sine_partial = self.magnitude * torch.sin(data)
 
-        loc = np.random.choice(range(length))
+        loc = self.R.choice(range(length))
         signal = paste(signal, sine_partial, (loc,))
 
         return signal
@@ -354,7 +354,7 @@ def __call__(self, signal: NdarrayOrTensor) -> NdarrayOrTensor:
         time_partial = np.arange(0, round(self.fracs * length), 1)
         squaredpulse_partial = self.magnitude * squarepulse(self.freqs * time_partial)
 
-        loc = np.random.choice(range(length))
+        loc = self.R.choice(range(length))
         signal = paste(signal, squaredpulse_partial, (loc,))
 
         return signal

monai/transforms/utility/array.py

@@ -1049,19 +1049,34 @@ class ConvertToMultiChannelBasedOnBratsClasses(Transform):
     which include TC (Tumor core), WT (Whole tumor) and ET (Enhancing tumor):
     label 1 is the necrotic and non-enhancing tumor core, which should be counted under TC and WT subregion,
     label 2 is the peritumoral edema, which is counted only under WT subregion,
-    label 4 is the GD-enhancing tumor, which should be counted under ET, TC, WT subregions.
+    the specified `et_label` (default 4) is the GD-enhancing tumor, which should be counted under ET, TC, WT subregions.
+
+    Args:
+        et_label: the label used for the GD-enhancing tumor (ET).
+        - Use 4 for BraTS 2018-2022.
+        - Use 3 for BraTS 2023.
+        Defaults to 4.
     """
 
     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, as these are reserved. Got {et_label}.")
+        self.et_label = et_label
+
     def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         # if img has channel dim, squeeze it
         if img.ndim == 4 and img.shape[0] == 1:
             img = img.squeeze(0)
 
-        result = [(img == 1) | (img == 4), (img == 1) | (img == 4) | (img == 2), img == 4]
-        # merge labels 1 (tumor non-enh) and 4 (tumor enh) and 2 (large edema) to WT
-        # label 4 is ET
+        result = [
+            (img == 1) | (img == self.et_label),
+            (img == 1) | (img == self.et_label) | (img == 2),
+            img == self.et_label,
+        ]
+        # merge labels 1 (tumor non-enh) and self.et_label (tumor enh) and 2 (large edema) to WT
+        # self.et_label is ET (4 or 3)
         return torch.stack(result, dim=0) if isinstance(img, torch.Tensor) else np.stack(result, axis=0)
 
 

monai/transforms/utility/dictionary.py

@@ -1297,19 +1297,27 @@ def __call__(self, data: Mapping[Hashable, Any]):
 class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.ConvertToMultiChannelBasedOnBratsClasses`.
-    Convert labels to multi channels based on brats18 classes:
+    Convert labels to multi channels based on brats classes:
     label 1 is the necrotic and non-enhancing tumor core
     label 2 is the peritumoral edema
-    label 4 is the GD-enhancing tumor
+    the specified `et_label` (default 4) is the GD-enhancing tumor
     The possible classes are TC (Tumor core), WT (Whole tumor)
     and ET (Enhancing tumor).
+
+    Args:
+        keys: keys of the corresponding items to be transformed.
+        et_label: the label used for the GD-enhancing tumor (ET).
+            - Use 4 for BraTS 2018-2022.
+            - Use 3 for BraTS 2023.
+            Defaults to 4.
+        allow_missing_keys: don't raise exception if key is missing.
     """
 
     backend = ConvertToMultiChannelBasedOnBratsClasses.backend
 
-    def __init__(self, keys: KeysCollection, allow_missing_keys: bool = False):
+    def __init__(self, keys: KeysCollection, allow_missing_keys: bool = False, et_label: int = 4):
         super().__init__(keys, allow_missing_keys)
-        self.converter = ConvertToMultiChannelBasedOnBratsClasses()
+        self.converter = ConvertToMultiChannelBasedOnBratsClasses(et_label=et_label)
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
         d = dict(data)

monai/transforms/utils.py

@@ -85,6 +85,9 @@
 cp, has_cp = optional_import("cupy")
 cp_ndarray, _ = optional_import("cupy", name="ndarray")
 exposure, has_skimage = optional_import("skimage.exposure")
+# NOTE: cucim is deliberately NOT imported at module level.
+# Module-level cucim imports caused very slow import times and other buggy behaviour.
+# Keep cucim imports inside the functions that need them.
 
 __all__ = [
     "allow_missing_keys_mode",

monai/utils/module.py

@@ -17,6 +17,7 @@
 import pdb
 import re
 import sys
+import traceback as traceback_mod
 import warnings
 from collections.abc import Callable, Collection, Hashable, Iterable, Mapping
 from functools import partial, wraps
@@ -368,8 +369,9 @@ def optional_import(
         OptionalImportError: from torch.nn.functional import conv1d (requires version '42' by 'min_version').
     """
 
-    tb = None
+    had_exception = False
     exception_str = ""
+    tb_str = ""
     if name:
         actual_cmd = f"from {module} import {name}"
     else:
@@ -384,8 +386,12 @@ def optional_import(
         if name:  # user specified to load class/function/... from the module
             the_module = getattr(the_module, name)
     except Exception as import_exception:  # any exceptions during import
-        tb = import_exception.__traceback__
+        tb_str = "".join(
+            traceback_mod.format_exception(type(import_exception), import_exception, import_exception.__traceback__)
+        )
+        import_exception.__traceback__ = None
         exception_str = f"{import_exception}"
+        had_exception = True
     else:  # found the module
         if version_args and version_checker(pkg, f"{version}", version_args):
             return the_module, True
@@ -394,7 +400,7 @@ def optional_import(
 
     # preparing lazy error message
     msg = descriptor.format(actual_cmd)
-    if version and tb is None:  # a pure version issue
+    if version and not had_exception:  # a pure version issue
         msg += f" (requires '{module} {version}' by '{version_checker.__name__}')"
     if exception_str:
         msg += f" ({exception_str})"
@@ -407,10 +413,9 @@ def __init__(self, *_args, **_kwargs):
                 + "\n\nFor details about installing the optional dependencies, please visit:"
                 + "\n    https://monai.readthedocs.io/en/latest/installation.html#installing-the-recommended-dependencies"
             )
-            if tb is None:
-                self._exception = OptionalImportError(_default_msg)
-            else:
-                self._exception = OptionalImportError(_default_msg).with_traceback(tb)
+            if tb_str:
+                _default_msg += f"\n\nOriginal traceback:\n{tb_str}"
+            self._exception = OptionalImportError(_default_msg)
 
         def __getattr__(self, name):
             """

setup.py

@@ -144,7 +144,7 @@ def get_cmds():
 setup(
     version=versioneer.get_version(),
     cmdclass=get_cmds(),
-    packages=find_packages(exclude=("docs", "examples", "tests")),
+    packages=find_packages(exclude=("docs", "examples", "tests", "tests.*")),
     zip_safe=False,
     package_data={"monai": ["py.typed", *jit_extension_source]},  # type: ignore[arg-type]
     ext_modules=get_extensions(),

tests/networks/nets/test_swin_unetr.py

@@ -90,6 +90,17 @@ def test_ill_arg(self):
         with self.assertRaises(ValueError):
             SwinUNETR(in_channels=1, out_channels=3, feature_size=24, norm_name="instance", drop_rate=-1)
 
+    @skipUnless(has_einops, "Requires einops")
+    def test_invalid_input_shape(self):
+        # spatial dims not divisible by patch_size**5 (default patch_size=2, so must be divisible by 32)
+        net = SwinUNETR(in_channels=1, out_channels=2, feature_size=24, spatial_dims=3)
+        with self.assertRaises(ValueError):
+            net(torch.randn(1, 1, 33, 64, 64))  # 33 is not divisible by 32
+
+        net_2d = SwinUNETR(in_channels=1, out_channels=2, feature_size=24, spatial_dims=2)
+        with self.assertRaises(ValueError):
+            net_2d(torch.randn(1, 1, 48, 33))  # 33 is not divisible by 32
+
     def test_patch_merging(self):
         dim = 10
         t = PatchMerging(dim)(torch.zeros((1, 21, 20, 20, dim)))

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,23 @@
         ]
     )
 
+# 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 +74,18 @@ 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()