torch_em.loss.cldice

  1import torch
  2import torch.nn as nn
  3import torch.nn.functional as F
  4
  5from .dice import dice_score
  6
  7# From "clDice -- A Novel Topology-Preserving Loss Function for Tubular Structure Segmentation":
  8# https://arxiv.org/abs/2003.07311
  9
 10
 11class SoftSkeletonize(torch.nn.Module):
 12    """`SoftSkeletonize` is a differentiable approximation for skeletonization,
 13        which applies iterative min- and max-pooling as a proxy for
 14        morphological erosion and dilation.
 15
 16    Args:
 17        num_iter: Number of iterations for soft-skeletonization.
 18            Should be greater or equal to than the maximum observed radius.
 19    """
 20    def __init__(self, num_iter: int = 5):
 21
 22        super(SoftSkeletonize, self).__init__()
 23        self.num_iter = num_iter
 24
 25    def soft_erode(self, input_: torch.Tensor):
 26
 27        if len(input_.shape) == 4:
 28            p1 = -F.max_pool2d(-input_, (3, 1), (1, 1), (1, 0))
 29            p2 = -F.max_pool2d(-input_, (1, 3), (1, 1), (0, 1))
 30            return torch.min(p1, p2)
 31        elif len(input_.shape) == 5:
 32            p1 = -F.max_pool3d(-input_, (3, 1, 1), (1, 1, 1), (1, 0, 0))
 33            p2 = -F.max_pool3d(-input_, (1, 3, 1), (1, 1, 1), (0, 1, 0))
 34            p3 = -F.max_pool3d(-input_, (1, 1, 3), (1, 1, 1), (0, 0, 1))
 35            return torch.min(torch.min(p1, p2), p3)
 36
 37    def soft_dilate(self, input_: torch.Tensor):
 38
 39        if len(input_.shape) == 4:
 40            return F.max_pool2d(input_, (3, 3), (1, 1), (1, 1))
 41        elif len(input_.shape) == 5:
 42            return F.max_pool3d(input_, (3, 3, 3), (1, 1, 1), (1, 1, 1))
 43
 44    def soft_open(self, input_: torch.Tensor):
 45
 46        return self.soft_dilate(self.soft_erode(input_))
 47
 48    def soft_skel(self, input_: torch.Tensor):
 49
 50        input1 = self.soft_open(input_)
 51        skel = F.relu(input_ - input1)
 52
 53        for j in range(self.num_iter):
 54            input_ = self.soft_erode(input_)
 55            input1 = self.soft_open(input_)
 56            delta = F.relu(input_-input1)
 57            skel = skel + F.relu(delta - skel * delta)
 58
 59        return skel
 60
 61    def forward(self, input_: torch.Tensor):
 62        """Skeletonize the input prediction.
 63
 64        Args:
 65            input_: The input logits.
 66
 67        Returns:
 68            The skeletonization.
 69        """
 70        return self.soft_skel(input_)
 71
 72
 73def cldice_score(
 74    input_: torch.Tensor,
 75    target: torch.Tensor,
 76    num_iter: int = 5,
 77    invert: bool = False,
 78    eps: float = 1e-7,
 79) -> torch.Tensor:
 80    """Adapted from .dice.py `dice_score`. Compute the soft clDice score between input and target.
 81
 82    Args:
 83        input_: The input tensor.
 84        target: The target tensor.
 85        num_iter: Number of iterations for soft-skeletonization.
 86        invert: Whether to invert the returned dice score to obtain the cldice error instead of the cldice score.
 87        channelwise: Not implemented; whether to return the dice score independently per channel.
 88        reduce_channel: Not implemented; how to return the dice score over the channel axis.
 89        eps: The epsilon value added to the denominator for numerical stability.
 90
 91    Returns:
 92        The clDice score.
 93    """
 94    if input_.shape != target.shape:
 95        raise ValueError(f"Expect input and target of same shape, got: {input_.shape}, {target.shape}.")
 96
 97    soft_skeletonize = SoftSkeletonize(num_iter=num_iter)
 98    skel_input = soft_skeletonize(input_)
 99    skel_target = soft_skeletonize(target)
100
101    t_prec = (skel_input * target).sum() / (skel_input.sum()).clamp(min=eps)
102    t_sens = (skel_target * input_).sum() / (skel_target.sum()).clamp(min=eps)
103    score = 2.*(t_prec*t_sens)/(t_prec+t_sens).clamp(min=eps)
104
105    if invert:
106        score = 1. - score
107
108    return score
109
110
111class SoftclDiceLoss(nn.Module):
112    """Combined soft Dice and clDice loss for segmentation of tubular structures.
113
114        The soft clDice loss computes topology-aware loss by computing the
115        soft skeleton of both the prediction and target
116        and measuring overlap of the two skeletons. Teaches the model to learn
117        skeletons directly. In the clDice paper, the authors recommend using
118        the combined soft-Dice and soft-clDice loss to learn topology-aware
119        segmentations, which is implemented below as `CombinedclDiceLoss`.
120
121    Args:
122        num_iter: Number of iterations for soft-skeletonization.
123        eps: The epsilon value added to the denominator for numerical
124            stability.
125        exclude_background: Whether to exclude background channel 0 from the
126            loss computation.
127            Useful for multi-class segmentation.
128        channelwise: Not implemented; Whether to return the dice score
129            independently per channel.
130        reduce_channel: Not implemented; The epsilon value added to the
131            denominator for numerical stability.
132    """
133    def __init__(self, num_iter: int = 5, eps: float = 1e-7,
134                 exclude_background: bool = False):
135        super(SoftclDiceLoss, self).__init__()
136
137        self.num_iter = num_iter
138        self.eps = eps
139        self.exclude_background = exclude_background
140        self.init_kwargs = {"num_iter": num_iter, "eps": eps, "exclude_background": exclude_background}
141
142    def forward(self, input_: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
143        """Compute soft clDice score between the input logits and binary target.
144
145        Args:
146            input_: The input logits.
147            target: The binary target.
148
149        Returns:
150            The soft clDice score.
151        """
152        if input_.shape != target.shape:
153            raise ValueError(f"Expect input and target of same shape, got: {input_.shape}, {target.shape}.")
154
155        if self.exclude_background:
156            target = target[:, 1:, :, :]
157            input_ = input_[:, 1:, :, :]
158
159        cldice = cldice_score(input_, target, num_iter=self.num_iter, invert=True, eps=self.eps)
160
161        return cldice
162
163
164# TODO implement `channelwise` for multiclass segmentation
165# TODO consider if `exclude_background` is needed for multiclass segmentation
166class CombinedclDiceLoss(SoftclDiceLoss):
167    """Combined soft-Dice and soft-clDice loss for segmentation of tubular structures.
168
169        The soft-clDice loss computes topology-aware loss by computing the
170        soft skeleton of both the prediction and target and measuring overlap
171        of the two skeletons. This encourages the model to preserve the
172        connectivity and topology of tubular structures. The final loss is a
173        weighted combination of soft Dice and clDice, controlled by alpha.
174
175    Args:
176        num_iter: Number of iterations for soft-skeletonization.
177        alpha: The weight for combining the soft Dice and soft clDice loss.
178        eps: The epsilon value added to the denominator for numerical
179            stability.
180        exclude_background: Whether to exclude background channel 0 from the
181            loss computation. Useful for multi-class segmentation.
182        invert: Not implemented; Whether to invert the returned dice score to
183            obtain the dice error instead of the dice score.
184        channelwise: Not implemented; Whether to return the dice score
185            independently per channel.
186        reduce_chnanel: Not implemented; How to return the dice score over the
187            channel axis.
188
189    """
190    def __init__(self, num_iter: int = 5, alpha: float = 0.5, eps: float = 1e-7,
191                 exclude_background: bool = False):
192        super(CombinedclDiceLoss, self).__init__(num_iter=num_iter, eps=eps, exclude_background=exclude_background)
193
194        self.alpha = alpha
195        self.init_kwargs = {"num_iter": num_iter, "alpha": alpha, "eps": eps, "exclude_background": exclude_background}
196
197    def forward(self, input_: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
198        """Compute combined soft Dice and clDice loss.
199
200        Args:
201            input_: The input logits.
202            target: The binary target.
203
204        Returns:
205            Combined clDice loss.
206        """
207        if input_.shape != target.shape:
208            raise ValueError(f"Expect input and target of same shape, got: {input_.shape}, {target.shape}.")
209
210        if self.exclude_background:
211            target = target[:, 1:, :, :]
212            input_ = input_[:, 1:, :, :]
213        dice = dice_score(input_, target, invert=True, channelwise=False, eps=self.eps)
214        cldice = cldice_score(input_, target, num_iter=self.num_iter, invert=True, eps=self.eps)
215
216        return (1.0-self.alpha)*dice+self.alpha*cldice