IoU score for Muilticlass segmentation

I bit confuse how to calculate IoU score and dice for multi class segmentation, is my code valid for multi-class?

def IoU_score(inputs, targets, num_classes=23, smooth=1e-5):
    with torch.no_grad():
        #soft = nn.Softmax2d()
        inputs = F.softmax(inputs, dim=1) #convert into probabilites 0-1
        targets = F.one_hot(targets, num_classes = n_classes).permute(0,3,1,2).contiguous()#convert target into one-hot

        inputs = inputs.contiguous().view(-1)
        targets = targets.view(-1)

        intersection = (inputs * targets).sum()
        total = (inputs + targets).sum()
        union = total - intersection 

        IoU = (intersection + smooth)/(union + smooth)
        return IoU.item()

Based on @tom’s post the usual way would be to calculate a “class agnostic” loss and it seems he also liked an implementation to it, so you might compare your current approach to it. :slight_smile: