Multi-class training

[clarkbab]

Hi there,

I'm running nnU-Net on an institutional CT head and neck dataset with multi-class labels that I'm adding as shown here: https://github.com/MIC-DKFZ/nnUNet/blob/master/documentation/dataset_format.md#what-do-training-cases-look-like.

When performing inference on the test set, the "Brain" prediction is missing for about a quarter of the samples. For the remaining samples, the Brain prediction looks reasonable. I'm surprised as this is usually an easy structure to segment, and thought the model might be learning to predict empty labels for some samples because of missing labels in the training dataset. However, the documentation seems to indicate that nnU-Net handles missing multi-class labels.

Do I need to specify which labels are present for a sample during multi-class training or is this handled implicitly by nnU-Net?

[sten2lu]

Hi @clarkbab,

I am not sure whether I correctly understand you but nnU-Net has as any other training based models issue if structures which are supposed to be annotated are not annotated always in the training set.

Missing annotations can be handled with the ignore label as documented here .

However to apply this you have to know which images are not properly annotated or where the annotations are missing inside the images.

Best regards,

Carsten

[clarkbab]

Hi @sten2lu, thanks for your reply!

To clarify, the issue is that I'm training on multi-class labels, for which a single sample (patient) may have missing classes (e.g. no Brain segmentation). The nnU-Net docs state:

Segmentations must share the same geometry with their corresponding images (same shape etc.). Segmentations are integer maps with each value representing a semantic class. The background must be 0. If there is no background, then do not use the label 0 for something else! Integer values of your semantic classes must be consecutive (0, 1, 2, 3, ...). Of course, not all labels have to be present in each training case. Segmentations are saved as {CASE_IDENTIFER}.{FILE_ENDING} .

Of course, not all labels have to be present in each training case. This implies that it is fine to train with some samples missing a Brain label, for example. However, the trained nnU-Net model fails to predict the Brain label for a number of test cases, which seems strange as the Brain is usually an easy prediction due to large volume. So I'm just double-checking that this is all correct?

The ignore label command is something different, this is for when parts of a label are missing, e.g. you only have half a Brain.

I'm also curious if any normalisation is performed by nnU-Net to account for volume differences between structures during multi-class training. For example, the Brain is the largest-volume structure by far, so perhaps it is being strongly down-weighted in the loss function?

Thanks, Brett

[sten2lu]

Hi @clarkbab, The documentation states that not all labels have to be present inside of each image. Given the following 3 classes:

0: background
1: brain
2: brain tumor

It is fine when class 2 (brain tumor) is not present for the scan of a patient where there is no tumor present. However, it is an issue in a case where there is a brain tumor present, and it is just not annotated correctly (so the tumor is mislabeled as background or brain).

As nnU-Net uses a mix of DICE and Cross Entropy Loss by default. This makes it focus well on both large and small structures. Generally, though, nnU-Net sometimes has issues with small structures.

I am not sure whether I understand your issue correctly. Based on what you wrote me, I would try to train the model without the training case (patient) with the missing class and see whether it solves the predictions get better.

Best regards, Carsten

[clarkbab]

Hi @sten2lu,

Thanks for your reply. Yes, your understanding is correct. In this instance I have whole CT images of the head and neck region, and each patient only has labels for a subset of the regions as these labels are generated from clinical treatment. So the Brain may be present in the CT image but there is no label for the Brain. Unfortunately, I can't use ignore_label as I'd need to know the spatial location of the missing labels, which we don't know without the Brain label. Also I can't train without the training cases that have missing labels as most (if not all) of the dataset is missing at least one label.

It would be great if we could specify per-patient which labels are missing. However, I'm not sure of the structure of your U-Net model predictions and it might be a problem depending on the model's output structure. nnU-Net stores labels using a single 3D volume with integer classes and using this structure, we need to know the spatial location of the class to be ignored - which is not possible here. However, it the U-Net models predict one-hot-encoded label maps and convert to a single 3D volume as postprocessing, then it may be that you could ignore a whole channel (e.g. the Brain channel) during training for a particular sample?

I'd be interested if you could share what the model prediction looks like before postprocessing, is it a 3D volume with integer classes per structure or a 4D one-hot-encoded label map with a channel per structure?

To get around this limitation I'm currently using TotalSegmentator to predict the missing labels but it doesn't cover all classes unfortunately.

Thanks, Brett