search

ellisdg / 3DUnetCNN

Pytorch 3D U-Net Convolution Neural Network (CNN) designed for medical image segmentation
MIT License
1.9k stars 653 forks source link

why ignore background label? #140

Closed amirgholami closed 6 years ago

amirgholami commented 6 years ago

Hi

Could you please comment on the design choice of not outputting probability map for background? I noticed that you are using a hard threshold of 0.5 to set a tissue as background. This may become problematic in case of multi class segmentation. So I was wondering why you made that choice. Was it done to address class imbalance, or saving memory?

Many Thanks,

ellisdg commented 6 years ago

I treat "background" as the absence of having something that I am interested in. If a pixel/voxel does not contain something I care about, it doesn't get any label assigned to it. Having no label assigned to it means that it is background. In my mind, it makes the most sense to train the model to only look for things that I am interested in.

Now if we want to get technical, background does actually have a label in the final label maps, and that label is 0. You could add 0 to the labels being trained on, and that might improve your performance. It might also decrease your performance. It would be an interesting experiment, but my guess is that it wouldn't really help.

As for using a hard threshold with multi class segmentation, I simply find the class that has maximum prediction for each voxel. If the maximum prediction for a given voxel is greater than the hard threshold (0.5 in this case) I assign the class label to that voxel. If the maximum prediction is less than the hard threshold, it does not get assigned a label (i.e. is labelled as background).

amirgholami commented 6 years ago

I think it depends on the task as in cases where there are multiple classes the thresholding method would require some tuning. But many thanks for your detailed answer, that explains the design choice

ellisdg commented 6 years ago

Could you give an example where the thresholding method would require some tuning? Thanks.

amirgholami commented 6 years ago

Sure. So let's consider the current brats application. If the threshold is set to a high value, say 0.8, then it would force the segmentation to predict tumors only when it is very confident (low entropy for logits). This could reduce false positives but on the other hand could lead to sub-optimal prediction (that is we will miss tumorous areas for which the network does not have a high probability for). Vice versa is true if we set this hard threshold to a small value. For the brats application, I have seen some tumorous areas, especially around skull boundary, that get missed if threshold is high.