Merge branch 'dev' into feature/global-coords-spatial-crop

Author: ericspod

.github/workflows/pythonapp.yml

@@ -82,7 +82,7 @@ jobs:
         find  /opt/hostedtoolcache/* -maxdepth 0 ! -name 'Python' -exec rm -rf {} \;
     - name: Install the dependencies
       run: |
-        python -m pip install --user --upgrade pip wheel
+        python -m pip install --user --upgrade pip wheel pybind11
         python -m pip install torch==2.5.1 torchvision==0.20.1
         cat "requirements-dev.txt"
         python -m pip install --no-build-isolation -r requirements-dev.txt

MANIFEST.in

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

docs/source/losses.rst

@@ -98,6 +98,11 @@ Segmentation Losses
 .. autoclass:: NACLLoss
     :members:
 
+`MCCLoss`
+~~~~~~~~~
+.. autoclass:: MCCLoss
+    :members:
+
 Registration Losses
 -------------------
 

docs/source/metrics.rst

@@ -158,6 +158,13 @@ Metrics
 
 `Fréchet Inception Distance`
 ------------------------------
+`Embedding Collapse`
+------------------------------
+.. autofunction:: compute_embedding_collapse
+
+.. autoclass:: EmbeddingCollapseMetric
+    :members:
+
 .. autofunction:: compute_frechet_distance
 
 .. autoclass:: FIDMetric

monai/apps/auto3dseg/auto_runner.py

@@ -229,7 +229,7 @@ def __init__(
             input = os.path.join(os.path.abspath(work_dir), "input.yaml")
             logger.info(f"Input config is not provided, using the default {input}")
 
-        self.data_src_cfg = dict()
+        self.data_src_cfg = {}
         if isinstance(input, dict):
             self.data_src_cfg = input
         elif isinstance(input, str) and os.path.isfile(input):

monai/apps/detection/transforms/box_ops.py

@@ -267,7 +267,7 @@ def convert_box_to_mask(
             boxes_only_mask = np.ones(box_size, dtype=np.int16) * np.int16(labels_np[b])
         # apply to global mask
         slicing = [b]
-        slicing.extend(slice(boxes_np[b, d], boxes_np[b, d + spatial_dims]) for d in range(spatial_dims))  # type:ignore
+        slicing.extend(slice(boxes_np[b, d], boxes_np[b, d + spatial_dims]) for d in range(spatial_dims))  # type: ignore
         boxes_mask_np[tuple(slicing)] = boxes_only_mask
     return convert_to_dst_type(src=boxes_mask_np, dst=boxes, dtype=torch.int16)[0]
 

monai/apps/detection/utils/anchor_utils.py

@@ -253,7 +253,7 @@ def grid_anchors(self, grid_sizes: list[list[int]], strides: list[list[Tensor]])
             # compute anchor centers regarding to the image.
             # shifts_centers is [x_center, y_center] or [x_center, y_center, z_center]
             shifts_centers = [
-                torch.arange(0, size[axis], dtype=torch.int32, device=device) * stride[axis]
+                torch.arange(0, size[axis], dtype=torch.int32, device=device) * stride[axis] + stride[axis] // 2
                 for axis in range(self.spatial_dims)
             ]
 

monai/apps/nuclick/transforms.py

@@ -367,14 +367,14 @@ def inclusion_map(self, mask, dtype):
 
     def exclusion_map(self, others, dtype, jitter_range, drop_rate):
         point_mask = torch.zeros_like(others, dtype=dtype)
-        if np.random.choice([True, False], p=[drop_rate, 1 - drop_rate]):
+        if self.R.choice([True, False], p=[drop_rate, 1 - drop_rate]):
             return point_mask
 
         max_x = point_mask.shape[0] - 1
         max_y = point_mask.shape[1] - 1
         stats = measure.regionprops(convert_to_numpy(others))
         for stat in stats:
-            if np.random.choice([True, False], p=[drop_rate, 1 - drop_rate]):
+            if self.R.choice([True, False], p=[drop_rate, 1 - drop_rate]):
                 continue
 
             # random jitter

monai/apps/reconstruction/transforms/dictionary.py

@@ -20,6 +20,7 @@
 from monai.apps.reconstruction.transforms.array import EquispacedKspaceMask, RandomKspaceMask
 from monai.config import DtypeLike, KeysCollection
 from monai.config.type_definitions import NdarrayOrTensor
+from monai.data.meta_tensor import MetaTensor
 from monai.transforms import InvertibleTransform
 from monai.transforms.croppad.array import SpatialCrop
 from monai.transforms.intensity.array import NormalizeIntensity
@@ -33,15 +34,36 @@ class ExtractDataKeyFromMetaKeyd(MapTransform):
     Moves keys from meta to data. It is useful when a dataset of paired samples
     is loaded and certain keys should be moved from meta to data.
 
+    This transform supports two modes:
+
+    1. When ``meta_key`` references a metadata dictionary in the data (e.g., when
+       ``image_only=False`` was used with ``LoadImaged``), the requested keys are
+       extracted directly from that dictionary.
+
+    2. When ``meta_key`` references a ``MetaTensor`` in the data (e.g., when
+       ``image_only=True`` was used with ``LoadImaged``), the requested keys are
+       extracted from its ``.meta`` attribute.
+
     Args:
         keys: keys to be transferred from meta to data
-        meta_key: the meta key where all the meta-data is stored
+        meta_key: the key in the data dictionary where the metadata source is
+            stored. This can be either a metadata dictionary or a ``MetaTensor``.
         allow_missing_keys: don't raise exception if key is missing
 
     Example:
         When the fastMRI dataset is loaded, "kspace" is stored in the data dictionary,
         but the ground-truth image with the key "reconstruction_rss" is stored in the meta data.
         In this case, ExtractDataKeyFromMetaKeyd moves "reconstruction_rss" to data.
+
+        When ``LoadImaged`` is used with ``image_only=True`` (the default), the loaded
+        data is a ``MetaTensor`` with metadata accessible via ``.meta``. In this case,
+        set ``meta_key`` to the key of the ``MetaTensor`` itself::
+
+            li = LoadImaged(keys="image")  # image_only=True by default
+            dat = li({"image": "image.nii"})
+            e = ExtractDataKeyFromMetaKeyd("filename_or_obj", meta_key="image")
+            dat = e(dat)
+            assert dat["image"].meta["filename_or_obj"] == dat["filename_or_obj"]
     """
 
     def __init__(self, keys: KeysCollection, meta_key: str, allow_missing_keys: bool = False) -> None:
@@ -58,9 +80,18 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, T
             the new data dictionary
         """
         d = dict(data)
+        meta_obj = d[self.meta_key]
+
+        # If meta_key references a MetaTensor, extract from its .meta attribute;
+        # otherwise treat it as a metadata dictionary directly.
+        if isinstance(meta_obj, MetaTensor):
+            meta_dict: dict = meta_obj.meta
+        else:
+            meta_dict = dict(meta_obj)
+
         for key in self.keys:
-            if key in d[self.meta_key]:
-                d[key] = d[self.meta_key][key]  # type: ignore
+            if key in meta_dict:
+                d[key] = meta_dict[key]  # type: ignore
             elif not self.allow_missing_keys:
                 raise KeyError(
                     f"Key `{key}` of transform `{self.__class__.__name__}` was missing in the meta data"

monai/auto3dseg/analyzer.py

@@ -15,7 +15,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Hashable, Mapping
 from copy import deepcopy
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import torch
@@ -105,7 +105,7 @@ def update_ops_nested_label(self, nested_key: str, op: Operations) -> None:
             raise ValueError("Nested_key input format is wrong. Please ensure it is like key1#0#key2")
         root: str
         child_key: str
-        (root, _, child_key) = keys
+        root, _, child_key = keys
         if root not in self.ops:
             self.ops[root] = [{}]
         self.ops[root][0].update({child_key: None})
@@ -216,50 +216,58 @@ def __init__(self, image_key: str, stats_name: str = DataStatsKeys.IMAGE_STATS)
         super().__init__(stats_name, report_format)
         self.update_ops(ImageStatsKeys.INTENSITY, SampleOperations())
 
+    @torch.no_grad()
     def __call__(self, data):
-        # Input Validation Addition
-        if not isinstance(data, dict):
-            raise TypeError(f"Input data must be a dict, but got {type(data).__name__}.")
-        if self.image_key not in data:
-            raise KeyError(f"Key '{self.image_key}' not found in input data.")
-        image = data[self.image_key]
-        if not isinstance(image, (np.ndarray, torch.Tensor, MetaTensor)):
-            raise TypeError(
-                f"Value for '{self.image_key}' must be a numpy array, torch.Tensor, or MetaTensor, "
-                f"but got {type(image).__name__}."
-            )
-        if image.ndim < 3:
-            raise ValueError(
-                f"Image data under '{self.image_key}' must have at least 3 dimensions, but got shape {image.shape}."
-            )
-            # --- End of validation ---
         """
-        Callable to execute the pre-defined functions
+        Callable to execute the pre-defined functions.
 
         Returns:
             A dictionary. The dict has the key in self.report_format. The value of
             ImageStatsKeys.INTENSITY is in a list format. Each element of the value list
             has stats pre-defined by SampleOperations (max, min, ....).
 
         Raises:
-            RuntimeError if the stats report generated is not consistent with the pre-
+            KeyError: if ``self.image_key`` is not present in the input data.
+            TypeError: if the input data is not a dictionary, or if the image value is
+                not a numpy array, torch.Tensor, or MetaTensor.
+            ValueError: if the image has fewer than 3 dimensions, or if pre-computed
+                ``nda_croppeds`` is not a list/tuple with one entry per image channel.
+            RuntimeError: if the stats report generated is not consistent with the pre-
                 defined report_format.
 
         Note:
             The stats operation uses numpy and torch to compute max, min, and other
             functions. If the input has nan/inf, the stats results will be nan/inf.
 
         """
+        if not isinstance(data, dict):
+            raise TypeError(f"Input data must be a dict, but got {type(data).__name__}.")
+        if self.image_key not in data:
+            raise KeyError(f"Key '{self.image_key}' not found in input data.")
+        image = data[self.image_key]
+        if not isinstance(image, (np.ndarray, torch.Tensor, MetaTensor)):
+            raise TypeError(
+                f"Value for '{self.image_key}' must be a numpy array, torch.Tensor, or MetaTensor, "
+                f"but got {type(image).__name__}."
+            )
+        if image.ndim < 3:
+            raise ValueError(
+                f"Image data under '{self.image_key}' must have at least 3 dimensions, but got shape {image.shape}."
+            )
+
         d = dict(data)
         start = time.time()
-        restore_grad_state = torch.is_grad_enabled()
-        torch.set_grad_enabled(False)
-
         ndas = [d[self.image_key][i] for i in range(d[self.image_key].shape[0])]
-        if "nda_croppeds" not in d:
+        if "nda_croppeds" in d:
+            nda_croppeds = d["nda_croppeds"]
+            if not isinstance(nda_croppeds, (list, tuple)) or len(nda_croppeds) != len(ndas):
+                raise ValueError(
+                    "Pre-computed 'nda_croppeds' must be a list or tuple with one entry per image channel "
+                    f"(expected {len(ndas)})."
+                )
+        else:
             nda_croppeds = [get_foreground_image(nda) for nda in ndas]
 
-        # perform calculation
         report = deepcopy(self.get_report_format())
 
         report[ImageStatsKeys.SHAPE] = [list(nda.shape) for nda in ndas]
@@ -284,7 +292,6 @@ def __call__(self, data):
 
         d[self.stats_name] = report
 
-        torch.set_grad_enabled(restore_grad_state)
         logger.debug(f"Get image stats spent {time.time() - start}")
         return d
 
@@ -321,6 +328,7 @@ def __init__(self, image_key: str, label_key: str, stats_name: str = DataStatsKe
         super().__init__(stats_name, report_format)
         self.update_ops(ImageStatsKeys.INTENSITY, SampleOperations())
 
+    @torch.no_grad()
     def __call__(self, data: Mapping) -> dict:
         """
         Callable to execute the pre-defined functions
@@ -341,9 +349,6 @@ def __call__(self, data: Mapping) -> dict:
 
         d = dict(data)
         start = time.time()
-        restore_grad_state = torch.is_grad_enabled()
-        torch.set_grad_enabled(False)
-
         ndas = [d[self.image_key][i] for i in range(d[self.image_key].shape[0])]
         ndas_label = d[self.label_key]  # (H,W,D)
 
@@ -353,7 +358,6 @@ def __call__(self, data: Mapping) -> dict:
         nda_foregrounds = [get_foreground_label(nda, ndas_label) for nda in ndas]
         nda_foregrounds = [nda if nda.numel() > 0 else MetaTensor([0.0]) for nda in nda_foregrounds]
 
-        # perform calculation
         report = deepcopy(self.get_report_format())
 
         report[ImageStatsKeys.INTENSITY] = [
@@ -365,7 +369,6 @@ def __call__(self, data: Mapping) -> dict:
 
         d[self.stats_name] = report
 
-        torch.set_grad_enabled(restore_grad_state)
         logger.debug(f"Get foreground image stats spent {time.time() - start}")
         return d
 
@@ -418,6 +421,7 @@ def __init__(
         id_seq = ID_SEP_KEY.join([LabelStatsKeys.LABEL, "0", LabelStatsKeys.IMAGE_INTST])
         self.update_ops_nested_label(id_seq, SampleOperations())
 
+    @torch.no_grad()
     def __call__(self, data: Mapping[Hashable, MetaTensor]) -> dict[Hashable, MetaTensor | dict]:
         """
         Callable to execute the pre-defined functions.
@@ -468,21 +472,31 @@ def __call__(self, data: Mapping[Hashable, MetaTensor]) -> dict[Hashable, MetaTe
         """
         d: dict[Hashable, MetaTensor] = dict(data)
         start = time.time()
-        if isinstance(d[self.image_key], (torch.Tensor, MetaTensor)) and d[self.image_key].device.type == "cuda":
-            using_cuda = True
-        else:
-            using_cuda = False
-        restore_grad_state = torch.is_grad_enabled()
-        torch.set_grad_enabled(False)
+        image_tensor = d[self.image_key]
+        label_tensor = d[self.label_key]
+        using_cuda = any(
+            isinstance(t, (torch.Tensor, MetaTensor)) and t.device.type == "cuda" for t in (image_tensor, label_tensor)
+        )
 
-        ndas: list[MetaTensor] = [d[self.image_key][i] for i in range(d[self.image_key].shape[0])]  # type: ignore
-        ndas_label: MetaTensor = d[self.label_key].astype(torch.int16)  # (H,W,D)
+        if isinstance(image_tensor, (MetaTensor, torch.Tensor)) and isinstance(
+            label_tensor, (MetaTensor, torch.Tensor)
+        ):
+            if label_tensor.device != image_tensor.device:
+                if using_cuda:
+                    cuda_device = image_tensor.device if image_tensor.device.type == "cuda" else label_tensor.device
+                    image_tensor = cast(MetaTensor, image_tensor.to(cuda_device))
+                    label_tensor = cast(MetaTensor, label_tensor.to(cuda_device))
+                else:
+                    label_tensor = cast(MetaTensor, label_tensor.to(image_tensor.device))
+
+        ndas: list[MetaTensor] = [image_tensor[i] for i in range(image_tensor.shape[0])]  # type: ignore
+        ndas_label: MetaTensor = label_tensor.astype(torch.int16)  # (H,W,D)
 
         if ndas_label.shape != ndas[0].shape:
             raise ValueError(f"Label shape {ndas_label.shape} is different from image shape {ndas[0].shape}")
 
         nda_foregrounds: list[torch.Tensor] = [get_foreground_label(nda, ndas_label) for nda in ndas]
-        nda_foregrounds = [nda if nda.numel() > 0 else torch.Tensor([0]) for nda in nda_foregrounds]
+        nda_foregrounds = [nda if nda.numel() > 0 else MetaTensor([0.0]) for nda in nda_foregrounds]
 
         unique_label = unique(ndas_label)
         if isinstance(ndas_label, (MetaTensor, torch.Tensor)):
@@ -534,7 +548,6 @@ def __call__(self, data: Mapping[Hashable, MetaTensor]) -> dict[Hashable, MetaTe
 
         d[self.stats_name] = report  # type: ignore[assignment]
 
-        torch.set_grad_enabled(restore_grad_state)
         logger.debug(f"Get label stats spent {time.time() - start}")
         return d  # type: ignore[return-value]