style: fix import sorting with isort

Signed-off-by: Rusheel Sharma <rusheelhere@gmail.com>

Author: Rusheel86

monai/losses/dice.py

@@ -23,9 +23,9 @@
 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.metrics.utils import create_ignore_mask
 from monai.networks import one_hot
 from monai.utils import DiceCEReduction, LossReduction, Weight, look_up_option
-from monai.metrics.utils import create_ignore_mask
 
 
 class DiceLoss(_Loss):

monai/losses/focal_loss.py

@@ -18,9 +18,9 @@
 import torch.nn.functional as F
 from torch.nn.modules.loss import _Loss
 
+from monai.metrics.utils import create_ignore_mask
 from monai.networks import one_hot
 from monai.utils import LossReduction
-from monai.metrics.utils import create_ignore_mask
 
 
 class FocalLoss(_Loss):

monai/losses/tversky.py

@@ -18,9 +18,9 @@
 from torch.nn.modules.loss import _Loss
 
 from monai.losses.utils import compute_tp_fp_fn
+from monai.metrics.utils import create_ignore_mask
 from monai.networks import one_hot
 from monai.utils import LossReduction
-from monai.metrics.utils import create_ignore_mask
 
 
 class TverskyLoss(_Loss):

monai/losses/unified_focal_loss.py

@@ -16,9 +16,9 @@
 import torch
 from torch.nn.modules.loss import _Loss
 
+from monai.metrics.utils import create_ignore_mask
 from monai.networks import one_hot
 from monai.utils import LossReduction
-from monai.metrics.utils import create_ignore_mask
 
 
 class AsymmetricFocalTverskyLoss(_Loss):
@@ -78,16 +78,19 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         mask = create_ignore_mask(original_y_true if original_y_true is not None else y_true, self.ignore_index)
         if mask is not None:
             mask = mask.expand_as(y_true)
-            y_pred = y_pred * mask
-            y_true = y_true * mask
 
         y_pred = torch.clamp(y_pred, self.epsilon, 1.0 - self.epsilon)
         axis = list(range(2, len(y_pred.shape)))
 
         # Calculate true positives (tp), false negatives (fn) and false positives (fp)
-        tp = torch.sum(y_true * y_pred, dim=axis)
-        fn = torch.sum(y_true * (1 - y_pred), dim=axis)
-        fp = torch.sum((1 - y_true) * y_pred, dim=axis)
+        if mask is not None:
+            tp = torch.sum(y_true * y_pred * mask, dim=axis)
+            fn = torch.sum(y_true * (1 - y_pred) * mask, dim=axis)
+            fp = torch.sum((1 - y_true) * y_pred * mask, dim=axis)
+        else:
+            tp = torch.sum(y_true * y_pred, dim=axis)
+            fn = torch.sum(y_true * (1 - y_pred), dim=axis)
+            fp = torch.sum((1 - y_true) * y_pred, dim=axis)
         dice_class = (tp + self.epsilon) / (tp + self.delta * fn + (1 - self.delta) * fp + self.epsilon)
 
         # Calculate losses separately for each class, enhancing both classes