torch_em.loss.affinity_side_loss

  1import numpy as np
  2import torch
  3import torch.nn as nn
  4from .dice import dice_score
  5
  6
  7def shift_tensor(tensor, offset):
  8    """ Shift a tensor by the given (spatial) offset.
  9    Arguments:
 10        tensor [torch.Tensor] - 4D (=2 spatial dims) or 5D (=3 spatial dims) tensor.
 11            Needs to be of float type.
 12        offset (tuple) - 2d or 3d spatial offset used for shifting the tensor
 13    """
 14
 15    ndim = len(offset)
 16    assert ndim in (2, 3)
 17    diff = tensor.dim() - ndim
 18
 19    # don't pad for the first dimensions
 20    # (usually batch and/or channel dimension)
 21    slice_ = diff * [slice(None)]
 22
 23    # torch padding behaviour is a bit weird.
 24    # we use nn.ReplicationPadND
 25    # (torch.nn.functional.pad is even weirder and ReflectionPad is not supported in 3d)
 26    # still, padding needs to be given in the inverse spatial order
 27
 28    # add padding in inverse spatial order
 29    padding = []
 30    for off in offset[::-1]:
 31        # if we have a negative offset, we need to shift "to the left",
 32        # which means padding at the right border
 33        # if we have a positive offset, we need to shift "to the right",
 34        # which means padding to the left border
 35        padding.extend([max(0, off), max(0, -off)])
 36
 37    # add slicing in the normal spatial order
 38    for off in offset:
 39        if off == 0:
 40            slice_.append(slice(None))
 41        elif off > 0:
 42            slice_.append(slice(None, -off))
 43        else:
 44            slice_.append(slice(-off, None))
 45
 46    # pad the spatial part of the tensor with replication padding
 47    slice_ = tuple(slice_)
 48    padding = tuple(padding)
 49    padder = nn.ReplicationPad2d if ndim == 2 else nn.ReplicationPad3d
 50    padder = padder(padding)
 51    shifted = padder(tensor)
 52
 53    # slice the oadded tensor to get the spatially shifted tensor
 54    shifted = shifted[slice_]
 55    assert shifted.shape == tensor.shape
 56
 57    return shifted
 58
 59
 60def invert_offsets(offsets):
 61    return [[-off for off in offset] for offset in offsets]
 62
 63
 64def segmentation_to_affinities(segmentation, offsets):
 65    """ Transform segmentation to affinities.
 66    Arguments:
 67        segmentation [torch.tensor] - 4D (2 spatial dims) or 5D (3 spatial dims) segmentation tensor.
 68            The channel axis (= dimension 1) needs to be a singleton.
 69        offsets [list[tuple]] - list of offsets for which to compute the affinities.
 70    """
 71    assert segmentation.shape[1] == 1, f"{segmentation.shape}"
 72    # shift the segmentation and substract the shifted tensor from the segmentation
 73    # we need to shift in the opposite direction of the offsets, so we invert them
 74    # before applying the shift
 75    offsets_ = invert_offsets(offsets)
 76    shifted = torch.cat([shift_tensor(segmentation.float(), off) for off in offsets_], dim=1)
 77    affs = (segmentation - shifted)
 78    # the affinities are 1, where we had the same segment id (the difference is 0)
 79    # and 0 otherwise
 80    affs.eq_(0.)
 81    return affs
 82
 83
 84def embeddings_to_affinities(embeddings, offsets, delta):
 85    """ Transform embeddings to affinities.
 86    """
 87    # shift the embeddings by the offsets and stack them along a new axis
 88    # we need to shift in the opposite direction of the offsets, so we invert them
 89    # before applying the shift
 90    offsets_ = invert_offsets(offsets)
 91    shifted = torch.cat([shift_tensor(embeddings, off).unsqueeze(1) for off in offsets_], dim=1)
 92    # substract the embeddings from the shifted embeddings, take the norm and
 93    # transform to affinities based on the delta distance
 94    affs = (2 * delta - torch.norm(embeddings.unsqueeze(1) - shifted, dim=2)) / (2 * delta)
 95    affs = torch.clamp(affs, min=0) ** 2
 96    return affs
 97
 98
 99class AffinitySideLoss(nn.Module):
100    def __init__(self, offset_ranges, n_samples, delta):
101        assert all(len(orange) == 2 for orange in offset_ranges)
102        super().__init__()
103        self.ndim = len(offset_ranges)
104        self.offset_ranges = offset_ranges
105        self.n_samples = n_samples
106        self.delta = delta
107
108    def __call__(
109        self,
110        input_,
111        target,
112        ignore_labels=None,
113        ignore_in_variance_term=None,
114        ignore_in_distance_term=None,
115    ):
116        assert input_.dim() == target.dim(), f"{input_.dim()}, {target.dim()}"
117        assert input_.shape[2:] == target.shape[2:]
118
119        # sample offsets
120        offsets = [[np.random.randint(orange[0], orange[1]) for orange in self.offset_ranges]
121                   for _ in range(self.n_samples)]
122
123        # we invert the affinities and the target affinities
124        # so that we get boundaries as foreground, which is benefitial for the dice loss.
125        # compute affinities from emebeddings
126        affs = 1. - embeddings_to_affinities(input_, offsets, self.delta)
127
128        # compute groundtruth affinities from target
129        target_affs = 1. - segmentation_to_affinities(target, offsets)
130        assert affs.shape == target_affs.shape, f"{affs.shape}, {target_affs.shape}"
131
132        # TODO implement masking the ignore labels
133        # compute the dice score between affinities and target affinities
134        loss = dice_score(affs, target_affs, invert=True)
135        return loss