torch_em.loss.wrapper

  1import torch
  2import torch.nn as nn
  3
  4
  5class LossWrapper(nn.Module):
  6    """ Wrapper around a torch loss function.
  7
  8    Applies transformations to prediction and/or target before passing it to the loss.
  9    """
 10    def __init__(self, loss, transform):
 11        super().__init__()
 12        self.loss = loss
 13
 14        if not callable(transform):
 15            raise ValueError("transform has to be callable.")
 16        self.transform = transform
 17        self.init_kwargs = {'loss': loss, 'transform': transform}
 18
 19    def apply_transform(self, prediction, target, **kwargs):
 20        # check if the tensors (prediction and target are lists)
 21        # if they are, apply the transform to each element inidvidually
 22        if isinstance(prediction, (list, tuple)):
 23            assert isinstance(target, (list, tuple))
 24            transformed_prediction, transformed_target = [], []
 25            for pred, targ in zip(prediction, target):
 26                tr_pred, tr_targ = self.transform(pred, targ, **kwargs)
 27                transformed_prediction.append(tr_pred)
 28                transformed_target.append(tr_targ)
 29            return transformed_prediction, transformed_target
 30        # tensor input
 31        else:
 32            prediction, target = self.transform(prediction, target, **kwargs)
 33            return prediction, target
 34
 35    def forward(self, prediction, target, **kwargs):
 36        prediction, target = self.apply_transform(prediction, target, **kwargs)
 37        loss = self.loss(prediction, target)
 38        return loss
 39
 40
 41#
 42# Loss transformations
 43#
 44
 45
 46class ApplyMask:
 47    def _crop(prediction, target, mask, channel_dim):
 48        if mask.shape[channel_dim] != 1:
 49            raise ValueError(
 50                "_crop only supports a mask with a singleton channel axis. \
 51                Please consider using masking_method=multiply."
 52            )
 53        mask = mask.type(torch.bool)
 54        # remove singleton axis
 55        mask = mask.squeeze(channel_dim)
 56        # move channel axis to end
 57        prediction = prediction.moveaxis(channel_dim, -1)
 58        target = target.moveaxis(channel_dim, -1)
 59        # output has shape N x C
 60        # correct for torch_em.loss.dice.flatten_samples
 61        return prediction[mask], target[mask]
 62
 63    def _multiply(prediction, target, mask, channel_dim):
 64        prediction = prediction * mask
 65        target = target * mask
 66        return prediction, target
 67
 68    MASKING_FUNCS = {
 69        "crop": _crop,
 70        "multiply": _multiply,
 71    }
 72
 73    def __init__(self, masking_method="crop", channel_dim=1):
 74        if masking_method not in self.MASKING_FUNCS.keys():
 75            raise ValueError(f"{masking_method} is not available, please use one of {list(self.MASKING_FUNCS.keys())}.")
 76        self.masking_func = self.MASKING_FUNCS[masking_method]
 77        self.channel_dim = channel_dim
 78
 79        self.init_kwargs = {
 80            "masking_method": masking_method,
 81            "channel_dim": channel_dim,
 82        }
 83
 84    def __call__(self, prediction, target, mask):
 85        mask.requires_grad = False
 86        return self.masking_func(prediction, target, mask, self.channel_dim)
 87
 88
 89class ApplyAndRemoveMask(ApplyMask):
 90    def __call__(self, prediction, target):
 91        assert target.dim() == prediction.dim(), f"{target.dim()}, {prediction.dim()}"
 92        assert target.size(1) == 2 * prediction.size(1), f"{target.size(1)}, {prediction.size(1)}"
 93        assert target.shape[2:] == prediction.shape[2:], f"{str(target.shape)}, {str(prediction.shape)}"
 94        seperating_channel = target.size(1) // 2
 95        mask = target[:, seperating_channel:]
 96        target = target[:, :seperating_channel]
 97        prediction, target = super().__call__(prediction, target, mask)
 98        return prediction, target
 99
100
101class MaskIgnoreLabel(ApplyMask):
102    def __init__(self, ignore_label=-1, masking_method="crop", channel_dim=1):
103        super().__init__(masking_method, channel_dim)
104        self.ignore_label = ignore_label
105        self.init_kwargs["ignore_label"] = ignore_label
106
107    def __call__(self, prediction, target):
108        mask = (target != self.ignore_label)
109        prediction, target = super().__call__(prediction, target, mask)
110        return prediction, target