torch_em.loss.distance_based

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