Merge branch 'dev' into feature/add-non-central-chi-noise

Author: ericspod

.github/workflows/codeql-analysis.yml

@@ -42,7 +42,7 @@ jobs:
 
     # Initializes the CodeQL tools for scanning.
     - name: Initialize CodeQL
-      uses: github/codeql-action/init@v3
+      uses: github/codeql-action/init@v4
       with:
         languages: ${{ matrix.language }}
         # If you wish to specify custom queries, you can do so here or in a config file.
@@ -72,4 +72,4 @@ jobs:
         BUILD_MONAI=1 ./runtests.sh --build
 
     - name: Perform CodeQL Analysis
-      uses: github/codeql-action/analyze@v3
+      uses: github/codeql-action/analyze@v4

.github/workflows/docker.yml

@@ -37,7 +37,7 @@ jobs:
           python setup.py build
           cat build/lib/monai/_version.py
       - name: Upload version
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v5
         with:
           name: _version.py
           path: build/lib/monai/_version.py

.github/workflows/release.yml

@@ -66,7 +66,7 @@ jobs:
 
     - if: matrix.python-version == '3.9' && startsWith(github.ref, 'refs/tags/')
       name: Upload artifacts
-      uses: actions/upload-artifact@v4
+      uses: actions/upload-artifact@v5
       with:
         name: dist
         path: dist/
@@ -109,7 +109,7 @@ jobs:
           python setup.py build
           cat build/lib/monai/_version.py
       - name: Upload version
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v5
         with:
           name: _version.py
           path: build/lib/monai/_version.py

README.md

@@ -1,5 +1,5 @@
 <p align="center">
-  <img src="https://raw.githubusercontent.com/Project-MONAI/MONAI/dev/docs/images/MONAI-logo-color.png" width="50%" alt='project-monai'>
+<img src="https://raw.githubusercontent.com/Project-MONAI/MONAI/dev/docs/images/MONAI-logo-color.png" width="50%" alt='project-monai'>
 </p>
 
 **M**edical **O**pen **N**etwork for **AI**

docs/source/transforms.rst

@@ -37,6 +37,11 @@ Generic Interfaces
 .. autoclass:: MultiSampleTrait
     :members:
 
+`ReduceTrait`
+^^^^^^^^^^^^^^^^^^
+.. autoclass:: ReduceTrait
+    :members:
+
 `Randomizable`
 ^^^^^^^^^^^^^^
 .. autoclass:: Randomizable
@@ -1252,6 +1257,12 @@ Utility
     :members:
     :special-members: __call__
 
+`FlattenSequence`
+""""""""""""""""""""""""
+.. autoclass:: FlattenSequence
+    :members:
+    :special-members: __call__
+
 Dictionary Transforms
 ---------------------
 
@@ -2337,6 +2348,12 @@ Utility (Dict)
     :members:
     :special-members: __call__
 
+`FlattenSequenced`
+"""""""""""""""""""""""""
+.. autoclass:: FlattenSequenced
+    :members:
+    :special-members: __call__
+
 
 MetaTensor
 ^^^^^^^^^^

monai/apps/pathology/transforms/stain/array.py

@@ -85,7 +85,12 @@ def _deconvolution_extract_stain(self, image: np.ndarray) -> np.ndarray:
         v_max = eigvecs[:, 1:3].dot(np.array([(np.cos(max_phi), np.sin(max_phi))], dtype=np.float32).T)
 
         # a heuristic to make the vector corresponding to hematoxylin first and the one corresponding to eosin second
-        if v_min[0] > v_max[0]:
+        # Hematoxylin: high blue, lower red (low R/B ratio)
+        # Eosin: high red, lower blue (high R/B ratio)
+        eps = np.finfo(np.float32).eps
+        v_min_rb_ratio = v_min[0] / (v_min[2] + eps)
+        v_max_rb_ratio = v_max[0] / (v_max[2] + eps)
+        if v_min_rb_ratio < v_max_rb_ratio:
             he = np.array((v_min[:, 0], v_max[:, 0]), dtype=np.float32).T
         else:
             he = np.array((v_max[:, 0], v_min[:, 0]), dtype=np.float32).T

monai/data/box_utils.py

@@ -826,7 +826,10 @@ def box_iou(boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor) -> NdarrayOrTensor
         boxes2: bounding boxes, Mx4 or Mx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``
 
     Returns:
-        IoU, with size of (N,M) and same data type as ``boxes1``
+        An array/tensor matching the container type of ``boxes1`` (NumPy ndarray or Torch tensor), always
+        floating-point with size ``(N, M)``:
+        - if ``boxes1`` has a floating-point dtype, the same dtype is used.
+        - if ``boxes1`` has an integer dtype, the result is returned as ``torch.float32``.
 
     """
 
@@ -842,16 +845,18 @@ def box_iou(boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor) -> NdarrayOrTensor
 
     inter, union = _box_inter_union(boxes1_t, boxes2_t, compute_dtype=COMPUTE_DTYPE)
 
-    # compute IoU and convert back to original box_dtype
+    # compute IoU and convert back to original box_dtype or torch.float32
     iou_t = inter / (union + torch.finfo(COMPUTE_DTYPE).eps)  # (N,M)
+    if not box_dtype.is_floating_point:
+        box_dtype = COMPUTE_DTYPE
     iou_t = iou_t.to(dtype=box_dtype)
 
     # check if NaN or Inf
     if torch.isnan(iou_t).any() or torch.isinf(iou_t).any():
         raise ValueError("Box IoU is NaN or Inf.")
 
     # convert tensor back to numpy if needed
-    iou, *_ = convert_to_dst_type(src=iou_t, dst=boxes1)
+    iou, *_ = convert_to_dst_type(src=iou_t, dst=boxes1, dtype=box_dtype)
     return iou
 
 
@@ -867,7 +872,10 @@ def box_giou(boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor) -> NdarrayOrTenso
         boxes2: bounding boxes, Mx4 or Mx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``
 
     Returns:
-        GIoU, with size of (N,M) and same data type as ``boxes1``
+        An array/tensor matching the container type of ``boxes1`` (NumPy ndarray or Torch tensor), always
+        floating-point with size ``(N, M)``:
+        - if ``boxes1`` has a floating-point dtype, the same dtype is used.
+        - if ``boxes1`` has an integer dtype, the result is returned as ``torch.float32``.
 
     Reference:
         https://giou.stanford.edu/GIoU.pdf
@@ -904,12 +912,15 @@ def box_giou(boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor) -> NdarrayOrTenso
 
     # GIoU
     giou_t = iou - (enclosure - union) / (enclosure + torch.finfo(COMPUTE_DTYPE).eps)
+    if not box_dtype.is_floating_point:
+        box_dtype = COMPUTE_DTYPE
     giou_t = giou_t.to(dtype=box_dtype)
+
     if torch.isnan(giou_t).any() or torch.isinf(giou_t).any():
         raise ValueError("Box GIoU is NaN or Inf.")
 
     # convert tensor back to numpy if needed
-    giou, *_ = convert_to_dst_type(src=giou_t, dst=boxes1)
+    giou, *_ = convert_to_dst_type(src=giou_t, dst=boxes1, dtype=box_dtype)
     return giou
 
 
@@ -925,7 +936,10 @@ def box_pair_giou(boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor) -> NdarrayOr
         boxes2: bounding boxes, same shape with boxes1. The box mode is assumed to be ``StandardMode``
 
     Returns:
-        paired GIoU, with size of (N,) and same data type as ``boxes1``
+        An array/tensor matching the container type of ``boxes1`` (NumPy ndarray or Torch tensor), always
+        floating-point with size ``(N, )``:
+        - if ``boxes1`` has a floating-point dtype, the same dtype is used.
+        - if ``boxes1`` has an integer dtype, the result is returned as ``torch.float32``.
 
     Reference:
         https://giou.stanford.edu/GIoU.pdf
@@ -982,12 +996,15 @@ def box_pair_giou(boxes1: NdarrayOrTensor, boxes2: NdarrayOrTensor) -> NdarrayOr
     enclosure = torch.prod(wh, dim=-1, keepdim=False)  # (N,)
 
     giou_t: torch.Tensor = iou - (enclosure - union) / (enclosure + torch.finfo(COMPUTE_DTYPE).eps)  # type: ignore
+    if not box_dtype.is_floating_point:
+        box_dtype = COMPUTE_DTYPE
     giou_t = giou_t.to(dtype=box_dtype)  # (N,spatial_dims)
+
     if torch.isnan(giou_t).any() or torch.isinf(giou_t).any():
         raise ValueError("Box GIoU is NaN or Inf.")
 
     # convert tensor back to numpy if needed
-    giou, *_ = convert_to_dst_type(src=giou_t, dst=boxes1)
+    giou, *_ = convert_to_dst_type(src=giou_t, dst=boxes1, dtype=box_dtype)
     return giou
 
 

monai/data/dataset.py

@@ -230,6 +230,8 @@ def __init__(
         pickle_protocol: int = DEFAULT_PROTOCOL,
         hash_transform: Callable[..., bytes] | None = None,
         reset_ops_id: bool = True,
+        track_meta: bool = False,
+        weights_only: bool = True,
     ) -> None:
         """
         Args:
@@ -264,7 +266,17 @@ def __init__(
                 When this is enabled, the traced transform instance IDs will be removed from the cached MetaTensors.
                 This is useful for skipping the transform instance checks when inverting applied operations
                 using the cached content and with re-created transform instances.
-
+            track_meta: whether to track the meta information, if `True`, will convert to `MetaTensor`.
+                default to `False`. Cannot be used with `weights_only=True`.
+            weights_only: keyword argument passed to `torch.load` when reading cached files.
+                default to `True`. When set to `True`, `torch.load` restricts loading to tensors and
+                other safe objects. Setting this to `False` is required for loading `MetaTensor`
+                objects saved with `track_meta=True`, however this creates the possibility of remote
+                code execution through `torch.load` so be aware of the security implications of doing so.
+
+        Raises:
+            ValueError: When both `track_meta=True` and `weights_only=True`, since this combination
+                prevents cached MetaTensors from being reloaded and causes perpetual cache regeneration.
         """
         super().__init__(data=data, transform=transform)
         self.cache_dir = Path(cache_dir) if cache_dir is not None else None
@@ -280,6 +292,13 @@ def __init__(
         if hash_transform is not None:
             self.set_transform_hash(hash_transform)
         self.reset_ops_id = reset_ops_id
+        if track_meta and weights_only:
+            raise ValueError(
+                "Invalid argument combination: `track_meta=True` cannot be used with `weights_only=True`. "
+                "To cache and reload MetaTensors, set `track_meta=True` and `weights_only=False`."
+            )
+        self.track_meta = track_meta
+        self.weights_only = weights_only
 
     def set_transform_hash(self, hash_xform_func: Callable[..., bytes]):
         """Get hashable transforms, and then hash them. Hashable transforms
@@ -377,7 +396,7 @@ def _cachecheck(self, item_transformed):
 
         if hashfile is not None and hashfile.is_file():  # cache hit
             try:
-                return torch.load(hashfile, weights_only=True)
+                return torch.load(hashfile, weights_only=self.weights_only)
             except PermissionError as e:
                 if sys.platform != "win32":
                     raise e
@@ -398,7 +417,7 @@ def _cachecheck(self, item_transformed):
             with tempfile.TemporaryDirectory() as tmpdirname:
                 temp_hash_file = Path(tmpdirname) / hashfile.name
                 torch.save(
-                    obj=convert_to_tensor(_item_transformed, convert_numeric=False),
+                    obj=convert_to_tensor(_item_transformed, convert_numeric=False, track_meta=self.track_meta),
                     f=temp_hash_file,
                     pickle_module=look_up_option(self.pickle_module, SUPPORTED_PICKLE_MOD),
                     pickle_protocol=self.pickle_protocol,

monai/networks/blocks/activation_checkpointing.py

@@ -0,0 +1,41 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import cast
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+
+
+class ActivationCheckpointWrapper(nn.Module):
+    """Wrapper applying activation checkpointing to a module during training.
+
+    Args:
+        module: The module to wrap with activation checkpointing.
+    """
+
+    def __init__(self, module: nn.Module) -> None:
+        super().__init__()
+        self.module = module
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass with optional activation checkpointing.
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            Output tensor from the wrapped module.
+        """
+        return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))

monai/networks/nets/unet.py

@@ -17,11 +17,12 @@
 import torch
 import torch.nn as nn
 
+from monai.networks.blocks.activation_checkpointing import ActivationCheckpointWrapper
 from monai.networks.blocks.convolutions import Convolution, ResidualUnit
 from monai.networks.layers.factories import Act, Norm
 from monai.networks.layers.simplelayers import SkipConnection
 
-__all__ = ["UNet", "Unet"]
+__all__ = ["UNet", "Unet", "CheckpointUNet"]
 
 
 class UNet(nn.Module):
@@ -298,4 +299,29 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
+class CheckpointUNet(UNet):
+    """UNet variant that wraps internal connection blocks with activation checkpointing.
+
+    See `UNet` for constructor arguments. During training with gradients enabled,
+    intermediate activations inside encoder-decoder connections are recomputed in
+    the backward pass to reduce peak memory usage at the cost of extra compute.
+    """
+
+    def _get_connection_block(self, down_path: nn.Module, up_path: nn.Module, subblock: nn.Module) -> nn.Module:
+        """Returns connection block with activation checkpointing applied to all components.
+
+        Args:
+            down_path: encoding half of the layer (will be wrapped with checkpointing).
+            up_path: decoding half of the layer (will be wrapped with checkpointing).
+            subblock: block defining the next layer (will be wrapped with checkpointing).
+
+        Returns:
+            Connection block with all components wrapped for activation checkpointing.
+        """
+        subblock = ActivationCheckpointWrapper(subblock)
+        down_path = ActivationCheckpointWrapper(down_path)
+        up_path = ActivationCheckpointWrapper(up_path)
+        return super()._get_connection_block(down_path, up_path, subblock)
+
+
 Unet = UNet