Classification metrics for multiclass image segmentation with unlabeled pixels

[perlmutter]

I'm doing multi-class semantic segmentation of images with only partial labeling. For example, a pixel label of 0/1/2/3 indicates background/class1/class2/unlabeled. Unlabeled/background pixels far outnumber either foreground class. During training I use torch.nn.CrossEntropyLoss with ignore_index=3, and the predicted pixel probabilities will only have dimension 3, not 4.

I'd like to use torchmetrics to calculate pixel classification statistics (accuracy, recall, precision, jaccard, etc) but the metrics won't accept the shape of my predictions. For example,

import torch
import torchmetrics

d1, d2 = 10, 10
N = 2 
C = 3 # background + 2 foreground classes 

y = torch.randint(low=0,high=C+1,size=(N,d1,d2)) # Note: C+1 includes unlabeled index
pred = torch.rand(N,C,d1,d2)
pred /= torch.unsqueeze(pred.sum(axis=1),1) # normalize to probabilities

torchmetrics.functional.recall(pred,y,num_classes=C+1,average='none',mdmc_average='global')  # same result if num_classes=C

produces

ValueError: The highest label in 'target' should be smaller than the size of the 'C' dimension of 'preds'.

As a workaround I could add a 4th "dummy" probability, set to zero for all pixels. However in that case, I'd like to ignore both the unlabeled "class" (3) as well as the background class (0), since background recall is not interesting and would overwhelm the foreground statistics, but ignore_index only accepts a single value. Is there a more elegant way to compute foreground class statistics in this scenario?