# AUTOGENERATED! DO NOT EDIT! File to edit: 04_opt.ipynb (unless otherwise specified).

__all__ = ['DiceLoss', 'dice_similarity']


# Cell
import torch
import torch.nn as nn


# Cell
class DiceLoss(nn.Module):
    """
    Module to compute the Dice segmentation loss. Based on the following discussion:
    https://discuss.pytorch.org/t/one-hot-encoding-with-autograd-dice-loss/9781
    """
    def __init__(self, weights=None, ignore_index=None, eps=0.0001):
        super(DiceLoss, self).__init__()
        self.weights = weights
        self.ignore_index = ignore_index
        self.eps = eps

    def forward(self, output, target):
        """
        Arguments:
            output: (N, C, H, W) tensor of probabilities for the predicted output
            target: (N, H, W) tensor corresponding to the pixel-wise labels
        Returns:
            loss: the Dice loss averaged over channels
        """
        encoded_target = output.detach() * 0
        if self.ignore_index is not None:
            mask = target == self.ignore_index
            target = target.clone()
            target[mask] = 0
            encoded_target.scatter_(1, target.unsqueeze(1), 1)
            mask = mask.unsqueeze(1).expand_as(encoded_target)
            encoded_target[mask] = 0
        else:
            encoded_target.scatter_(1, target.unsqueeze(1), 1)

        if self.weights is None:
            self.weights = 1

        intersection = output * encoded_target
        numerator = 2 * intersection.sum(0).sum(1).sum(1)
        denominator = output + encoded_target

        if self.ignore_index is not None:
            denominator[mask] = 0
        denominator = denominator.sum(0).sum(1).sum(1) + self.eps
        loss_per_channel = self.weights * (1 - (numerator / denominator))

        return loss_per_channel.sum() / output.size(1)


# Cell
def dice_similarity(output, target, weights=None, ignore_index=None, eps=1e-8):
    """
    Arguments:
        output: (N, C, H, W) tensor of model output
        target: (N, H, W) tensor corresponding to the pixel-wise labels
    Returns:
        loss: the Dice loss averaged over channels
    """
    prediction = torch.argmax(output, dim=1)
    encoded_prediction = output.detach() * 0
    encoded_prediction.scatter_(1, prediction.unsqueeze(1), 1)

    encoded_target = output.detach() * 0
    if ignore_index is not None:
        mask = target == ignore_index
        target = target.clone()
        target[mask] = 0
        encoded_target.scatter_(1, target.unsqueeze(1), 1)
        mask = mask.unsqueeze(1).expand_as(encoded_target)
        encoded_target[mask] = 0
    else:
        encoded_target.scatter_(1, target.unsqueeze(1), 1)

    if weights is None:
        weights = 1

    intersection = encoded_prediction * encoded_target
    numerator = 2 * intersection.sum(0).sum(1).sum(1) + eps
    denominator = intersection + encoded_target
    if ignore_index is not None:
        denominator[mask] = 0
    denominator = denominator.sum(0).sum(1).sum(1) + eps
    acc_per_channel = weights * ((numerator / denominator))

    return acc_per_channel.sum() / output.size(1)