<gh_stars>0
import torch
from torch import nn
from torch.nn import functional as F
from torch.autograd import Variable
from torch.nn import MSELoss, SmoothL1Loss, L1Loss


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)


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 expand_as_one_hot(input, C, ignore_index=None):
    """
    Converts NxSPATIAL label image to NxCxSPATIAL, where each label gets
        converted to its corresponding one-hot vector. It is assumed that
        the batch dimension is present.
    Args:
        input (torch.Tensor): 3D/4D input image
        C (int): number of channels/labels
        ignore_index (int): ignore index to be kept during the expansion
    Returns:
        4D/5D output torch.Tensor (NxCxSPATIAL)
    """
    assert input.dim() == 4

    # expand the input tensor to Nx1xSPATIAL before scattering
    input = input.unsqueeze(1)
    # create output tensor shape (NxCxSPATIAL)
    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 src 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)


class _AbstractDiceLoss(nn.Module):
    """
    Base class for different implementations of Dice loss.
    """

    def __init__(self, opt):
        super(_AbstractDiceLoss, self).__init__()
        self.register_buffer('weight', opt.get('weight', None))
        self.num_classes = opt.num_classes
        normalization = opt.get('normalization', 'sigmoid')
        # 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 
        # `normalization=Softmax`
        assert normalization in ['sigmoid', 'softmax', 'none']
        if normalization == 'sigmoid':
            self.normalization = nn.Sigmoid()
        elif normalization == 'softmax':
            self.normalization = nn.Softmax(dim=1)
        else:
            self.normalization = lambda x: x

    def dice(self, input, target, weight):
        # actual Dice score computation; to be implemented by the subclass
        raise NotImplementedError

    def forward(self, input, target):
        # get probabilities from logits
        input = self.normalization(input)
        if input.dim() == target.dim() + 1:
            target = expand_as_one_hot(target, self.num_classes)

        # compute per channel Dice coefficient
        per_channel_dice = self.dice(input, target, weight=self.weight)

        # average Dice score across all channels/classes
        return 1. - torch.mean(per_channel_dice)


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.
    The input to the loss function is assumed to be a logit and will be 
        normalized by the Sigmoid function.
    """

    def __init__(self, opt):
        super().__init__(opt)

    def dice(self, input, target, *args, **kwargs):
        return compute_per_channel_dice(input, target, weight=self.weight)


class GeneralizedDiceLoss(_AbstractDiceLoss):
    epsilon=1e-6

    """Computes Generalized Dice Loss (GDL) as described in https://arxiv.org/pdf/1707.03237.pdf.
    """

    def __init__(self, opt):
        super().__init__(opt)

    def dice(self, input, target, *args, **kwargs):
        assert input.size() == target.size(), "'input' and 'target' must have the same shape"

        input = flatten(input)
        target = flatten(target)
        target = target.float()

        if input.size(0) == 1:
            # for GDL to make sense we need at least 2 channels 
            #   (see https://arxiv.org/pdf/1707.03237.pdf)
            # put foreground and background voxels in separate channels
            input = torch.cat((input, 1 - input), dim=0)
            target = torch.cat((target, 1 - target), dim=0)

        # GDL weighting: the contribution of each label is corrected 
        # by the inverse of its volume
        w_l = target.sum(-1)
        w_l = 1 / (w_l * w_l).clamp(min=self.epsilon)
        w_l.requires_grad = False

        intersect = (input * target).sum(-1)
        intersect = intersect * w_l

        denominator = (input + target).sum(-1)
        denominator = (denominator * w_l).clamp(min=self.epsilon)

        return 2 * (intersect.sum() / denominator.sum())


class BCEDiceLoss(nn.Module):
    """Linear combination of BCE and Dice losses"""

    def __init__(self, opt):
        super(BCEDiceLoss, self).__init__()
        self.alpha = opt.bce_weight
        self.bce = nn.BCEWithLogitsLoss()
        self.beta = opt.dice_weight
        self.dice = DiceLoss(opt)

    def forward(self, input, target):
        return self.alpha * self.bce(input, target) \
            + self.beta * self.dice(input, target)


class WeightedCrossEntropyLoss(nn.Module):
    """WeightedCrossEntropyLoss (WCE) as described in 
        https://arxiv.org/pdf/1707.03237.pdf
    """

    def __init__(self, opt):
        super(WeightedCrossEntropyLoss, self).__init__()
        self.ignore_index = opt.get('ignore_index', -1)

    def forward(self, input, target):
        weight = self._class_weights(input)
        return F.cross_entropy(input, target,
                               weight=weight,
                               ignore_index=self.ignore_index)

    @staticmethod
    def _class_weights(input):
        # normalize the input first
        input = F.softmax(input, dim=1)
        flattened = flatten(input)
        nominator = (1. - flattened).sum(-1)
        denominator = flattened.sum(-1)
        class_weights = Variable(nominator / denominator, requires_grad=False)
        return class_weights


class PixelWiseCrossEntropyLoss(nn.Module):
    def __init__(self, opt):
        super(PixelWiseCrossEntropyLoss, self).__init__()
        self.register_buffer('class_weights', opt.get('class_weights', None))
        self.ignore_index = opt.get('ignore_index', None)
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, input, target, weights):
        assert target.size() == weights.size()
        # normalize the input
        log_probabilities = self.log_softmax(input)
        # standard CrossEntropyLoss requires the target to be (NxDxHxW),
        # so we need to expand it to (NxCxDxHxW)
        target = expand_as_one_hot(target,
                                   C=input.size()[1],
                                   ignore_index=self.ignore_index)
        # expand weights
        weights = weights.unsqueeze(0)
        weights = weights.expand_as(input)

        # create default class_weights if None
        if self.class_weights is None:
            class_weights = torch.ones(input.size()[1]).float().to(input.device)
        else:
            class_weights = self.class_weights

        # resize class_weights to be broadcastable into the weights
        class_weights = class_weights.view(1, -1, 1, 1, 1)

        # multiply weights tensor by class weights
        weights = class_weights * weights

        # compute the losses
        result = -weights * target * log_probabilities
        # average the losses
        return result.mean()


class WeightedSmoothL1Loss(nn.SmoothL1Loss):
    def __init__(self, opt):
        super().__init__(reduction="none")
        self.threshold = opt.threshold
        self.apply_below_threshold = opt.get('apply_below_threshold', True)
        self.weight = opt.initial_weight

    def forward(self, input, target):
        l1 = super().forward(input, target)

        if self.apply_below_threshold:
            mask = target < self.threshold
        else:
            mask = target >= self.threshold

        l1[mask] = l1[mask] * self.weight

        return l1.mean()

import torch
from torch import nn
from torch.nn import functional as F
from torch.autograd import Variable
from torch.nn import MSELoss, SmoothL1Loss, L1Loss


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)


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 expand_as_one_hot(input, C, ignore_index=None):
    """
    Converts NxSPATIAL label image to NxCxSPATIAL, where each label gets
        converted to its corresponding one-hot vector. It is assumed that
        the batch dimension is present.
    Args:
        input (torch.Tensor): 3D/4D input image
        C (int): number of channels/labels
        ignore_index (int): ignore index to be kept during the expansion
    Returns:
        4D/5D output torch.Tensor (NxCxSPATIAL)
    """
    assert input.dim() == 4

    # expand the input tensor to Nx1xSPATIAL before scattering
    input = input.unsqueeze(1)
    # create output tensor shape (NxCxSPATIAL)
    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 src 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)


class _AbstractDiceLoss(nn.Module):
    """
    Base class for different implementations of Dice loss.
    """

    def __init__(self, opt):
        super(_AbstractDiceLoss, self).__init__()
        self.register_buffer('weight', opt.get('weight', None))
        self.num_classes = opt.num_classes
        normalization = opt.get('normalization', 'sigmoid')
        # 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 
        # `normalization=Softmax`
        assert normalization in ['sigmoid', 'softmax', 'none']
        if normalization == 'sigmoid':
            self.normalization = nn.Sigmoid()
        elif normalization == 'softmax':
            self.normalization = nn.Softmax(dim=1)
        else:
            self.normalization = lambda x: x

    def dice(self, input, target, weight):
        # actual Dice score computation; to be implemented by the subclass
        raise NotImplementedError

    def forward(self, input, target):
        # get probabilities from logits
        input = self.normalization(input)
        if input.dim() == target.dim() + 1:
            target = expand_as_one_hot(target, self.num_classes)

        # compute per channel Dice coefficient
        per_channel_dice = self.dice(input, target, weight=self.weight)

        # average Dice score across all channels/classes
        return 1. - torch.mean(per_channel_dice)


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.
    The input to the loss function is assumed to be a logit and will be 
        normalized by the Sigmoid function.
    """

    def __init__(self, opt):
        super().__init__(opt)

    def dice(self, input, target, *args, **kwargs):
        return compute_per_channel_dice(input, target, weight=self.weight)


class GeneralizedDiceLoss(_AbstractDiceLoss):
    epsilon=1e-6

    """Computes Generalized Dice Loss (GDL) as described in https://arxiv.org/pdf/1707.03237.pdf.
    """

    def __init__(self, opt):
        super().__init__(opt)

    def dice(self, input, target, *args, **kwargs):
        assert input.size() == target.size(), "'input' and 'target' must have the same shape"

        input = flatten(input)
        target = flatten(target)
        target = target.float()

        if input.size(0) == 1:
            # for GDL to make sense we need at least 2 channels 
            #   (see https://arxiv.org/pdf/1707.03237.pdf)
            # put foreground and background voxels in separate channels
            input = torch.cat((input, 1 - input), dim=0)
            target = torch.cat((target, 1 - target), dim=0)

        # GDL weighting: the contribution of each label is corrected 
        # by the inverse of its volume
        w_l = target.sum(-1)
        w_l = 1 / (w_l * w_l).clamp(min=self.epsilon)
        w_l.requires_grad = False

        intersect = (input * target).sum(-1)
        intersect = intersect * w_l

        denominator = (input + target).sum(-1)
        denominator = (denominator * w_l).clamp(min=self.epsilon)

        return 2 * (intersect.sum() / denominator.sum())


class BCEDiceLoss(nn.Module):
    """Linear combination of BCE and Dice losses"""

    def __init__(self, opt):
        super(BCEDiceLoss, self).__init__()
        self.alpha = opt.bce_weight
        self.bce = nn.BCEWithLogitsLoss()
        self.beta = opt.dice_weight
        self.dice = DiceLoss(opt)

    def forward(self, input, target):
        return self.alpha * self.bce(input, target) \
            + self.beta * self.dice(input, target)


class WeightedCrossEntropyLoss(nn.Module):
    """WeightedCrossEntropyLoss (WCE) as described in 
        https://arxiv.org/pdf/1707.03237.pdf
    """

    def __init__(self, opt):
        super(WeightedCrossEntropyLoss, self).__init__()
        self.ignore_index = opt.get('ignore_index', -1)

    def forward(self, input, target):
        weight = self._class_weights(input)
        return F.cross_entropy(input, target,
                               weight=weight,
                               ignore_index=self.ignore_index)

    @staticmethod
    def _class_weights(input):
        # normalize the input first
        input = F.softmax(input, dim=1)
        flattened = flatten(input)
        nominator = (1. - flattened).sum(-1)
        denominator = flattened.sum(-1)
        class_weights = Variable(nominator / denominator, requires_grad=False)
        return class_weights


class PixelWiseCrossEntropyLoss(nn.Module):
    def __init__(self, opt):
        super(PixelWiseCrossEntropyLoss, self).__init__()
        self.register_buffer('class_weights', opt.get('class_weights', None))
        self.ignore_index = opt.get('ignore_index', None)
        self.log_softmax = nn.LogSoftmax(dim=1)

    def forward(self, input, target, weights):
        assert target.size() == weights.size()
        # normalize the input
        log_probabilities = self.log_softmax(input)
        # standard CrossEntropyLoss requires the target to be (NxDxHxW),
        # so we need to expand it to (NxCxDxHxW)
        target = expand_as_one_hot(target,
                                   C=input.size()[1],
                                   ignore_index=self.ignore_index)
        # expand weights
        weights = weights.unsqueeze(0)
        weights = weights.expand_as(input)

        # create default class_weights if None
        if self.class_weights is None:
            class_weights = torch.ones(input.size()[1]).float().to(input.device)
        else:
            class_weights = self.class_weights

        # resize class_weights to be broadcastable into the weights
        class_weights = class_weights.view(1, -1, 1, 1, 1)

        # multiply weights tensor by class weights
        weights = class_weights * weights

        # compute the losses
        result = -weights * target * log_probabilities
        # average the losses
        return result.mean()


class WeightedSmoothL1Loss(nn.SmoothL1Loss):
    def __init__(self, opt):
        super().__init__(reduction="none")
        self.threshold = opt.threshold
        self.apply_below_threshold = opt.get('apply_below_threshold', True)
        self.weight = opt.initial_weight

    def forward(self, input, target):
        l1 = super().forward(input, target)

        if self.apply_below_threshold:
            mask = target < self.threshold
        else:
            mask = target >= self.threshold

        l1[mask] = l1[mask] * self.weight

        return l1.mean()
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 # new axis order # Transpose: (N, C, D, H, W) -> (C, N, D, H, W) # Flatten: (C, N, D, H, W) -> (C, N * D * H * W) 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 # compute per channel Dice Coefficient # here we can use standard dice (input + target).sum(-1) or extension # (see V-Net) (input^2 + target^2).sum(-1) Converts NxSPATIAL label image to NxCxSPATIAL, where each label gets         converted to its corresponding one-hot vector. It is assumed that         the batch dimension is present.     Args:         input (torch.Tensor): 3D/4D input image         C (int): number of channels/labels         ignore_index (int): ignore index to be kept during the expansion     Returns:         4D/5D output torch.Tensor (NxCxSPATIAL) # expand the input tensor to Nx1xSPATIAL before scattering # create output tensor shape (NxCxSPATIAL) # create ignore_index mask for the result # clone the src tensor and zero out ignore_index in the input # scatter to get the one-hot tensor # bring back the ignore_index in the result # scatter to get the one-hot tensor Base class for different implementations of Dice loss. # 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 # `normalization=Softmax` # actual Dice score computation; to be implemented by the subclass # get probabilities from logits # compute per channel Dice coefficient # average Dice score across all channels/classes 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.     The input to the loss function is assumed to be a logit and will be          normalized by the Sigmoid function. Computes Generalized Dice Loss (GDL) as described in https://arxiv.org/pdf/1707.03237.pdf. # for GDL to make sense we need at least 2 channels #   (see https://arxiv.org/pdf/1707.03237.pdf) # put foreground and background voxels in separate channels # GDL weighting: the contribution of each label is corrected # by the inverse of its volume Linear combination of BCE and Dice losses WeightedCrossEntropyLoss (WCE) as described in          https://arxiv.org/pdf/1707.03237.pdf # normalize the input first # normalize the input # standard CrossEntropyLoss requires the target to be (NxDxHxW), # so we need to expand it to (NxCxDxHxW) # expand weights # create default class_weights if None # resize class_weights to be broadcastable into the weights # multiply weights tensor by class weights # compute the losses # average the losses