Merge branch 'dev' of https://github.com/Project-MONAI/MONAI into fix/metatensor-einops

Author: aymuos15

monai/losses/dice.py

@@ -13,7 +13,6 @@
 
 import warnings
 from collections.abc import Callable, Sequence
-from typing import Any
 
 import numpy as np
 import torch
@@ -239,11 +238,52 @@ class MaskedDiceLoss(DiceLoss):
 
     """
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        include_background: bool = True,
+        to_onehot_y: bool = False,
+        sigmoid: bool = False,
+        softmax: bool = False,
+        other_act: Callable | None = None,
+        squared_pred: bool = False,
+        jaccard: bool = False,
+        reduction: LossReduction | str = LossReduction.MEAN,
+        smooth_nr: float = 1e-5,
+        smooth_dr: float = 1e-5,
+        batch: bool = False,
+        weight: Sequence[float] | float | int | torch.Tensor | None = None,
+        soft_label: bool = False,
+    ) -> None:
         """
         Args follow :py:class:`monai.losses.DiceLoss`.
         """
-        super().__init__(*args, **kwargs)
+        if other_act is not None and not callable(other_act):
+            raise TypeError(f"other_act must be None or callable but is {type(other_act).__name__}.")
+        if sigmoid and softmax:
+            raise ValueError("Incompatible values: sigmoid=True and softmax=True.")
+        if other_act is not None and (sigmoid or softmax):
+            raise ValueError("Incompatible values: other_act is not None and sigmoid=True or softmax=True.")
+
+        self.pre_sigmoid = sigmoid
+        self.pre_softmax = softmax
+        self.pre_other_act = other_act
+
+        super().__init__(
+            include_background=include_background,
+            to_onehot_y=to_onehot_y,
+            sigmoid=False,
+            softmax=False,
+            other_act=None,
+            squared_pred=squared_pred,
+            jaccard=jaccard,
+            reduction=reduction,
+            smooth_nr=smooth_nr,
+            smooth_dr=smooth_dr,
+            batch=batch,
+            weight=weight,
+            soft_label=soft_label,
+        )
+
         self.spatial_weighted = MaskedLoss(loss=super().forward)
 
     def forward(self, input: torch.Tensor, target: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
@@ -253,6 +293,19 @@ def forward(self, input: torch.Tensor, target: torch.Tensor, mask: torch.Tensor
             target: the shape should be BNH[WD].
             mask: the shape should B1H[WD] or 11H[WD].
         """
+
+        if self.pre_sigmoid:
+            input = torch.sigmoid(input)
+
+        n_pred_ch = input.shape[1]
+        if self.pre_softmax:
+            if n_pred_ch == 1:
+                warnings.warn("single channel prediction, `softmax=True` ignored.", stacklevel=2)
+            else:
+                input = torch.softmax(input, 1)
+
+        if self.pre_other_act is not None:
+            input = self.pre_other_act(input)
         return self.spatial_weighted(input=input, target=target, mask=mask)  # type: ignore[no-any-return]
 
 

monai/networks/blocks/fft_utils_t.py

@@ -28,9 +28,6 @@ def roll_1d(x: Tensor, shift: int, shift_dim: int) -> Tensor:
 
     Returns:
         1d-shifted version of x
-
-    Note:
-        This function is called when fftshift and ifftshift are not available in the running pytorch version
     """
     shift = shift % x.size(shift_dim)
     if shift == 0:
@@ -55,9 +52,6 @@ def roll(x: Tensor, shift: list[int], shift_dims: list[int]) -> Tensor:
 
     Returns:
         shifted version of x
-
-    Note:
-        This function is called when fftshift and ifftshift are not available in the running pytorch version
     """
     if len(shift) != len(shift_dims):
         raise ValueError(f"len(shift) != len(shift_dims), got f{len(shift)} and f{len(shift_dims)}.")
@@ -78,9 +72,6 @@ def fftshift(x: Tensor, shift_dims: list[int]) -> Tensor:
 
     Returns:
         fft-shifted version of x
-
-    Note:
-        This function is called when fftshift is not available in the running pytorch version
     """
     shift = [0] * len(shift_dims)
     for i, dim_num in enumerate(shift_dims):
@@ -100,9 +91,6 @@ def ifftshift(x: Tensor, shift_dims: list[int]) -> Tensor:
 
     Returns:
         ifft-shifted version of x
-
-    Note:
-        This function is called when ifftshift is not available in the running pytorch version
     """
     shift = [0] * len(shift_dims)
     for i, dim_num in enumerate(shift_dims):

monai/transforms/utils.py

@@ -1878,11 +1878,7 @@ def shift_fourier(x: NdarrayOrTensor, spatial_dims: int, as_contiguous: bool = F
         dims = tuple(range(-spatial_dims, 0))
         k: NdarrayOrTensor
         if isinstance(x, torch.Tensor):
-            if hasattr(torch.fft, "fftshift"):  # `fftshift` is new in torch 1.8.0
-                k = torch.fft.fftshift(torch.fft.fftn(x, dim=dims), dim=dims)
-            else:
-                # if using old PyTorch, will convert to numpy array and return
-                k = np.fft.fftshift(np.fft.fftn(x.cpu().numpy(), axes=dims), axes=dims)
+            k = torch.fft.fftshift(torch.fft.fftn(x, dim=dims), dim=dims)
         else:
             k = np.fft.fftshift(np.fft.fftn(x, axes=dims), axes=dims)
         return ascontiguousarray(k) if as_contiguous else k
@@ -1906,11 +1902,7 @@ def inv_shift_fourier(
         dims = tuple(range(-spatial_dims, 0))
         out: NdarrayOrTensor
         if isinstance(k, torch.Tensor):
-            if hasattr(torch.fft, "ifftshift"):  # `ifftshift` is new in torch 1.8.0
-                out = torch.fft.ifftn(torch.fft.ifftshift(k, dim=dims), dim=dims, norm="backward").real
-            else:
-                # if using old PyTorch, will convert to numpy array and return
-                out = np.fft.ifftn(np.fft.ifftshift(k.cpu().numpy(), axes=dims), axes=dims).real
+            out = torch.fft.ifftn(torch.fft.ifftshift(k, dim=dims), dim=dims, norm="backward").real
         else:
             out = np.fft.ifftn(np.fft.ifftshift(k, axes=dims), axes=dims).real
         return ascontiguousarray(out) if as_contiguous else out

tests/losses/test_masked_dice_loss.py

@@ -27,7 +27,7 @@
             "target": torch.tensor([[[[1.0, 0.0], [1.0, 1.0]]]]),
             "mask": torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]),
         },
-        0.500,
+        0.333333,
     ],
     [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
         {"include_background": True, "sigmoid": True, "smooth_nr": 1e-4, "smooth_dr": 1e-4},
@@ -36,7 +36,7 @@
             "target": torch.tensor([[[[1.0, 1.0], [1.0, 1.0]]], [[[1.0, 0.0], [1.0, 0.0]]]]),
             "mask": torch.tensor([[[[1.0, 1.0], [1.0, 1.0]]], [[[1.0, 1.0], [0.0, 0.0]]]]),
         },
-        0.422969,
+        0.301128,
     ],
     [  # shape: (2, 2, 3), (2, 1, 3)
         {"include_background": False, "to_onehot_y": True, "smooth_nr": 0, "smooth_dr": 0},
@@ -54,7 +54,7 @@
             "target": torch.tensor([[[1.0, 0.0, 0.0]], [[1.0, 1.0, 0.0]]]),
             "mask": torch.tensor([[[1.0, 1.0, 0.0]]]),
         },
-        0.47033,
+        0.579184,
     ],
     [  # shape: (2, 2, 3), (2, 1, 3)
         {