What should be the shape to be used in Recall and JaccardIndex for binary masks?

[pini-kop]

Hi,
I'm working on segmentation task and at the moment my model's output is a batch binary masks. and its shape is (B, 1, H, W).
when calling Precision/Recall/F!, if i'm not specifying num_classes everything works. but if i'm trying to set num_classes=1 or 2,
I have to flatten the tensor. otherwise I get ValueError: The implied number of classes (from shape of inputs) does not match num_classes.

Also, shouldn't there be at least the possibility to calculate the metrics per image/mask and then take the mean over the batch rather then accumulate the confusion matrix of the entire batch/epoch and only then calculating the result?

[digital-idiot]

if i'm trying to set num_classes=1 or 2

num_classes cannot be 1 for obvious reasons. See the following example:

import torch
from torchmetrics import JaccardIndex, Recall

iou = JaccardIndex(num_classes=2)
recall = Recall(num_classes=2, mdmc_average='global')

target1 = torch.randint(0, 2, (10, 15, 16))
target2 = target1.unsqueeze(1)

pred1 = torch.rand(10, 1, 15, 16)
pred2 = torch.cat([pred1, (1 - pred1)], dim=1)

print(f"Prediction-1 Shape: {tuple(pred1.shape)}, Prediction-2 Shape: {tuple(pred2.shape)}")
print(f"Target-1 Shape: {tuple(target1.shape)}, Target-2 Shape: {tuple(target2.shape)}")

print(f"Pred-2 : Target-1 >>> IoU: {iou(pred2, target1):.3f}, Recall: {recall(pred2, target1):.3f}")
print(f"Pred-2 : Target-2 >>> IoU: {iou(pred2, target2):.3f}, Recall: {recall(pred2, target2):.3f}")

'''
Following will result in ValueError. When num_classes is specified pred.size(1) == num_classes must be satisfied
After removing num_classes from the metric initializations above folowwing will work

'''
# print(f"Pred-1 : Target-1 >>> IoU: {iou(pred1, target1):.3f}, Recall: {recall(pred1, target1):.3f}")
# print(f"Pred-1 : Target-2 >>> IoU: {iou(pred1, target2):.3f}, Recall: {recall(pred1, target2):.3f}")