''' Dice loss Adapted from MedicalZooPytorch: https://github.com/black0017/MedicalZooPytorch ''' import torch from torch import nn as nn from typing import Optional # Code was adapted and mofified from https://github.com/wolny/pytorch-3dunet/blob/master/pytorch3dunet/unet3d/losses.py class _AbstractDiceLoss(nn.Module): """ Base class for different implementations of Dice loss. """ def __init__(self, weight=None, sigmoid_normalization=True, classes=2): super(_AbstractDiceLoss, self).__init__() self.register_buffer('weight', weight) self.classes = None self.skip_index_after = None # The output from the network during training is assumed to be un-normalized probabilities and we would # like to normalize the logits. Since Dice (or soft Dice in this case) is usually used for binary data, # normalizing the channels with Sigmoid is the default choice even for multi-class segmentation problems. # However if one would like to apply Softmax in order to get the proper probability distribution from the # output, just specify sigmoid_normalization=False. if sigmoid_normalization: self.normalization = nn.Sigmoid() else: self.normalization = nn.Softmax(dim=1) def dice(self, input, target, weight): # actual Dice score computation; to be implemented by the subclass raise NotImplementedError def skip_target_channels(self, target, index): """ Assuming dim 1 is the classes dim , it skips all the indexes after the desired class """ assert index >= 2 return target[:, 0:index, ...] def forward(self, input, target): """ Expand to one hot added extra for consistency reasons """ # target = expand_as_one_hot(target.long(), self.classes) shape = list(target.long().size()) shape[1] = self.classes # One hot encoder target = torch.zeros(shape).to(target.long()).scatter_(1, target.long(), 1) # assert input.dim() == target.dim() == 5, "'input' and 'target' have different number of dims" if self.skip_index_after is not None: before_size = target.size() target = self.skip_target_channels(target, self.skip_index_after) print("Target {} after skip index {}".format(before_size, target.size())) assert input.size() == target.size(), "'input' and 'target' must have the same shape" # get probabilities from logits input = self.normalization(input) # compute per channel Dice coefficient per_channel_dice = self.dice(input, target, weight=self.weight) loss = (1. - torch.mean(per_channel_dice)) per_channel_dice = per_channel_dice.detach().cpu().numpy() # average Dice score across all channels/classes return loss, per_channel_dice # Code was adapted and mofified from https://github.com/wolny/pytorch-3dunet/blob/master/pytorch3dunet/unet3d/losses.py class DiceLoss(_AbstractDiceLoss): """Computes Dice Loss according to https://arxiv.org/abs/1606.04797. For multi-class segmentation `weight` parameter can be used to assign different weights per class. """ def __init__(self, classes=2, skip_index_after=None, weight=None, sigmoid_normalization=True ): super().__init__(weight, sigmoid_normalization) self.classes = classes if skip_index_after is not None: self.skip_index_after = skip_index_after def dice(self, input, target, weight): return compute_per_channel_dice(input, target, weight=self.weight) def expand_as_one_hot(input, C, ignore_index=None): """ Converts NxDxHxW label image to NxCxDxHxW, where each label gets converted to its corresponding one-hot vector :param input: 4D input image (NxDxHxW) :param C: number of channels/labels :param ignore_index: ignore index to be kept during the expansion :return: 5D output image (NxCxDxHxW) """ if input.dim() == 5: return input assert input.dim() == 4 # expand the input tensor to Nx1xDxHxW before scattering input = input.unsqueeze(1) # create result tensor shape (NxCxDxHxW) shape = list(input.size()) shape[1] = C if ignore_index is not None: # create ignore_index mask for the result mask = input.expand(shape) == ignore_index # clone the lib tensor and zero out ignore_index in the input input = input.clone() input[input == ignore_index] = 0 # scatter to get the one-hot tensor result = torch.zeros(shape).to(input.device).scatter_(1, input, 1) # bring back the ignore_index in the result result[mask] = ignore_index return result else: # scatter to get the one-hot tensor return torch.zeros(shape).to(input.device).scatter_(1, input, 1) def compute_per_channel_dice(input, target, epsilon=1e-6, weight=None): """ Computes DiceCoefficient as defined in https://arxiv.org/abs/1606.04797 given a multi channel input and target. Assumes the input is a normalized probability, e.g. a result of Sigmoid or Softmax function. Args: input (torch.Tensor): NxCxSpatial input tensor target (torch.Tensor): NxCxSpatial target tensor epsilon (float): prevents division by zero weight (torch.Tensor): Cx1 tensor of weight per channel/class """ # input and target shapes must match assert input.size() == target.size(), "'input' and 'target' must have the same shape" input = flatten(input) target = flatten(target) target = target.float() # compute per channel Dice Coefficient intersect = (input * target).sum(-1) if weight is not None: intersect = weight * intersect # here we can use standard dice (input + target).sum(-1) or extension (see V-Net) (input^2 + target^2).sum(-1) denominator = (input * input).sum(-1) + (target * target).sum(-1) return 2 * (intersect / denominator.clamp(min=epsilon)) def flatten(tensor): """Flattens a given tensor such that the channel axis is first. The shapes are transformed as follows: (N, C, D, H, W) -> (C, N * D * H * W) """ # number of channels C = tensor.size(1) # new axis order axis_order = (1, 0) + tuple(range(2, tensor.dim())) # Transpose: (N, C, D, H, W) -> (C, N, D, H, W) transposed = tensor.permute(axis_order) # Flatten: (C, N, D, H, W) -> (C, N * D * H * W) return transposed.contiguous().view(C, -1) class _BaseProbLoss(nn.Module): def __init__(self, classes: int, weight: Optional[torch.Tensor] = None, sigmoid_normalization: bool = True): super().__init__() self.classes = classes self.register_buffer("weight", weight if weight is not None else torch.ones(classes)) self.norm = nn.Sigmoid() if sigmoid_normalization else nn.Softmax(dim=1) def _one_hot(self, target: torch.Tensor, shape) -> torch.Tensor: out = torch.zeros(shape, dtype=torch.float32, device=target.device) out.scatter_(1, target.long(), 1) return out def forward(self, logits: torch.Tensor, target: torch.Tensor): # -------- one-hot -------- if target.dim() == logits.dim() - 1: target = target.unsqueeze(1) oh = self._one_hot(target, logits.shape) prob = self.norm(logits) # logits → prob(0-1) # print("prob.shape =", prob.shape) # print("oh.shape =", oh.shape) # one-hot 后 # print("prob.type", prob.dtype) # 原始标签 # print("oh.type", oh.dtype) score = self.score(prob, oh) # 由子类实现 loss = 1.0 - score.mean() return loss, score.detach().cpu().numpy() def score(self, prob: torch.Tensor, target: torch.Tensor): raise NotImplementedError