Roc multiclass problem

Hey community,

I am trying to create a ROC (weighted and per class) in my multiclass segmentation model. But, in the code example of sklearn they binarize the y (before even train test split) to do this multiclass task. So my question is: how to do this binarize my model, which I have the label (y_true, a tensor Bx1xHxW) and the model prediction (y_pred, a tensor BxCxHxW).


If I understand this example correctly, label_binarize will create a one-hot encoded array for a multi-class target.
Would flattening the target tensor before passing it to label_binarize work?
This would create a flat tensor containing the one-hot-encoded targets for each sample in the batch.
roc_curve also uses the ravel() operation, so it would flatten the arrays anyway.
Let me know, if this would work.