monai-weekly 1.5.dev2447__py3-none-any.whl → 1.5.dev2449__py3-none-any.whl

monai/__init__.py

@@ -136,4 +136,4 @@ except BaseException:
 
     if MONAIEnvVars.debug():
         raise
-__commit_id__ = "b1e915c323a8065cfe9e92de3013476f2f67c1b2"
+__commit_id__ = "e604d1841fe60c0ffb6978ae4116535ca8d8f34f"

monai/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-11-24T02:30:43+0000",
+ "date": "2024-12-08T02:32:52+0000",
  "dirty": false,
  "error": null,
- "full-revisionid": "2d6751b01bf78fddabe03b2c53645c6bc9808ed8",
- "version": "1.5.dev2447"
+ "full-revisionid": "8cad248c8b374702245989507da1dd8430ef863f",
+ "version": "1.5.dev2449"
 }
 '''  # END VERSION_JSON
 

monai/bundle/__init__.py

@@ -43,4 +43,4 @@ from .utils import (
     MACRO_KEY,
     load_bundle_config,
 )
-from .workflows import BundleWorkflow, ConfigWorkflow
+from .workflows import BundleWorkflow, ConfigWorkflow, PythonicWorkflow

monai/bundle/reference_resolver.py

@@ -192,6 +192,16 @@ class ReferenceResolver:
         """
         return self._resolve_one_item(id=id, **kwargs)
 
+    def remove_resolved_content(self, id: str) -> Any | None:
+        """
+        Remove the resolved ``ConfigItem`` by id.
+
+        Args:
+            id: id name of the expected item.
+
+        """
+        return self.resolved_content.pop(id) if id in self.resolved_content else None
+
     @classmethod
     def normalize_id(cls, id: str | int) -> str:
         """

monai/bundle/workflows.py

@@ -44,12 +44,18 @@ class BundleWorkflow(ABC):
         workflow_type: specifies the workflow type: "train" or "training" for a training workflow,
             or "infer", "inference", "eval", "evaluation" for a inference workflow,
             other unsupported string will raise a ValueError.
-            default to `train` for train workflow.
+            default to `None` for only using meta properties.
         workflow: specifies the workflow type: "train" or "training" for a training workflow,
             or "infer", "inference", "eval", "evaluation" for a inference workflow,
             other unsupported string will raise a ValueError.
             default to `None` for common workflow.
-        properties_path: the path to the JSON file of properties.
+        properties_path: the path to the JSON file of properties. If `workflow_type` is specified, properties will be
+            loaded from the file based on the provided `workflow_type` and meta. If no `workflow_type` is specified,
+            properties will default to loading from "meta". If `properties_path` is None, default properties
+            will be sourced from "monai/bundle/properties.py" based on the workflow_type:
+            For a training workflow, properties load from `TrainProperties` and `MetaProperties`.
+            For a inference workflow, properties load from `InferProperties` and `MetaProperties`.
+            For workflow_type = None : only `MetaProperties` will be loaded.
         meta_file: filepath of the metadata file, if this is a list of file paths, their contents will be merged in order.
         logging_file: config file for `logging` module in the program. for more details:
             https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig.
@@ -97,29 +103,50 @@ class BundleWorkflow(ABC):
                         meta_file = None
 
         workflow_type = workflow if workflow is not None else workflow_type
-        if workflow_type is None and properties_path is None:
-            self.properties = copy(MetaProperties)
-            self.workflow_type = None
-            self.meta_file = meta_file
-            return
+        if workflow_type is not None:
+            if workflow_type.lower() in self.supported_train_type:
+                workflow_type = "train"
+            elif workflow_type.lower() in self.supported_infer_type:
+                workflow_type = "infer"
+            else:
+                raise ValueError(f"Unsupported workflow type: '{workflow_type}'.")
+
         if properties_path is not None:
             properties_path = Path(properties_path)
             if not properties_path.is_file():
                 raise ValueError(f"Property file {properties_path} does not exist.")
             with open(properties_path) as json_file:
-                self.properties = json.load(json_file)
-            self.workflow_type = None
-            self.meta_file = meta_file
-            return
-        if workflow_type.lower() in self.supported_train_type:  # type: ignore[union-attr]
-            self.properties = {**TrainProperties, **MetaProperties}
-            self.workflow_type = "train"
-        elif workflow_type.lower() in self.supported_infer_type:  # type: ignore[union-attr]
-            self.properties = {**InferProperties, **MetaProperties}
-            self.workflow_type = "infer"
+                try:
+                    properties = json.load(json_file)
+                    self.properties: dict = {}
+                    if workflow_type is not None and workflow_type in properties:
+                        self.properties = properties[workflow_type]
+                        if "meta" in properties:
+                            self.properties.update(properties["meta"])
+                    elif workflow_type is None:
+                        if "meta" in properties:
+                            self.properties = properties["meta"]
+                            logger.info(
+                                "No workflow type specified, default to load meta properties from property file."
+                            )
+                        else:
+                            logger.warning("No 'meta' key found in properties while workflow_type is None.")
+                except KeyError as e:
+                    raise ValueError(f"{workflow_type} not found in property file {properties_path}") from e
+                except json.JSONDecodeError as e:
+                    raise ValueError(f"Error decoding JSON from property file {properties_path}") from e
         else:
-            raise ValueError(f"Unsupported workflow type: '{workflow_type}'.")
+            if workflow_type == "train":
+                self.properties = {**TrainProperties, **MetaProperties}
+            elif workflow_type == "infer":
+                self.properties = {**InferProperties, **MetaProperties}
+            elif workflow_type is None:
+                self.properties = copy(MetaProperties)
+                logger.info("No workflow type and property file specified, default to 'meta' properties.")
+            else:
+                raise ValueError(f"Unsupported workflow type: '{workflow_type}'.")
 
+        self.workflow_type = workflow_type
         self.meta_file = meta_file
 
     @abstractmethod
@@ -226,6 +253,124 @@ class BundleWorkflow(ABC):
         return [n for n, p in self.properties.items() if p.get(BundleProperty.REQUIRED, False) and not hasattr(self, n)]
 
 
+class PythonicWorkflow(BundleWorkflow):
+    """
+    Base class for the pythonic workflow specification in bundle, it can be a training, evaluation or inference workflow.
+    It defines the basic interfaces for the bundle workflow behavior: `initialize`, `finalize`, etc.
+    This also provides the interface to get / set public properties to interact with a bundle workflow through
+    defined `get_<property>` accessor methods or directly defining members of the object.
+    For how to set the properties, users can define the `_set_<property>` methods or directly set the members of the object.
+    The `initialize` method is called to set up the workflow before running. This method sets up internal state
+    and prepares properties. If properties are modified after the workflow has been initialized, `self._is_initialized`
+    is set to `False`. Before running the workflow again, `initialize` should be called to ensure that the workflow is
+    properly set up with the new property values.
+
+    Args:
+        workflow_type: specifies the workflow type: "train" or "training" for a training workflow,
+            or "infer", "inference", "eval", "evaluation" for a inference workflow,
+            other unsupported string will raise a ValueError.
+            default to `None` for only using meta properties.
+        workflow: specifies the workflow type: "train" or "training" for a training workflow,
+            or "infer", "inference", "eval", "evaluation" for a inference workflow,
+            other unsupported string will raise a ValueError.
+            default to `None` for common workflow.
+        properties_path: the path to the JSON file of properties. If `workflow_type` is specified, properties will be
+            loaded from the file based on the provided `workflow_type` and meta. If no `workflow_type` is specified,
+            properties will default to loading from "meta". If `properties_path` is None, default properties
+            will be sourced from "monai/bundle/properties.py" based on the workflow_type:
+            For a training workflow, properties load from `TrainProperties` and `MetaProperties`.
+            For a inference workflow, properties load from `InferProperties` and `MetaProperties`.
+            For workflow_type = None : only `MetaProperties` will be loaded.
+        config_file: path to the config file, typically used to store hyperparameters.
+        meta_file: filepath of the metadata file, if this is a list of file paths, their contents will be merged in order.
+        logging_file: config file for `logging` module in the program. for more details:
+            https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig.
+
+    """
+
+    supported_train_type: tuple = ("train", "training")
+    supported_infer_type: tuple = ("infer", "inference", "eval", "evaluation")
+
+    def __init__(
+        self,
+        workflow_type: str | None = None,
+        properties_path: PathLike | None = None,
+        config_file: str | Sequence[str] | None = None,
+        meta_file: str | Sequence[str] | None = None,
+        logging_file: str | None = None,
+        **override: Any,
+    ):
+        meta_file = str(Path(os.getcwd()) / "metadata.json") if meta_file is None else meta_file
+        super().__init__(
+            workflow_type=workflow_type, properties_path=properties_path, meta_file=meta_file, logging_file=logging_file
+        )
+        self._props_vals: dict = {}
+        self._set_props_vals: dict = {}
+        self.parser = ConfigParser()
+        if config_file is not None:
+            self.parser.read_config(f=config_file)
+        if self.meta_file is not None:
+            self.parser.read_meta(f=self.meta_file)
+
+        # the rest key-values in the _args are to override config content
+        self.parser.update(pairs=override)
+        self._is_initialized: bool = False
+
+    def initialize(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Initialize the bundle workflow before running.
+        """
+        self._props_vals = {}
+        self._is_initialized = True
+
+    def _get_property(self, name: str, property: dict) -> Any:
+        """
+        With specified property name and information, get the expected property value.
+        If the property is already generated, return from the bucket directly.
+        If user explicitly set the property, return it directly.
+        Otherwise, generate the expected property as a class private property with prefix "_".
+
+        Args:
+            name: the name of target property.
+            property: other information for the target property, defined in `TrainProperties` or `InferProperties`.
+        """
+        if not self._is_initialized:
+            raise RuntimeError("Please execute 'initialize' before getting any properties.")
+        value = None
+        if name in self._set_props_vals:
+            value = self._set_props_vals[name]
+        elif name in self._props_vals:
+            value = self._props_vals[name]
+        elif name in self.parser.config[self.parser.meta_key]:  # type: ignore[index]
+            id = self.properties.get(name, None).get(BundlePropertyConfig.ID, None)
+            value = self.parser[id]
+        else:
+            try:
+                value = getattr(self, f"get_{name}")()
+            except AttributeError as e:
+                if property[BundleProperty.REQUIRED]:
+                    raise ValueError(
+                        f"unsupported property '{name}' is required in the bundle properties,"
+                        f"need to implement a method 'get_{name}' to provide the property."
+                    ) from e
+            self._props_vals[name] = value
+        return value
+
+    def _set_property(self, name: str, property: dict, value: Any) -> Any:
+        """
+        With specified property name and information, set value for the expected property.
+        Stores user-reset initialized objects that should not be re-initialized and marks the workflow as not initialized.
+
+        Args:
+            name: the name of target property.
+            property: other information for the target property, defined in `TrainProperties` or `InferProperties`.
+            value: value to set for the property.
+
+        """
+        self._set_props_vals[name] = value
+        self._is_initialized = False
+
+
 class ConfigWorkflow(BundleWorkflow):
     """
     Specification for the config-based bundle workflow.
@@ -262,7 +407,13 @@ class ConfigWorkflow(BundleWorkflow):
             or "infer", "inference", "eval", "evaluation" for a inference workflow,
             other unsupported string will raise a ValueError.
             default to `None` for common workflow.
-        properties_path: the path to the JSON file of properties.
+        properties_path: the path to the JSON file of properties. If `workflow_type` is specified, properties will be
+            loaded from the file based on the provided `workflow_type` and meta. If no `workflow_type` is specified,
+            properties will default to loading from "train". If `properties_path` is None, default properties
+            will be sourced from "monai/bundle/properties.py" based on the workflow_type:
+            For a training workflow, properties load from `TrainProperties` and `MetaProperties`.
+            For a inference workflow, properties load from `InferProperties` and `MetaProperties`.
+            For workflow_type = None : only `MetaProperties` will be loaded.
         override: id-value pairs to override or add the corresponding config content.
             e.g. ``--net#input_chns 42``, ``--net %/data/other.json#net_arg``
 
@@ -324,7 +475,6 @@ class ConfigWorkflow(BundleWorkflow):
         self.parser.read_config(f=config_file)
         if self.meta_file is not None:
             self.parser.read_meta(f=self.meta_file)
-
         # the rest key-values in the _args are to override config content
         self.parser.update(pairs=override)
         self.init_id = init_id
@@ -394,8 +544,23 @@ class ConfigWorkflow(BundleWorkflow):
             ret.extend(wrong_props)
         return ret
 
-    def _run_expr(self, id: str, **kwargs: dict) -> Any:
-        return self.parser.get_parsed_content(id, **kwargs) if id in self.parser else None
+    def _run_expr(self, id: str, **kwargs: dict) -> list[Any]:
+        """
+        Evaluate the expression or expression list given by `id`. The resolved values from the evaluations are not stored,
+        allowing this to be evaluated repeatedly (eg. in streaming applications) without restarting the hosting process.
+        """
+        ret = []
+        if id in self.parser:
+            # suppose all the expressions are in a list, run and reset the expressions
+            if isinstance(self.parser[id], list):
+                for i in range(len(self.parser[id])):
+                    sub_id = f"{id}{ID_SEP_KEY}{i}"
+                    ret.append(self.parser.get_parsed_content(sub_id, **kwargs))
+                    self.parser.ref_resolver.remove_resolved_content(sub_id)
+            else:
+                ret.append(self.parser.get_parsed_content(id, **kwargs))
+                self.parser.ref_resolver.remove_resolved_content(id)
+        return ret
 
     def _get_prop_id(self, name: str, property: dict) -> Any:
         prop_id = property[BundlePropertyConfig.ID]

monai/losses/dice.py

@@ -23,6 +23,7 @@ from torch.nn.modules.loss import _Loss
 
 from monai.losses.focal_loss import FocalLoss
 from monai.losses.spatial_mask import MaskedLoss
+from monai.losses.utils import compute_tp_fp_fn
 from monai.networks import one_hot
 from monai.utils import DiceCEReduction, LossReduction, Weight, look_up_option, pytorch_after
 
@@ -39,8 +40,16 @@ class DiceLoss(_Loss):
     The `smooth_nr` and `smooth_dr` parameters are values added to the intersection and union components of
     the inter-over-union calculation to smooth results respectively, these values should be small.
 
-    The original paper: Milletari, F. et. al. (2016) V-Net: Fully Convolutional Neural Networks forVolumetric
-    Medical Image Segmentation, 3DV, 2016.
+    The original papers:
+
+        Milletari, F. et. al. (2016) V-Net: Fully Convolutional Neural Networks for Volumetric
+        Medical Image Segmentation. 3DV 2016.
+
+        Wang, Z. et. al. (2023) Jaccard Metric Losses: Optimizing the Jaccard Index with
+        Soft Labels. NeurIPS 2023.
+
+        Wang, Z. et. al. (2023) Dice Semimetric Losses: Optimizing the Dice Score with
+        Soft Labels. MICCAI 2023.
 
     """
 
@@ -58,6 +67,7 @@ class DiceLoss(_Loss):
         smooth_dr: float = 1e-5,
         batch: bool = False,
         weight: Sequence[float] | float | int | torch.Tensor | None = None,
+        soft_label: bool = False,
     ) -> None:
         """
         Args:
@@ -89,6 +99,8 @@ class DiceLoss(_Loss):
                 of the sequence should be the same as the number of classes. If not ``include_background``,
                 the number of classes should not include the background category class 0).
                 The value/values should be no less than 0. Defaults to None.
+            soft_label: whether the target contains non-binary values (soft labels) or not.
+                If True a soft label formulation of the loss will be used.
 
         Raises:
             TypeError: When ``other_act`` is not an ``Optional[Callable]``.
@@ -114,6 +126,7 @@ class DiceLoss(_Loss):
         weight = torch.as_tensor(weight) if weight is not None else None
         self.register_buffer("class_weight", weight)
         self.class_weight: None | torch.Tensor
+        self.soft_label = soft_label
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
@@ -174,21 +187,15 @@ class DiceLoss(_Loss):
             # reducing spatial dimensions and batch
             reduce_axis = [0] + reduce_axis
 
-        intersection = torch.sum(target * input, dim=reduce_axis)
-
-        if self.squared_pred:
-            ground_o = torch.sum(target**2, dim=reduce_axis)
-            pred_o = torch.sum(input**2, dim=reduce_axis)
-        else:
-            ground_o = torch.sum(target, dim=reduce_axis)
-            pred_o = torch.sum(input, dim=reduce_axis)
-
-        denominator = ground_o + pred_o
-
-        if self.jaccard:
-            denominator = 2.0 * (denominator - intersection)
+        ord = 2 if self.squared_pred else 1
+        tp, fp, fn = compute_tp_fp_fn(input, target, reduce_axis, ord, self.soft_label)
+        if not self.jaccard:
+            fp *= 0.5
+            fn *= 0.5
+        numerator = 2 * tp + self.smooth_nr
+        denominator = 2 * (tp + fp + fn) + self.smooth_dr
 
-        f: torch.Tensor = 1.0 - (2.0 * intersection + self.smooth_nr) / (denominator + self.smooth_dr)
+        f: torch.Tensor = 1 - numerator / denominator
 
         num_of_classes = target.shape[1]
         if self.class_weight is not None and num_of_classes != 1:
@@ -272,6 +279,7 @@ class GeneralizedDiceLoss(_Loss):
         smooth_nr: float = 1e-5,
         smooth_dr: float = 1e-5,
         batch: bool = False,
+        soft_label: bool = False,
     ) -> None:
         """
         Args:
@@ -295,6 +303,8 @@ class GeneralizedDiceLoss(_Loss):
             batch: whether to sum the intersection and union areas over the batch dimension before the dividing.
                 Defaults to False, intersection over union is computed from each item in the batch.
                 If True, the class-weighted intersection and union areas are first summed across the batches.
+            soft_label: whether the target contains non-binary values (soft labels) or not.
+                If True a soft label formulation of the loss will be used.
 
         Raises:
             TypeError: When ``other_act`` is not an ``Optional[Callable]``.
@@ -319,6 +329,7 @@ class GeneralizedDiceLoss(_Loss):
         self.smooth_nr = float(smooth_nr)
         self.smooth_dr = float(smooth_dr)
         self.batch = batch
+        self.soft_label = soft_label
 
     def w_func(self, grnd):
         if self.w_type == str(Weight.SIMPLE):
@@ -370,13 +381,13 @@ class GeneralizedDiceLoss(_Loss):
         reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist()
         if self.batch:
             reduce_axis = [0] + reduce_axis
-        intersection = torch.sum(target * input, reduce_axis)
 
-        ground_o = torch.sum(target, reduce_axis)
-        pred_o = torch.sum(input, reduce_axis)
-
-        denominator = ground_o + pred_o
+        tp, fp, fn = compute_tp_fp_fn(input, target, reduce_axis, 1, self.soft_label)
+        fp *= 0.5
+        fn *= 0.5
+        denominator = 2 * (tp + fp + fn)
 
+        ground_o = torch.sum(target, reduce_axis)
         w = self.w_func(ground_o.float())
         infs = torch.isinf(w)
         if self.batch:
@@ -388,7 +399,7 @@ class GeneralizedDiceLoss(_Loss):
             w = w + infs * max_values
 
         final_reduce_dim = 0 if self.batch else 1
-        numer = 2.0 * (intersection * w).sum(final_reduce_dim, keepdim=True) + self.smooth_nr
+        numer = 2.0 * (tp * w).sum(final_reduce_dim, keepdim=True) + self.smooth_nr
         denom = (denominator * w).sum(final_reduce_dim, keepdim=True) + self.smooth_dr
         f: torch.Tensor = 1.0 - (numer / denom)
 

monai/losses/tversky.py

@@ -17,6 +17,7 @@ from collections.abc import Callable
 import torch
 from torch.nn.modules.loss import _Loss
 
+from monai.losses.utils import compute_tp_fp_fn
 from monai.networks import one_hot
 from monai.utils import LossReduction
 
@@ -28,6 +29,9 @@ class TverskyLoss(_Loss):
         Sadegh et al. (2017) Tversky loss function for image segmentation
         using 3D fully convolutional deep networks. (https://arxiv.org/abs/1706.05721)
 
+        Wang, Z. et. al. (2023) Dice Semimetric Losses: Optimizing the Dice Score with
+        Soft Labels. MICCAI 2023.
+
     Adapted from:
         https://github.com/NifTK/NiftyNet/blob/v0.6.0/niftynet/layer/loss_segmentation.py#L631
 
@@ -46,6 +50,7 @@ class TverskyLoss(_Loss):
         smooth_nr: float = 1e-5,
         smooth_dr: float = 1e-5,
         batch: bool = False,
+        soft_label: bool = False,
     ) -> None:
         """
         Args:
@@ -70,6 +75,8 @@ class TverskyLoss(_Loss):
             batch: whether to sum the intersection and union areas over the batch dimension before the dividing.
                 Defaults to False, a Dice loss value is computed independently from each item in the batch
                 before any `reduction`.
+            soft_label: whether the target contains non-binary values (soft labels) or not.
+                If True a soft label formulation of the loss will be used.
 
         Raises:
             TypeError: When ``other_act`` is not an ``Optional[Callable]``.
@@ -93,6 +100,7 @@ class TverskyLoss(_Loss):
         self.smooth_nr = float(smooth_nr)
         self.smooth_dr = float(smooth_dr)
         self.batch = batch
+        self.soft_label = soft_label
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
@@ -134,20 +142,15 @@ class TverskyLoss(_Loss):
         if target.shape != input.shape:
             raise AssertionError(f"ground truth has differing shape ({target.shape}) from input ({input.shape})")
 
-        p0 = input
-        p1 = 1 - p0
-        g0 = target
-        g1 = 1 - g0
-
         # reducing only spatial dimensions (not batch nor channels)
         reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist()
         if self.batch:
             # reducing spatial dimensions and batch
             reduce_axis = [0] + reduce_axis
 
-        tp = torch.sum(p0 * g0, reduce_axis)
-        fp = self.alpha * torch.sum(p0 * g1, reduce_axis)
-        fn = self.beta * torch.sum(p1 * g0, reduce_axis)
+        tp, fp, fn = compute_tp_fp_fn(input, target, reduce_axis, 1, self.soft_label, False)
+        fp *= self.alpha
+        fn *= self.beta
         numerator = tp + self.smooth_nr
         denominator = tp + fp + fn + self.smooth_dr
 

monai/losses/utils.py

@@ -0,0 +1,68 @@
+# 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
+
+import torch
+import torch.linalg as LA
+
+
+def compute_tp_fp_fn(
+    input: torch.Tensor,
+    target: torch.Tensor,
+    reduce_axis: list[int],
+    ord: int,
+    soft_label: bool,
+    decoupled: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Args:
+        input: the shape should be BNH[WD], where N is the number of classes.
+        target: the shape should be BNH[WD] or B1H[WD], where N is the number of classes.
+        reduce_axis: the axis to be reduced.
+        ord: the order of the vector norm.
+        soft_label: whether the target contains non-binary values (soft labels) or not.
+            If True a soft label formulation of the loss will be used.
+        decoupled: whether the input and the target should be decoupled when computing fp and fn.
+            Only for the original implementation when soft_label is False.
+
+    Adapted from:
+        https://github.com/zifuwanggg/JDTLosses
+    """
+
+    # the original implementation that is erroneous with soft labels
+    if ord == 1 and not soft_label:
+        tp = torch.sum(input * target, dim=reduce_axis)
+        # the original implementation of Dice and Jaccard loss
+        if decoupled:
+            fp = torch.sum(input, dim=reduce_axis) - tp
+            fn = torch.sum(target, dim=reduce_axis) - tp
+        # the original implementation of Tversky loss
+        else:
+            fp = torch.sum(input * (1 - target), dim=reduce_axis)
+            fn = torch.sum((1 - input) * target, dim=reduce_axis)
+    # the new implementation that is correct with soft labels
+    # and it is identical to the original implementation with hard labels
+    else:
+        pred_o = LA.vector_norm(input, ord=ord, dim=reduce_axis)
+        ground_o = LA.vector_norm(target, ord=ord, dim=reduce_axis)
+        difference = LA.vector_norm(input - target, ord=ord, dim=reduce_axis)
+
+        if ord > 1:
+            pred_o = torch.pow(pred_o, exponent=ord)
+            ground_o = torch.pow(ground_o, exponent=ord)
+            difference = torch.pow(difference, exponent=ord)
+
+        tp = (pred_o + ground_o - difference) / 2
+        fp = pred_o - tp
+        fn = ground_o - tp
+
+    return tp, fp, fn

monai/networks/blocks/pos_embed_utils.py

@@ -56,7 +56,7 @@ def build_sincos_position_embedding(
         grid_h = torch.arange(h, dtype=torch.float32)
         grid_w = torch.arange(w, dtype=torch.float32)
 
-        grid_h, grid_w = torch.meshgrid(grid_h, grid_w, indexing="ij")
+        grid_h, grid_w = torch.meshgrid(grid_h, grid_w)
 
         if embed_dim % 4 != 0:
             raise AssertionError("Embed dimension must be divisible by 4 for 2D sin-cos position embedding")
@@ -75,7 +75,7 @@ def build_sincos_position_embedding(
         grid_w = torch.arange(w, dtype=torch.float32)
         grid_d = torch.arange(d, dtype=torch.float32)
 
-        grid_h, grid_w, grid_d = torch.meshgrid(grid_h, grid_w, grid_d, indexing="ij")
+        grid_h, grid_w, grid_d = torch.meshgrid(grid_h, grid_w, grid_d)
 
         if embed_dim % 6 != 0:
             raise AssertionError("Embed dimension must be divisible by 6 for 3D sin-cos position embedding")

monai/networks/blocks/selfattention.py

@@ -11,7 +11,7 @@
 
 from __future__ import annotations
 
-from typing import Tuple, Union
+from typing import Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -154,10 +154,12 @@ class SABlock(nn.Module):
         )
         self.input_size = input_size
 
-    def forward(self, x):
+    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
         """
         Args:
             x (torch.Tensor): input tensor. B x (s_dim_1 * ... * s_dim_n) x C
+            attn_mask (torch.Tensor, optional): mask to apply to the attention matrix.
+            B x (s_dim_1 * ... * s_dim_n). Defaults to None.
 
         Return:
             torch.Tensor: B x (s_dim_1 * ... * s_dim_n) x C
@@ -176,7 +178,13 @@ class SABlock(nn.Module):
 
         if self.use_flash_attention:
             x = F.scaled_dot_product_attention(
-                query=q, key=k, value=v, scale=self.scale, dropout_p=self.dropout_rate, is_causal=self.causal
+                query=q,
+                key=k,
+                value=v,
+                attn_mask=attn_mask,
+                scale=self.scale,
+                dropout_p=self.dropout_rate,
+                is_causal=self.causal,
             )
         else:
             att_mat = torch.einsum("blxd,blyd->blxy", q, k) * self.scale
@@ -186,10 +194,16 @@ class SABlock(nn.Module):
                 att_mat = self.rel_positional_embedding(x, att_mat, q)
 
             if self.causal:
+                if attn_mask is not None:
+                    raise ValueError("Causal attention does not support attention masks.")
                 att_mat = att_mat.masked_fill(self.causal_mask[:, :, : x.shape[-2], : x.shape[-2]] == 0, float("-inf"))
 
-            att_mat = att_mat.softmax(dim=-1)
+            if attn_mask is not None:
+                attn_mask = attn_mask.unsqueeze(1).unsqueeze(2)
+                attn_mask = attn_mask.expand(-1, self.num_heads, -1, -1)
+                att_mat = att_mat.masked_fill(attn_mask == 0, float("-inf"))
 
+            att_mat = att_mat.softmax(dim=-1)
             if self.save_attn:
                 # no gradients and new tensor;
                 # https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html

monai/networks/blocks/transformerblock.py

@@ -90,8 +90,10 @@ class TransformerBlock(nn.Module):
             use_flash_attention=use_flash_attention,
         )
 
-    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None) -> torch.Tensor:
-        x = x + self.attn(self.norm1(x))
+    def forward(
+        self, x: torch.Tensor, context: Optional[torch.Tensor] = None, attn_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), attn_mask=attn_mask)
         if self.with_cross_attention:
             x = x + self.cross_attn(self.norm_cross_attn(x), context=context)
         x = x + self.mlp(self.norm2(x))