Repository: Lightning-AI/torchmetrics

Issue:
Incorrect results with ignore_index in IoU
## 🐛 Bug

The ignore_index option in IoU does not behave as it should

### To Reproduce

Run the above code

#### Code sample

```python
import torch
import torchmetrics

iou = torchmetrics.IoU(3, ignore_index=0)
print(iou(torch.tensor([[1,2,1,1]]), torch.tensor([[2,1,1,0]])))
print(iou(torch.tensor([[1,2,1,0]]), torch.tensor([[2,1,1,0]])))
> tensor(0.1250)
tensor(0.1667)
```


### Expected behavior

Predicting a '0' here should not modify the score, both should give 0.1667

### Environment

 - PyTorch Version (e.g., 1.0): 1.8.1
 - OS (e.g., Linux): Linux & Windows
 - How you installed PyTorch (`conda`, `pip`, source): both
 - Build command you used (if compiling from source):
 - Python version: 3.8
 - CUDA/cuDNN version: 
 - GPU models and configuration:
 - Any other relevant information: 

### Additional context

The ignore_index is implemented like [this](https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/functional/classification/iou.py#L39):
```
if ignore_index is not None and 0 <= ignore_index < num_classes:
    scores = torch.cat([
        scores[:ignore_index],
        scores[ignore_index + 1:],
    ])
```
This is incorrect as it does not provide the right union value. Instead, the confmat should be stripped of its ignore_index column and line.

This can be done with this line provided it is put before the computation of the intersection and union:
`confmat = confmat[torch.arange(confmat.size(0)) != ignore_index].T[[ignore_index != torch.arange(confmat.size(0))]].T`



Discussion:
Hi! thanks for your contribution!, great first issue!
Hi @CSautier,
Thanks for raising this issue. Would you be up for sending a PR? :]

Resolving patch:
diff --git a/torchmetrics/functional/classification/iou.py b/torchmetrics/functional/classification/iou.py
--- a/torchmetrics/functional/classification/iou.py
+++ b/torchmetrics/functional/classification/iou.py
@@ -28,6 +28,11 @@ def _iou_from_confmat(
     absent_score: float = 0.0,
     reduction: str = 'elementwise_mean',
 ) -> Tensor:
+
+    # Remove the ignored class index from the scores.
+    if ignore_index is not None and 0 <= ignore_index < num_classes:
+        confmat[ignore_index] = 0.
+
     intersection = torch.diag(confmat)
     union = confmat.sum(0) + confmat.sum(1) - intersection
 
@@ -35,12 +40,12 @@ def _iou_from_confmat(
     scores = intersection.float() / union.float()
     scores[union == 0] = absent_score
 
-    # Remove the ignored class index from the scores.
     if ignore_index is not None and 0 <= ignore_index < num_classes:
         scores = torch.cat([
             scores[:ignore_index],
             scores[ignore_index + 1:],
         ])
+
     return reduce(scores, reduction=reduction)
 
 


Test patch:
diff --git a/tests/classification/test_iou.py b/tests/classification/test_iou.py
--- a/tests/classification/test_iou.py
+++ b/tests/classification/test_iou.py
@@ -153,7 +153,7 @@ def test_iou_differentiability(self, reduction, preds, target, sk_metric, num_cl
     pytest.param(False, 'none', 0, Tensor([1, 1])),
     pytest.param(True, 'none', None, Tensor([0.5, 0.5, 0.5])),
     pytest.param(True, 'elementwise_mean', None, Tensor([0.5])),
-    pytest.param(True, 'none', 0, Tensor([0.5, 0.5])),
+    pytest.param(True, 'none', 0, Tensor([2 / 3, 1 / 2])),
 ])
 def test_iou(half_ones, reduction, ignore_index, expected):
     preds = (torch.arange(120) % 3).view(-1, 1)
@@ -226,7 +226,7 @@ def test_iou_absent_score(pred, target, ignore_index, absent_score, num_classes,
         # Ignoring a valid index drops only that index from the result.
         pytest.param([0, 1, 1, 2, 2], [0, 1, 2, 2, 2], 0, 3, 'none', [1 / 2, 2 / 3]),
         pytest.param([0, 1, 1, 2, 2], [0, 1, 2, 2, 2], 1, 3, 'none', [1, 2 / 3]),
-        pytest.param([0, 1, 1, 2, 2], [0, 1, 2, 2, 2], 2, 3, 'none', [1, 1 / 2]),
+        pytest.param([0, 1, 1, 2, 2], [0, 1, 2, 2, 2], 2, 3, 'none', [1, 1]),
         # When reducing to mean or sum, the ignored index does not contribute to the output.
         pytest.param([0, 1, 1, 2, 2], [0, 1, 2, 2, 2], 0, 3, 'elementwise_mean', [7 / 12]),
         pytest.param([0, 1, 1, 2, 2], [0, 1, 2, 2, 2], 0, 3, 'sum', [7 / 6]),