Pytorch: How to compute IoU (Jaccard Index) for semantic segmentation

Question

Can someone provide a toy example of how to compute IoU (intersection over union) for semantic segmentation in pytorch?

Answer 1

As of 2021, there's no need to implement your own IoU, as torchmetrics comes equipped with it - here's the link. It is named torchmetrics.JaccardIndex (previously torchmetrics.IoU) and calculates what you want. It works with PyTorch and PyTorch Lightning, also with distributed training.

From the documentation:

torchmetrics.JaccardIndex(num_classes, ignore_index=None, absent_score=0.0, threshold=0.5, multilabel=False, reduction='elementwise_mean', compute_on_step=None, **kwargs)

Computes Intersection over union, or Jaccard index calculation:

J(A,B) = \frac{|A\cap B|}{|A\cup B|}

Where: A and B are both tensors of the same size, containing integer class values. They may be subject to conversion from input data (see description below). Note that it is different from box IoU.

Works with binary, multiclass and multi-label data. Accepts probabilities from a model output or integer class values in prediction. Works with multi-dimensional preds and target.

Forward accepts

• preds (float or long tensor): (N, ...) or (N, C, ...) where C is the number of classes
• target (long tensor): (N, ...) If preds and target are the same shape and preds is a float tensor, we use the self.threshold argument to convert into integer labels. This is the case for binary and multi-label probabilities.

If preds has an extra dimension as in the case of multi-class scores we perform an argmax on dim=1.

Official example:

>>> from torchmetrics import JaccardIndex
>>> target = torch.randint(0, 2, (10, 25, 25))
>>> pred = torch.tensor(target)
>>> pred[2:5, 7:13, 9:15] = 1 - pred[2:5, 7:13, 9:15]
>>> jaccard = JaccardIndex(num_classes=2)
>>> jaccard(pred, target)
tensor(0.9660)