soft_dice_coefficient incorrectly flattens all classes together (#125)

* compute soft dice per-class instead of flattening all classes

* add do_bg parameter to DiceBCELossWithLogits

* rename do_bg to include_bg

* add include_bg argument in SegmentationTrainer

Author: perctrix

mipcandy/common/optim/loss.py

@@ -26,12 +26,13 @@ def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
 
 class DiceBCELossWithLogits(nn.Module):
     def __init__(self, num_classes: int, *, lambda_bce: float = .5, lambda_soft_dice: float = 1,
-                 smooth: float = 1e-5) -> None:
+                 smooth: float = 1e-5, include_bg: bool = True) -> None:
         super().__init__()
         self.num_classes: int = num_classes
         self.lambda_bce: float = lambda_bce
         self.lambda_soft_dice: float = lambda_soft_dice
         self.smooth: float = smooth
+        self.include_bg: bool = include_bg
 
     def forward(self, masks: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, dict[str, float]]:
         if self.num_classes != 1 and labels.shape[1] == 1:
@@ -42,6 +43,6 @@ def forward(self, masks: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tens
         labels = labels.float()
         bce = nn.functional.binary_cross_entropy_with_logits(masks, labels)
         masks = masks.sigmoid()
-        soft_dice = soft_dice_coefficient(masks, labels, smooth=self.smooth)
+        soft_dice = soft_dice_coefficient(masks, labels, smooth=self.smooth, include_bg=self.include_bg)
         c = self.lambda_bce * bce + self.lambda_soft_dice * (1 - soft_dice)
         return c, {"soft dice": soft_dice.item(), "bce loss": bce.item()}

mipcandy/metrics.py

@@ -60,14 +60,15 @@ def dice_similarity_coefficient_multiclass(output: torch.Tensor, label: torch.Te
     return apply_multiclass_to_binary(dice_similarity_coefficient_binary, output, label, num_classes, if_empty)
 
 
-def soft_dice_coefficient(output: torch.Tensor, label: torch.Tensor, *, smooth: float = 1e-5) -> torch.Tensor:
+def soft_dice_coefficient(output: torch.Tensor, label: torch.Tensor, *,
+                          smooth: float = 1e-5, include_bg: bool = True) -> torch.Tensor:
     _args_check(output, label)
-    num = label.size(0)
-    output = output.view(num, -1)
-    label = label.view(num, -1)
-    intersection = (output * label)
-    dice = (2 * intersection.sum(1) + smooth) / (output.sum(1) + label.sum(1) + smooth)
-    return dice.sum() / num
+    axes = tuple(range(2, output.ndim))
+    intersection = (output * label).sum(dim=axes)
+    dice = (2 * intersection + smooth) / (output.sum(dim=axes) + label.sum(dim=axes) + smooth)
+    if not include_bg:
+        dice = dice[:, 1:]
+    return dice.mean()
 
 
 def accuracy_binary(output: torch.Tensor, label: torch.Tensor, *, if_empty: float = 1) -> torch.Tensor:

mipcandy/presets/segmentation.py

@@ -13,6 +13,7 @@
 
 class SegmentationTrainer(Trainer, metaclass=ABCMeta):
     num_classes: int = 1
+    include_bg: bool = True
 
     def _save_preview(self, x: torch.Tensor, title: str, quality: float) -> None:
         path = f"{self.experiment_folder()}/{title} (preview).png"
@@ -36,7 +37,7 @@ def save_preview(self, image: torch.Tensor, label: torch.Tensor, output: torch.T
 
     @override
     def build_criterion(self) -> nn.Module:
-        return DiceBCELossWithLogits(self.num_classes)
+        return DiceBCELossWithLogits(self.num_classes, include_bg=self.include_bg)
 
     @override
     def build_optimizer(self, params: Params) -> optim.Optimizer: