Fixed handling of ignored pixels in dice loss

Author: PierreMarza

src/thunder/tasks/train_eval_probe.py

@@ -661,7 +661,7 @@ def train_eval(
                 )
 
                 # Finding background-only masks
-                bg_only = np.array([l.sum().item() == 0 for l in all_label])
+                bg_only = np.array([l.sum().item() == 0 for l in all_label if len(l) > 0])
                 freq_bg_only = bg_only.sum().item() / len(bg_only)
                 no_bg_only_weight = max(
                     1.0, freq_bg_only * cfg.task.no_bg_only_weight_test

src/thunder/utils/dice_loss.py

@@ -4,12 +4,12 @@
 
 # Adapted from https://medium.com/data-scientists-diary/implementation-of-dice-loss-vision-pytorch-7eef1e438f68
 def multiclass_dice_loss(
-    pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor, smooth: float = 1
+    pred: torch.Tensor, label: torch.Tensor, mask: torch.Tensor, smooth: float = 1
 ) -> torch.Tensor:
     """
     Computes Dice Loss for multi-class segmentation.
     :param pred: Tensor of predictions (B, C, H, W).
-    :param target: Ground truth labels (B, H, W).
+    :param label: Ground truth labels (B, H, W).
     :param mask: Mask to apply to pred and target.
     :param smooth: Smoothing factor.
 
@@ -19,6 +19,7 @@ def multiclass_dice_loss(
     pred = F.softmax(pred, dim=1)  # Converting logits to probabilities
     num_classes = pred.shape[1]  # Number of classes (C)
 
+    target = label.clone()
     target[~mask] = (
         num_classes  # Adding a dummy class to account for masked pixels (-1 label values)
     )