from .dice import DiceLoss
import torch.nn as nn

__all__ = ['DiceLoss', 'get_loss']

def get_loss(name: str):
    """Factory function for segmentation losses."""
    if name == 'ce':
        return nn.CrossEntropyLoss()
    elif name == 'dice':
        return DiceLoss()
    elif name == 'combined':
        return lambda p, t: nn.CrossEntropyLoss()(p, t) + DiceLoss()(p, t)
    else:
        raise ValueError(f"Unknown loss: {name}")