Lightning-AI / Lightning-AI/torchmetrics

Segmentation IOU compute Ignore some tagged values that don't need to be recorded (such as 255)

Open
#2,747 3 comments 5 reactions 1 assignee View on GitHub

@SkafteNicki is already working on this.

Since Apr 16, 2025.

enhancement
Dominant language
Python
Stars
2.5k
Forks
526
Avg merge
6d 11h
Merged PRs (30d)
5

Description

## 🚀 Feature

when we compute IOU
```python
import torch

_ = torch.manual_seed(0)
from torchmetrics.segmentation import MeanIoU

miou = MeanIoU(num_classes=3)
preds = torch.randint(0, 2, (5,))
target = torch.as_tensor((0, 1, 2, 0, 255)) # An index of 255 is a tag to be ignored.
miou(preds, target)
>>> This will result in an error
```

### Motivation

When I generate the sample pairs, the opposite mask (assuming 3 classes), but not all pixels in the entire mask should be classified into a particular class, so I set these pixels to 255. The pixel is then ignored in the loss calculation using `torch.nn.CrossEntropyLoss(ignore_index=255)`. However, the IOU calculation does not have this feature, which leads to errors in the IOU calculation, so I wondered if it could be made to support the ignore_index parameter as well, to ignore certain pixels.

### Pitch

```python
import torch
_ = torch.manual_seed(0)
from torchmetrics.segmentation import MeanIoU
miou = MeanIoU(num_classes=3, ignore_index=255) # support ignore_index param to ignore index 255
preds = torch.randint(0, 2, (5,))
target = torch.as_tensor((0, 1, 2, 0, 255)) # An index of 255 is a tag to be ignored.
miou(preds, target)
```

### Alternatives

https://github.com/Lightning-AI/torchmetrics/blob/62d9d32280ba365f1c2c14e0bd8a5adc959a1a6e/src/torchmetrics/functional/segmentation/mean_iou.py#L42

https://github.com/Lightning-AI/torchmetrics/blob/62d9d32280ba365f1c2c14e0bd8a5adc959a1a6e/src/torchmetrics/functional/segmentation/mean_iou.py#L52-L55

```python
def _mean_iou_update(
preds: Tensor,
target: Tensor,
num_classes: int,
include_background: bool = False,
input_format: Literal["one-hot", "index"] = "one-hot",
ignore_index=255

) -> Tuple[Tensor, Tensor]:
...

if input_format == "index":
preds = torch.nn.functional.one_hot(preds, num_classes=num_classes).movedim(-1, 1)
mask = torch.where(target == ignore_index) # Add removal of ignored labels
target[mask] = 0
target = torch.nn.functional.one_hot(target, num_classes=num_classes)
target[mask] = 0 # set ont-hot to zero-hot from ignored labels
target = target.movedim(-1, 1)
...
```
### Additional context

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.