MeanIoU and GeneralizedDiceScore doesn't work properly with index tensors when per_class=True

[PRFina]

🐛 Bug

torchmetrics.segmentation.MeanIoU and torchmetrics.segmentation.GeneralizedDiceScore seem broken when passing:

• per_class=True
• num_classes=2
• preds or targets is an index tensor

I receive the following error RuntimeError: The size of tensor a (2) must match the size of tensor b (4) at non-singleton dimension 0. However, if preds and targets are both one-hot tensors everything works as expected.

To Reproduce

Code sample

import torch
from torchmetrics.segmentation import MeanIoU, GeneralizedDiceScore

dice = GeneralizedDiceScore(num_classes=2, per_class=True)
miou = MeanIoU(num_classes=2, per_class=True)

batch_size = 5

preds = torch.tensor([[0,0,1,1], [0,0,1,1],[0,0,1,1],[0,0,1,1]])
preds = torch.stack([preds]*batch_size, dim=0)

target = torch.tensor([[0,0,1,1], [0,0,1,1],[0,0,1,1],[0,0,1,1]])
target = torch.stack([target]*batch_size, dim=0)


miou.update(preds, target)
miou.compute()

However, if run a similar snippet changing only the shape of the input tensors, everything works as expected.

import torch
from torchmetrics.segmentation import MeanIoU, GeneralizedDiceScore

metric = MeanIoU(num_classes=2, per_class=True) # or GeneralizedDiceScore(num_classes=2, per_class=False)

batch_size = 5

preds = torch.tensor([[False,False,True,True], 
                      [False,False,True,True], 
                      [False,False,True,True], 
                      [False,False,True,True]]) # slice: 4x4 

preds = torch.stack([preds,~preds], dim=0) # mask: 2 x 4 x 4
preds = torch.stack([preds]*batch_size, dim=0) # batch: 5 x 2 x 4 x 4
preds = preds.int()

targets = torch.tensor([[False,False,True,True], 
                      [False,False,True,True], 
                      [False,False,True,True], 
                      [False,False,True,True]]) # slice: 4x4 

targets = torch.stack([targets,~targets], dim=0) # mask: 2 x 4 x 4
targets = torch.stack([targets]*batch_size, dim=0) # batch: 5 x 2 x 4 x 4
targets = targets.int()

metric.update(preds, targets)
metric.compute()

Expected behavior

As stated in the doc both functions support one-hot and index tensors: "An one-hot boolean tensor of shape (N, C, ...) with N being the number of samples and C the number of classes. Alternatively, an integer tensor of shape (N, ...) can be provided, where the integer values correspond to the class index. That format will be automatically converted to a one-hot tensor."

Environment

• TorchMetrics version (if build from source, add commit SHA): 1.4.1
• Python & PyTorch Version (e.g., 1.0): 3.11 and 2.4.1
• Any other relevant information such as OS (e.g., Linux): linux Ubuntu 22.04

Additional context

While trying to understand the problem from the error message I started a debugging session, this is what I discovered:

• The problem is not strictly related only to the metrics stateful classes but goes down to the relative functional version.
• I think that the root problem is this check to infer if input tensors are index tensors or not. Giving an index tensor that encodes classes with just 0 and 1 index, results in the check to evaluate to False, therefore the tensor is not reshaped into one-hot format.
• Since input tensors are still index tensors, the reduction here is wrong since it will reduce also one of the spatial dim, and will output a tensor with wrong shape that can't be added to self.score, raising the RuntimeError exception reported before.

I also noticed that this behavior has been changed in the master branch passing the index_format argument. This change will probably solve the issue, are there any plans to release soon?

[github-actions]

Hi! thanks for your contribution!, great first issue!

[SkafteNicki]

Hi @PRFina, thanks for raising this issue.
As you have already noted this has been fixed on master already by introducing the index_format argument because we at runtime cannot really distinguish between the two modes for particular inputs. We just had a release yesterday where this was not included, which was most likely a mistake on our part. We can do another release soon-ish.
@Borda just for reference we need PR #2572 in the next patch because it fixes this issue.
For now I recommend installing directly from the master branch:

pip install https://github.com/Lightning-AI/torchmetrics/archive/master.zip

Closing issue.