# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/loss_functions.ipynb (unless otherwise specified).

__all__ = ['FocalLoss', 'FocalLossFlat']

# Cell
from fastai.vision.all import *

# Cell
class FocalLoss(nn.Module):
    # implementation adapted from https://amaarora.github.io/2020/06/29/FocalLoss.html
    # paper: https://arxiv.org/abs/1708.02002
    "Focal Loss"
    def __init__(self, alpha=.25, gamma=2, reduction='none', pos_weight=None):
        super(FocalLoss, self).__init__()
        self.alpha = torch.tensor([alpha, 1-alpha])#.cuda()
        self.gamma = gamma
        self.reduction  = reduction
        self.pos_weight = pos_weight

    def forward(self, inputs, targets):
        BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets,
                                                      reduction=self.reduction,
                                                      pos_weight=self.pos_weight)
        targets = targets.type(torch.long)
        at = self.alpha.gather(0, targets.data.view(-1))
        pt = torch.exp(-BCE_loss)
        #F_loss = at*(1-pt)**self.gamma * BCE_loss
        F_loss = at*(1-pt.view(-1))**self.gamma * BCE_loss.view(-1)
        return F_loss.mean()
@log_args
@delegates()
class FocalLossFlat(BaseLoss):
    @use_kwargs_dict(keep=True, alpha=.25, gamma=2., reduction='none', pos_weight=None)
    def __init__(self, thresh=0.5, axis=-1, floatify=True, **kwargs):
        super().__init__(FocalLoss, axis=axis, floatify=floatify, is_2d=False, flatten=False, **kwargs)
        self.thresh = thresh

    def decodes(self, x):    return x > self.thresh
    def activation(self, x): return torch.sigmoid(x)