Merge branch 'dev' into dev

Author: ericspod

MANIFEST.in

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

monai/losses/image_dissimilarity.py

@@ -15,7 +15,7 @@
 from torch.nn import functional as F
 from torch.nn.modules.loss import _Loss
 
-from monai.networks.layers import gaussian_1d, separable_filtering
+from monai.networks.layers import separable_filtering
 from monai.utils import LossReduction
 from monai.utils.module import look_up_option
 
@@ -34,11 +34,11 @@ def make_triangular_kernel(kernel_size: int) -> torch.Tensor:
 
 
 def make_gaussian_kernel(kernel_size: int) -> torch.Tensor:
-    sigma = torch.tensor(kernel_size / 3.0)
-    kernel = gaussian_1d(sigma=sigma, truncated=kernel_size // 2, approx="sampled", normalize=False) * (
-        2.5066282 * sigma
-    )
-    return kernel[:kernel_size]
+    sigma = kernel_size / 3.0
+    half = kernel_size // 2
+    x = torch.arange(-half, half + 1, dtype=torch.float)
+    kernel = torch.exp(-0.5 / (sigma * sigma) * x**2)
+    return kernel
 
 
 kernel_dict = {

monai/networks/nets/autoencoderkl.py

@@ -680,6 +680,7 @@ def load_old_state_dict(self, old_state_dict: dict, verbose=False) -> None:
 
         Args:
             old_state_dict: state dict from the old AutoencoderKL model.
+            verbose: if True, print diagnostic information about key mismatches.
         """
 
         new_state_dict = self.state_dict()
@@ -715,13 +716,39 @@ def load_old_state_dict(self, old_state_dict: dict, verbose=False) -> None:
             new_state_dict[f"{block}.attn.to_k.bias"] = old_state_dict.pop(f"{block}.to_k.bias")
             new_state_dict[f"{block}.attn.to_v.bias"] = old_state_dict.pop(f"{block}.to_v.bias")
 
-            # old version did not have a projection so set these to the identity
-            new_state_dict[f"{block}.attn.out_proj.weight"] = torch.eye(
-                new_state_dict[f"{block}.attn.out_proj.weight"].shape[0]
-            )
-            new_state_dict[f"{block}.attn.out_proj.bias"] = torch.zeros(
-                new_state_dict[f"{block}.attn.out_proj.bias"].shape
-            )
+            out_w = f"{block}.attn.out_proj.weight"
+            out_b = f"{block}.attn.out_proj.bias"
+            proj_w = f"{block}.proj_attn.weight"
+            proj_b = f"{block}.proj_attn.bias"
+
+            if out_w in new_state_dict:
+                if proj_w in old_state_dict:
+                    new_state_dict[out_w] = old_state_dict.pop(proj_w)
+                    if proj_b in old_state_dict:
+                        new_state_dict[out_b] = old_state_dict.pop(proj_b)
+                    else:
+                        new_state_dict[out_b] = torch.zeros(
+                            new_state_dict[out_b].shape,
+                            dtype=new_state_dict[out_b].dtype,
+                            device=new_state_dict[out_b].device,
+                        )
+                else:
+                    # No legacy proj_attn - initialize out_proj to identity/zero
+                    new_state_dict[out_w] = torch.eye(
+                        new_state_dict[out_w].shape[0],
+                        dtype=new_state_dict[out_w].dtype,
+                        device=new_state_dict[out_w].device,
+                    )
+                    new_state_dict[out_b] = torch.zeros(
+                        new_state_dict[out_b].shape,
+                        dtype=new_state_dict[out_b].dtype,
+                        device=new_state_dict[out_b].device,
+                    )
+            elif proj_w in old_state_dict:
+                # new model has no out_proj at all - discard the legacy keys so they
+                # don't surface as "unexpected keys" during load_state_dict
+                old_state_dict.pop(proj_w)
+                old_state_dict.pop(proj_b, None)
 
         # fix the upsample conv blocks which were renamed postconv
         for k in new_state_dict:

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(),