from torch import nn

from torch import optim

from rocaseg.components.losses import CrossEntropyLoss

from rocaseg.components.metrics import (confusion_matrix, dice_score,

                                        dice_score_from_cm)

from rocaseg.components.checkpoint import CheckpointHandler

dict_losses = {

    'bce_loss': nn.BCEWithLogitsLoss,

    'multi_ce_loss': CrossEntropyLoss,

}

dict_metrics = {

    'confusion_matrix': confusion_matrix,

    'dice_score': dice_score,

    'bce_loss': nn.BCELoss(),

}

dict_optimizers = {

    'sgd': optim.SGD,

    'adam': optim.Adam,

}

__all__ = [

    'dict_losses',

    'dict_metrics',

    'dict_optimizers',

    'confusion_matrix',

    'dice_score',

    'dice_score_from_cm',

    'CheckpointHandler',

]