Merge branch 'dev' into 8620-modulenotfounderror-no-module-named-onnxscript-in-test-py3x-311-pipeline

Signed-off-by: Rafael Garcia-Dias <[email protected]>

Author: garciadias

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/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

monai/networks/schedulers/ddim.py

@@ -117,18 +117,14 @@ def set_timesteps(self, num_inference_steps: int, device: str | torch.device | N
             )
 
         self.num_inference_steps = num_inference_steps
-        step_ratio = self.num_train_timesteps // self.num_inference_steps
-        if self.steps_offset >= step_ratio:
-            raise ValueError(
-                f"`steps_offset`: {self.steps_offset} cannot be greater than or equal to "
-                f"`num_train_timesteps // num_inference_steps : {step_ratio}` as this will cause timesteps to exceed"
-                f" the max train timestep."
-            )
-
-        # creates integer timesteps by multiplying by ratio
-        # casting to int to avoid issues when num_inference_step is power of 3
-        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64)
-        self.timesteps = torch.from_numpy(timesteps).to(device)
+        if self.steps_offset < 0 or self.steps_offset >= self.num_train_timesteps:
+            raise ValueError(f"`steps_offset`: {self.steps_offset} must be in range [0, {self.num_train_timesteps}).")
+
+        self.timesteps = (
+            torch.linspace((self.num_train_timesteps - 1) - self.steps_offset, 0, num_inference_steps, device=device)
+            .round()
+            .long()
+        )
         self.timesteps += self.steps_offset
 
     def _get_variance(self, timestep: int, prev_timestep: int) -> torch.Tensor:

monai/networks/schedulers/ddpm.py

@@ -31,7 +31,6 @@
 
 from __future__ import annotations
 
-import numpy as np
 import torch
 
 from monai.utils import StrEnum
@@ -122,11 +121,9 @@ def set_timesteps(self, num_inference_steps: int, device: str | torch.device | N
             )
 
         self.num_inference_steps = num_inference_steps
-        step_ratio = self.num_train_timesteps // self.num_inference_steps
-        # creates integer timesteps by multiplying by ratio
-        # casting to int to avoid issues when num_inference_step is power of 3
-        timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].astype(np.int64)
-        self.timesteps = torch.from_numpy(timesteps).to(device)
+        self.timesteps = (
+            torch.linspace(self.num_train_timesteps - 1, 0, self.num_inference_steps, device=device).round().long()
+        )
 
     def _get_mean(self, timestep: int, x_0: torch.Tensor, x_t: torch.Tensor) -> torch.Tensor:
         """

monai/transforms/croppad/functional.py

@@ -144,7 +144,7 @@ def crop_or_pad_nd(img: torch.Tensor, translation_mat, spatial_size: tuple[int,
         _mode = _convert_pt_pad_mode(mode)
         img = pad_nd(img, to_pad, mode=_mode, **kwargs)
     if do_crop:
-        img = img[to_crop]
+        img = img[tuple(to_crop)]
     return img
 
 

requirements-dev.txt

@@ -51,7 +51,7 @@ h5py
 nni==2.10.1; platform_system == "Linux" and "arm" not in platform_machine and "aarch" not in platform_machine
 optuna
 git+https://github.com/Project-MONAI/MetricsReloaded@monai-support#egg=MetricsReloaded
-onnx>=1.13.0, <1.19.1
+onnx>=1.13.0
 onnxscript
 onnxruntime; python_version <= '3.10'
 typeguard<3  # https://github.com/microsoft/nni/issues/5457

tests/apps/pathology/transforms/test_pathology_he_stain.py

@@ -48,7 +48,7 @@
 # input pixels not uniformly filled, leading to two different stains extracted
 EXTRACT_STAINS_TEST_CASE_5 = [
     np.array([[[100, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]),
-    np.array([[0.70710677, 0.18696113], [0.0, 0.0], [0.70710677, 0.98236734]]),
+    np.array([[0.18696113, 0.70710677], [0.0, 0.0], [0.98236734, 0.70710677]]),
 ]
 
 # input pixels all transparent and below the beta absorbance threshold
@@ -68,7 +68,7 @@
 NORMALIZE_STAINS_TEST_CASE_4 = [
     {"target_he": np.full((3, 2), 1)},
     np.array([[[100, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]),
-    np.array([[[87, 87, 87], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]]]),
+    np.array([[[31, 31, 31], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]]]),
 ]
 
 
@@ -135,7 +135,7 @@ def test_result_value(self, image, expected_data):
           [[0.18696113],[0],[0.98236734]] and
           [[0.70710677],[0],[0.70710677]] respectively
         - the resulting extracted stain should be
-          [[0.70710677,0.18696113],[0,0],[0.70710677,0.98236734]]
+          [[0.18696113,0.70710677],[0,0],[0.98236734,0.70710677]]
         """
         if image is None:
             with self.assertRaises(TypeError):
@@ -206,17 +206,17 @@ def test_result_value(self, arguments, image, expected_data):
 
         For test case 4:
         - For this non-uniformly filled image, the stain extracted should be
-          [[0.70710677,0.18696113],[0,0],[0.70710677,0.98236734]], as validated for the
+          [[0.18696113,0.70710677],[0,0],[0.98236734,0.70710677]], as validated for the
           ExtractHEStains class. Solving the linear least squares problem (since
           absorbance matrix = stain matrix * concentration matrix), we obtain the concentration
-          matrix that should be [[-0.3101, 7.7508, 7.7508, 7.7508, 7.7508, 7.7508],
-          [5.8022, 0, 0, 0, 0, 0]]
+          matrix that should be [[5.8022, 0, 0, 0, 0, 0],
+          [-0.3101, 7.7508, 7.7508, 7.7508, 7.7508, 7.7508]]
         - Normalizing the concentration matrix, taking the matrix product of the
           target stain matrix and the concentration matrix, using the inverse
           Beer-Lambert transform to obtain the RGB image from the absorbance
           image, and finally converting to uint8, we get that the stain normalized
-          image should be [[[87, 87, 87], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]],
-          [[33, 33, 33], [33, 33, 33]]]
+          image should be [[[31, 31, 31], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]],
+          [[85, 85, 85], [85, 85, 85]]]
         """
         if image is None:
             with self.assertRaises(TypeError):

tests/apps/pathology/transforms/test_pathology_he_stain_dict.py

@@ -42,7 +42,7 @@
 # input pixels not uniformly filled, leading to two different stains extracted
 EXTRACT_STAINS_TEST_CASE_5 = [
     np.array([[[100, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]),
-    np.array([[0.70710677, 0.18696113], [0.0, 0.0], [0.70710677, 0.98236734]]),
+    np.array([[0.18696113, 0.70710677], [0.0, 0.0], [0.98236734, 0.70710677]]),
 ]
 
 # input pixels all transparent and below the beta absorbance threshold
@@ -62,7 +62,7 @@
 NORMALIZE_STAINS_TEST_CASE_4 = [
     {"target_he": np.full((3, 2), 1)},
     np.array([[[100, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]),
-    np.array([[[87, 87, 87], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]]]),
+    np.array([[[31, 31, 31], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]]]),
 ]
 
 
@@ -129,7 +129,7 @@ def test_result_value(self, image, expected_data):
           [[0.18696113],[0],[0.98236734]] and
           [[0.70710677],[0],[0.70710677]] respectively
         - the resulting extracted stain should be
-          [[0.70710677,0.18696113],[0,0],[0.70710677,0.98236734]]
+          [[0.18696113,0.70710677],[0,0],[0.98236734,0.70710677]]
         """
         key = "image"
         if image is None:
@@ -200,17 +200,17 @@ def test_result_value(self, arguments, image, expected_data):
 
         For test case 4:
         - For this non-uniformly filled image, the stain extracted should be
-          [[0.70710677,0.18696113],[0,0],[0.70710677,0.98236734]], as validated for the
+          [[0.18696113,0.70710677],[0,0],[0.98236734,0.70710677]], as validated for the
           ExtractHEStains class. Solving the linear least squares problem (since
           absorbance matrix = stain matrix * concentration matrix), we obtain the concentration
-          matrix that should be [[-0.3101, 7.7508, 7.7508, 7.7508, 7.7508, 7.7508],
-          [5.8022, 0, 0, 0, 0, 0]]
+          matrix that should be  [[5.8022, 0, 0, 0, 0, 0],
+          [-0.3101, 7.7508, 7.7508, 7.7508, 7.7508, 7.7508]]
         - Normalizing the concentration matrix, taking the matrix product of the
           target stain matrix and the concentration matrix, using the inverse
           Beer-Lambert transform to obtain the RGB image from the absorbance
           image, and finally converting to uint8, we get that the stain normalized
-          image should be [[[87, 87, 87], [33, 33, 33]], [[33, 33, 33], [33, 33, 33]],
-          [[33, 33, 33], [33, 33, 33]]]
+          image should be [[[31, 31, 31], [85, 85, 85]], [[85, 85, 85], [85, 85, 85]],
+          [[85, 85, 85], [85, 85, 85]]]
         """
         key = "image"
         if image is None: