MRI Segmentation Losses#

class atommic.collections.segmentation.losses.CategoricalCrossEntropyLoss(*args: Any, **kwargs: Any)[source]#

Bases: Loss

Wrapper around PyTorch’s CrossEntropyLoss to support 2D and 3D inputs.

__init__(include_background: bool = True, num_samples: int = 50, ignore_index: int = -100, reduction: str = 'mean', label_smoothing: float = 0.0, weight: Optional[List] = None, to_onehot_y: bool = False, num_segmentation_classes: Optional[int] = None)[source]#

Inits CategoricalCrossEntropyLoss.

Parameters

  • include_background (bool) – Whether to include the computation on the first channel of the predicted output. Default is True.

  • num_samples (int, optional) – Number of Monte Carlo samples. Default is 50.

  • ignore_index (int, optional) – Index to ignore. Default is -100.

  • reduction (Union[str, None]) – Specifies the reduction to apply: none: no reduction will be applied. mean: reduction with averaging over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D sum: reduction with summing over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D Default is mean.

  • label_smoothing (float, optional) – Label smoothing. Default is 0.0.

  • weight (list of floats, optional) – List with weights for each class. Default is None.

  • to_onehot_y (bool) – Whether to convert y into the one-hot format. Default is False.

  • num_segmentation_classes (int) – Total number of segmentation classes. Default is None.

forward(target: torch.Tensor, _input: torch.Tensor, pred_log_var: Optional[torch.Tensor] = None) → torch.Tensor[source]#

Forward pass of CategoricalCrossEntropyLoss.

Parameters

  • target (torch.Tensor) – Target tensor. Shape: (batch_size, num_classes, *spatial_dims)

  • _input (torch.Tensor) – Prediction tensor. Shape: (batch_size, num_classes, *spatial_dims)

  • pred_log_var (torch.Tensor, optional) – Prediction log variance tensor. Shape: (batch_size, num_classes, *spatial_dims). Default is None.

Returns

CategoricalCrossEntropy Loss

Return type

torch.Tensor

class atommic.collections.segmentation.losses.BinaryCrossEntropyLoss(*args: Any, **kwargs: Any)[source]#

Bases: Loss

Wrapper around PyTorch’s BinaryCrossEntropyLoss to support 2D and 3D inputs.

__init__(include_background: bool = True, num_samples: int = 50, weight: Optional[List] = None, reduction: str = 'mean', to_onehot_y: bool = False, num_segmentation_classes: Optional[int] = None)[source]#

Inits BinaryCrossEntropyLoss.

Parameters

  • include_background (bool) – Whether to include the computation on the first channel of the predicted output. Default is False.

  • num_samples (int, optional) – Number of Monte Carlo samples. Default is 50.

  • weight (list of floats, optional) – List of weight for each sample. Default is None.

  • reduction (Union[str, None]) – Specifies the reduction to apply: none: no reduction will be applied. mean: reduction with averaging over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D sum: reduction with summing over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D Default is mean.

  • to_onehot_y (bool) – Whether to convert y into the one-hot format. Default is False.

  • num_segmentation_classes (int) – Total number of segmentation classes. Default is None.

forward(target: torch.Tensor, _input: torch.Tensor, pred_log_var: Optional[torch.Tensor] = None) → torch.Tensor[source]#

Forward pass of BinaryCrossEntropyLoss.

Parameters

  • target (torch.Tensor) – Target tensor. Shape: (batch_size, num_classes, *spatial_dims)

  • _input (torch.Tensor) – Prediction tensor. Shape: (batch_size, num_classes, *spatial_dims)

  • pred_log_var (torch.Tensor, optional) – Prediction log variance tensor. Shape: (batch_size, num_classes, *spatial_dims). Default is None.

Returns

BinaryCrossEntropy Loss

Return type

torch.Tensor

class atommic.collections.segmentation.losses.FocalLoss(*args: Any, **kwargs: Any)[source]#

Bases: Loss

FocalLoss is an extension of BCEWithLogitsLoss that down-weights loss from high confidence correct predictions.

Reimplementation of the Focal Loss described in:

  • [“Focal Loss for Dense Object Detection”](https://arxiv.org/abs/1708.02002), T. Lin et al., ICCV 2017

  • “AnatomyNet: Deep learning for fast and fully automated whole-volume segmentation of head and neck anatomy”, Zhu et al., Medical Physics 2018

__init__(include_background: bool = True, gamma: float = 2.0, alpha: Optional[float] = None, weight: Optional[Union[Sequence[float], float, int, torch.Tensor]] = None, reduction: str = 'mean', use_softmax: bool = False, to_onehot_y: bool = False, num_segmentation_classes: Optional[int] = None) → None[source]#

Inits FocalLoss

Parameters

  • include_background (bool) – whether to include the computation on the first channel of the predicted output. Default is True.

  • gamma (float, optional) – Value of the exponent gamma in the definition of the Focal loss. Default is 2

  • alpha (float, optional) – Value of the alpha: [0,1] in the definition of the alpha-balanced Focal loss. Default is None

  • weight (torch.Tensor, optional) – Weight for each class. Default is None

  • reduction (Union[str, None]) – Specifies the reduction to apply: none: no reduction will be applied. mean: reduction with averaging over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D sum: reduction with summing over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D Default is mean.

  • use_softmax (bool, optional) – option to compute the focal loss as a categorical cross-entropy. Default is False

  • to_onehot_y (bool) – Whether to convert y into the one-hot format. Default is False.

  • num_segmentation_classes (int) – Total number of segmentation classes. Default is None.

forward(target: torch.Tensor, _input: torch.Tensor) → Tuple[Union[torch.Tensor, Any], torch.Tensor][source]#

Forward pass of FocalLoss.

Parameters

  • target (torch.Tensor) – Target tensor. Shape: (batch_size, num_classes, *spatial_dims)

  • _input (torch.Tensor) – Prediction tensor. Shape: (batch_size, num_classes, *spatial_dims)

Returns

Focal Loss

Return type

torch.Tensor

class atommic.collections.segmentation.losses.Dice(*args: Any, **kwargs: Any)[source]#

Bases: Loss

Wrapper for monai.losses.DiceLoss. Computes the average Dice loss between two tensors.

The data input (BNHW[D] where N is number of classes) is compared with ground truth target (BNHW[D]).

Note that axis N of input is expected to be logits or probabilities for each class, if passing logits as input, must set sigmoid=True or softmax=True, or specifying other_act. And the same axis of target can be 1 or N (one-hot format).

The smooth_nr and smooth_dr parameters are values added to the intersection and union components of the inter-over-union calculation to smooth results respectively, these values should be small.

The original paper: Milletari, F. et. al. (2016) V-Net: Fully Convolutional Neural Networks forVolumetric Medical Image Segmentation, 3DV, 2016.

Examples

>>> import torch
>>> from atommic.collections.segmentation.losses.dice import Dice
>>> pred = torch.tensor([[[[0.1, 0.2, 0.3, 0.4, 0.5],
...                        [0.1, 0.2, 0.3, 0.4, 0.5],
...                        [0.1, 0.2, 0.3, 0.4, 0.5],
...                        [0.1, 0.2, 0.3, 0.4, 0.5],
...                        [0.1, 0.2, 0.3, 0.4, 0.5]]],
...                       [[[0.1, 0.2, 0.3, 0.4, 0.5],
...                         [0.1, 0.2, 0.3, 0.4, 0.5],
...                         [0.1, 0.2, 0.3, 0.4, 0.5],
...                         [0.1, 0.2, 0.3, 0.4, 0.5],
...                         [0.1, 0.2, 0.3, 0.4, 0.5]]]])
>>> target = torch.tensor([[[[0, 0, 0, 0, 0],
...                          [0, 0, 0, 0, 0],
...                          [0, 0, 0, 0, 0],
...                          [0, 0, 0, 0, 0],
...                          [0, 0, 0, 0, 0]]],
...                         [[[1, 1, 1, 1, 1],
...                           [1, 1, 1, 1, 1],
...                           [1, 1, 1, 1, 1],
...                           [1, 1, 1, 1, 1],
...                           [1, 1, 1, 1, 1]]]])
>>> dice = Dice(include_background=False, to_onehot_y=True, sigmoid=False, softmax=False)
>>> dice(pred, target)
tensor(0.5000)
__init__(include_background: bool = True, to_onehot_y: bool = False, sigmoid: bool = True, softmax: bool = False, other_act: Optional[Callable] = None, squared_pred: bool = False, jaccard: bool = False, flatten: bool = False, reduction: str = 'mean', smooth_nr: float = 1e-05, smooth_dr: float = 1e-05, batch: bool = True, num_segmentation_classes: Optional[int] = None)[source]#

Inits Dice.

Parameters

  • include_background (bool) – whether to skip Dice computation on the first channel of the predicted output. Default is True.

  • to_onehot_y (bool) – Whether to convert y into the one-hot format. Default is False.

  • sigmoid (bool) – Whether to add sigmoid function to the input data. Default is True.

  • softmax (bool) – Whether to add softmax function to the input data. Default is False.

  • other_act (Callable) – Use this parameter if you want to apply another type of activation layer. Default is None.

  • squared_pred (bool) – Whether to square the prediction before calculating Dice. Default is False.

  • jaccard (bool) – Whether to compute Jaccard Index as a loss. Default is False.

  • flatten (bool) – Whether to flatten input data. Default is False.

  • reduction (Union[str, None]) – Specifies the reduction to apply: none: no reduction will be applied. mean: reduction with averaging over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D sum: reduction with summing over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D Default is mean.

  • smooth_nr (float) – A small constant added to the numerator to avoid nan when all items are 0. Default is 1e-5.

  • smooth_dr (float) – A small constant added to the denominator to avoid nan when all items are 0. Default is 1e-5.

  • batch (bool) – If True, compute Dice loss for each batch and return a tensor with shape (batch_size,). If False, compute Dice loss for the whole batch and return a tensor with shape (1,). Default is True.

  • num_segmentation_classes (int) – Total number of segmentation classes. Default is None.

forward(target: torch.Tensor, _input: torch.Tensor) → Tuple[Union[torch.Tensor, Any], torch.Tensor][source]#

Forward pass of Dice.

Parameters

  • target (torch.Tensor) – Target tensor. Shape: (batch_size, num_classes, *spatial_dims) or (batch_size, 1, *spatial_dims)

  • _input (torch.Tensor) – Prediction tensor. Shape: (batch_size, num_classes, *spatial_dims)

Returns

Dice Loss

Return type

torch.Tensor

class atommic.collections.segmentation.losses.GeneralisedDice(*args: Any, **kwargs: Any)[source]#

Bases: Loss

Compute the Generalised Dice loss, as presented in [Sudre2017].

Adapted from: NifTK/NiftyNet

References

__init__(include_background: bool = True, to_onehot_y: bool = False, sigmoid: bool = True, softmax: bool = False, other_act: Optional[Callable] = None, w_type: str = 'square', reduction: str = 'mean', smooth_nr: float = 1e-05, smooth_dr: float = 1e-05, batch: bool = True, num_segmentation_classes: Optional[int] = None) → None[source]#

Inits GeneralisedDiceLoss.

Parameters

  • include_background (bool) – whether to skip Dice computation on the first channel of the predicted output. Default is True.

  • to_onehot_y (bool) – Whether to convert y into the one-hot format. Default is False.

  • sigmoid (bool) – Whether to add sigmoid function to the input data. Default is True.

  • softmax (bool) – Whether to add softmax function to the input data. Default is False.

  • other_act (Callable) – Use this parameter if you want to apply another type of activation layer. Default is None.

  • w_type ({"square", "simple", "uniform"}) – Type of function to transform ground truth volume to a weight factor. Defaults to "square".

  • reduction (Union[str, None]) – Specifies the reduction to apply: none: no reduction will be applied. mean: reduction with averaging over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D sum: reduction with summing over both batch and channel dimensions if input is 2D, or batch dimension only if input is 1D Default is mean.

  • smooth_nr (float) – A small constant added to the numerator to avoid nan when all items are 0. Default is 1e-5.

  • smooth_dr (float) – A small constant added to the denominator to avoid nan when all items are 0. Default is 1e-5.

  • batch (bool) – If True, compute Dice loss for each batch and return a tensor with shape (batch_size,). If False, compute Dice loss for the whole batch and return a tensor with shape (1,). Default is True.

  • num_segmentation_classes (int) – Total number of segmentation classes. Default is None.

forward(target: torch.Tensor, _input: torch.Tensor) → Tuple[Union[torch.Tensor, Any], torch.Tensor][source]#

Forward pass of GeneralisedDice.

Parameters

  • target (torch.Tensor) – Target tensor. Shape: (batch_size, num_classes, *spatial_dims) or (batch_size, 1, *spatial_dims)

  • _input (torch.Tensor) – Prediction tensor. Shape: (batch_size, num_classes, *spatial_dims)

Returns

GeneralizedDice Loss

Return type

torch.Tensor

w_func(grnd)[source]#

Compute the weight factor for the Generalised Dice loss.