torch_em.loss.distance_based

 1import torch
 2import torch.nn as nn
 3
 4from .dice import DiceLoss
 5
 6
 7class DistanceLoss(nn.Module):
 8    """Loss for distance based instance segmentation.
 9
10    Expects input and targets with three channels: foreground and two distance channels.
11    Typically the distance channels are centroid and inverted boundary distance.
12
13    Args:
14        mask_distances_in_bg: whether to mask the loss for distance predictions in the background.
15        foreground_loss: the loss for comparing foreground predictions and target.
16        distance_loss: the loss for comparing distance predictions and target.
17    """
18    def __init__(
19        self,
20        mask_distances_in_bg: bool = True,
21        foreground_loss: nn.Module = DiceLoss(),
22        distance_loss: nn.Module = nn.MSELoss(reduction="mean")
23    ) -> None:
24        super().__init__()
25
26        self.foreground_loss = foreground_loss
27        self.distance_loss = distance_loss
28        self.mask_distances_in_bg = mask_distances_in_bg
29
30        self.init_kwargs = {"mask_distances_in_bg": mask_distances_in_bg}
31
32    def forward(self, input_: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
33        assert input_.shape == target.shape, input_.shape
34        assert input_.shape[1] == 3, input_.shape
35
36        # IMPORTANT: preserve the channels!
37        # Otherwise the Dice Loss will do all kinds of shennanigans.
38        # Because it always interprets the first axis as channel,
39        # and treats it differently (sums over it independently).
40        # This will lead to a very large dice loss that dominates over everything else.
41        fg_input, fg_target = input_[:, 0:1], target[:, 0:1]
42        fg_loss = self.foreground_loss(fg_input, fg_target)
43
44        cdist_input, cdist_target = input_[:, 1:2], target[:, 1:2]
45        if self.mask_distances_in_bg:
46            mask = fg_target
47            cdist_loss = self.distance_loss(cdist_input * mask, cdist_target * mask)
48        else:
49            cdist_loss = self.distance_loss(cdist_input, cdist_target)
50
51        bdist_input, bdist_target = input_[:, 2:3], target[:, 2:3]
52        if self.mask_distances_in_bg:
53            mask = fg_target
54            bdist_loss = self.distance_loss(bdist_input * mask, bdist_target * mask)
55        else:
56            bdist_loss = self.distance_loss(bdist_input, bdist_target)
57
58        overall_loss = fg_loss + cdist_loss + bdist_loss
59        return overall_loss
60
61
62class DiceBasedDistanceLoss(DistanceLoss):
63    """Similar to `DistanceLoss`, using the dice score for all losses.
64
65    Args:
66        mask_distances_in_bg: whether to mask the loss for distance predictions in the background.
67    """
68    def __init__(self, mask_distances_in_bg: bool) -> None:
69        super().__init__(mask_distances_in_bg, foreground_loss=DiceLoss(), distance_loss=DiceLoss())