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}")

[Borda]

@digital-idiot's clarification on shapes was helpful. The API has changed significantly since 2022 — here's the current approach (v1.9.0).

For binary segmentation masks, use the task-specific classes:

from torchmetrics.classification import BinaryRecall, BinaryJaccardIndex

# preds: (N, H, W) — binary predictions (0 or 1)
# target: (N, H, W) — binary ground truth
recall = BinaryRecall()
jaccard = BinaryJaccardIndex()

# Both accept flat or spatial tensors — they get flattened internally
score_r = recall(preds, target)
score_j = jaccard(preds, target)

For per-image scores (which was unanswered in the original thread):

TorchMetrics computes globally across the batch by default. To get per-image metrics, you have two options:

Option A — loop over batch:

from torchmetrics.functional.classification import binary_jaccard_index

per_image = torch.stack([
    binary_jaccard_index(preds[i], target[i])
    for i in range(preds.shape[0])
])
mean_per_image = per_image.mean()

Option B — use the segmentation module which has native samplewise support:

from torchmetrics.segmentation import MeanIoU

miou = MeanIoU(
    num_classes=2,
    input_format="one-hot",   # (N, C, H, W) with C=2
    per_class=True,
)
# Stack binary masks into one-hot: (N, 2, H, W)
preds_oh = torch.stack([1 - preds, preds], dim=1)
target_oh = torch.stack([1 - target, target], dim=1)
result = miou(preds_oh, target_oh)  # per-class IoU

Docs: BinaryJaccardIndex | MeanIoU