How to handle multi-class segmentation?

Hi everyone, How does TorchIO handle multi class segmentation? For example, if I am training a 3D Semantic Segmentation Model with multiple classes, my ground truth label would be a 4D one-hot encoded label like this: num_classes * H * W * D. For example in this snippet:

one_subject = Subject(
        T1=Image('../BRATS2018_crop_renamed/LGG75_T1.nii.gz', torchio.INTENSITY),
        T2=Image('../BRATS2018_crop_renamed/LGG75_T2.nii.gz', torchio.INTENSITY),
        label=Image('../BRATS2018_crop_renamed/LGG75_Label.nii.gz', torchio.LABEL),
    )

What is the kind of mask does the label expect? Should it be a one-hot encoded mask or can it be a 3D mask itself, and tochio would handle the one-hot encoding part internally? Thank you.

[fepegar]

You can use torchio.LabelMap for your segmentations. They can be one-hot or not, as everything is 4D after #238.

[GFabien]

Sorry, I was wrong in my last message, transforms are ok.

Problems come when you use patch based approaches and you want to define some sampling strategies based on the values of your PV label maps. We wanted to have the same probabilty to draw a patch for every class, to do so you just have to take your different PV label maps, get the total volume occupied per class (the sum over a PV label map) and sum all your maps weighted by the inverse of the volumes. The current solution to do so would be to create a key in your sample with the probability map and call the WeightedSampler with this key.

Maybe an implementation like the one used for the masking_method used in some scaling transforms would be a good idea. But it would need to accept the whole sample and not only the Image. This would bring the flexibility to deal with multiple PV label maps and to address the possibly complicated ways of creating the probability map out of (PV) label map(s).