convert_ids_to_logits() does not work with single-channel IDs

mipcandy/common/optim/loss.py

@@ -37,29 +37,11 @@ def __init__(self, num_classes: int, include_background: bool) -> None:
         self.include_background: bool = include_background
 
     def logitfy_no_grad(self, ids: torch.Tensor) -> torch.Tensor:
-        with torch.no_grad():
-            if self.num_classes != 1 and ids.shape[1] == 1:
-                if (d := ids.ndim - 2) not in (1, 2, 3):
-                    raise ValueError(f"Expected labels to be 1D, 2D, or 3D, got {d} spatial dimensions")
-                return convert_ids_to_logits(ids.int(), d, self.num_classes)
+        if self.num_classes != 1 and ids.shape[1] == 1:
+            with torch.no_grad():
+                return convert_ids_to_logits(ids.int(), self.num_classes)
         return ids.float()
 
-    def forward(self, outputs: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, dict[str, float]]:
-        if not self.validation_mode:
-            return self._forward(outputs, labels)
-        with torch.no_grad():
-            c, metrics = self._forward(outputs, labels)
-            outputs = convert_logits_to_ids(outputs)
-            dice = 0
-            for i in range(0 if self.include_background else 1, self.num_classes):
-                class_dice = binary_dice(outputs == i, labels == i).item()
-                dice += class_dice
-                metrics[f"dice {i}"] = class_dice
-            metrics["dice"] = dice_similarity_coefficient(
-                self.logitfy_no_grad(outputs), self.logitfy_no_grad(labels)
-            ).item()
-            return c, metrics
-
 
 class DiceCELossWithLogits(_SegmentationLoss):
     def __init__(self, num_classes: int, *, lambda_ce: float = 1, lambda_soft_dice: float = 1,
@@ -81,6 +63,20 @@ def _forward(self, outputs: torch.Tensor, labels: torch.Tensor) -> tuple[torch.T
         c = self.lambda_ce * ce + self.lambda_soft_dice * (1 - dice)
         return c, metrics
 
+    def forward(self, outputs: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, dict[str, float]]:
+        if not self.validation_mode:
+            return self._forward(outputs, labels)
+        with torch.no_grad():
+            c, metrics = self._forward(outputs, labels)
+            outputs = convert_logits_to_ids(outputs)
+            for i in range(0 if self.include_background else 1, self.num_classes):
+                class_dice = binary_dice(outputs == i, labels == i).item()
+                metrics[f"dice {i}"] = class_dice
+            metrics["dice"] = dice_similarity_coefficient(
+                self.logitfy_no_grad(outputs), self.logitfy_no_grad(labels)
+            ).item()
+            return c, metrics
+
 
 class DiceBCELossWithLogits(_SegmentationLoss):
     def __init__(self, *, lambda_bce: float = 1, lambda_soft_dice: float = 1,
@@ -99,3 +95,12 @@ def _forward(self, outputs: torch.Tensor, labels: torch.Tensor) -> tuple[torch.T
         metrics = {"soft dice": dice.item(), "bce loss": bce.item()}
         c = self.lambda_bce * bce + self.lambda_soft_dice * (1 - dice)
         return c, metrics
+
+    def forward(self, outputs: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, dict[str, float]]:
+        if not self.validation_mode:
+            return self._forward(outputs, labels)
+        with torch.no_grad():
+            c, metrics = self._forward(outputs, labels)
+            outputs = convert_logits_to_ids(outputs).bool()
+            metrics["dice"] = binary_dice(outputs, labels.bool()).item()
+            return c, metrics

mipcandy/data/convertion.py

@@ -1,25 +1,28 @@
-from typing import Literal
-
 import torch
 
 from mipcandy.common import Normalize
 
 
-def convert_ids_to_logits(ids: torch.Tensor, d: Literal[1, 2, 3], num_classes: int) -> torch.Tensor:
-    if ids.dtype != torch.int or ids.min() < 0:
-        raise TypeError("`ids` should be positive integers")
-    d += 1
-    if ids.ndim != d:
-        if ids.ndim == d + 1 and ids.shape[1] == 1:
-            ids = ids.squeeze(1)
-        else:
-            raise ValueError(f"`ids` should be {d} dimensional or {d + 1} dimensional with single channel")
-    logits = torch.zeros((ids.shape[0], num_classes, *ids.shape[1:]), device=ids.device, dtype=torch.float32)
-    logits.scatter_(1, ids.unsqueeze(1).long(), 1)
+def convert_ids_to_logits(ids: torch.Tensor, num_classes: int, *, channel_dim: int = 1) -> torch.Tensor:
+    """
+    :param ids: class ids (..., 1, ...)
+    :param num_classes: number of classes
+    :param channel_dim: the index of the channel dimension
+    :return: logits (..., num_classes, ...)
+    """
+    shape = list(ids.shape)
+    shape[channel_dim] = num_classes
+    logits = torch.zeros(shape, device=ids.device, dtype=torch.float32)
+    logits.scatter_(channel_dim, ids.long(), 1)
     return logits
 
 
 def convert_logits_to_ids(logits: torch.Tensor, *, channel_dim: int = 1) -> torch.Tensor:
+    """
+    :param logits: logits (..., num_classes, ...)
+    :param channel_dim: the index of the channel dimension
+    :return: class ids (..., 1, ...)
+    """
     return logits.round().int() if logits.shape[channel_dim] < 2 else logits.argmax(channel_dim, keepdim=True)
 
 

mipcandy/metrics.py

@@ -79,8 +79,4 @@ def soft_dice(outputs: torch.Tensor, labels: torch.Tensor, *, smooth: float = 1,
     label_sum = labels.sum(axes)
     intersection = (outputs * labels).sum(axes)
     output_sum = outputs.sum(axes)
-    if batch_dice:
-        intersection = intersection.sum(0)
-        output_sum = output_sum.sum(0)
-        label_sum = label_sum.sum(0)
     return do_reduction((2 * intersection + smooth) / (label_sum + output_sum + smooth), reduction)

mipcandy/presets/segmentation.py

@@ -158,7 +158,10 @@ def validate_case(self, idx: int, image: torch.Tensor, label: torch.Tensor, tool
         if self.deep_supervision:
             if not isinstance(toolbox.criterion, DeepSupervisionWrapper):
                 raise TypeError("Deep supervision is enabled but criterion is not a `DeepSupervisionWrapper`")
-            output = output[0] if isinstance(output, (list, tuple)) else output[:, 0]
+            if isinstance(output, (list, tuple)):
+                output = output[0]
+            elif output.ndim > label.ndim:
+                output = output[:, 0]
             loss, metrics = toolbox.criterion([output], [label])
         else:
             loss, metrics = toolbox.criterion(output, label)