# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_losses.ipynb.

# %% auto 0
__all__ = ['CombinedLoss']

# %% ../nbs/04_losses.ipynb 3
from fastai.losses import FocalLossFlat, DiceLoss
from fastcore.basics import store_attr

# %% ../nbs/04_losses.ipynb 4
class CombinedLoss:
    "Dice and Focal combined"
    def __init__(self, axis=1, smooth=1.0, alpha=1.0, gamma=2.0):
        store_attr()
        self.focal_loss = FocalLossFlat(gamma = gamma, axis=axis)
        self.dice_loss =  DiceLoss(axis, smooth, reduction='mean')
        
    def __call__(self, pred, targ, reduction=None):
        return (1. - self.alpha) * self.focal_loss(pred, targ) + self.alpha * self.dice_loss(pred, targ)
        #return self.dice_loss(pred, targ)
    
    def decodes(self, x):    return x.argmax(dim=self.axis)
    def activation(self, x): return F.softmax(x, dim=self.axis)
