Soft labels to reflect uncertainty on boundary of ground truth annotation

[jwutsetro]

Hey all,

I am currently using monai to participate in the grand challenge for COVID segmentation. As a baseline model I use the DynUnet with parameters adapted from nn-Net. This works great and gave me a validation dice of around 0.7. To further improve the results I wanted to focus on handeling the noisy ground truth annotations. Since the ground truth annotations from this project are not really clean, I want to implement some form of 'soft labels'. By gaussian smoothing the masks, the probability drops below 1 on the borders of the lesions reflecting the uncertainty of the ground truth annotation.

I tried implementing this with monai building blocks, but I got stuck while using dice-loss since the one_hot function that is called in there expects binary masks input and doesn't work as expected for probabilistic masks. I now wrote my own 'soft_label_dice' that handles probabilistic labels in the case of only 2 class labels. I thought this might be an interesting feature for monai since multiple segmentations problems have uncertain ground truth boundaries.

I was wondering what you guys think of this soft labeling strategy. I know other methods exist for increasing noise robustness, but it seemed my model was being punished to hard for making mistakes during training on regions that are only coarsely annotated.
Below I added a snippet with my soft_label_dice function.

Kind regards,
Joris Wuts

def soft_label_dice(preds, label):
`

preds = torch.softmax(preds, 1)
# label is of shape (B*1*H[WD]) having float values ranging from 0-1 
label=torch.cat((label,(1-label)),1)

reduce_axis = list(range(2, len(preds.shape)))
nom=torch.sum(torch.pow((preds -label),2), dim=reduce_axis)

ground_o = torch.sum(preds, dim=reduce_axis)
pred_o = torch.sum(label, dim=reduce_axis)

denominator = ground_o + pred_o +0.00001
f: torch.Tensor = nom / denominator 
f = torch.mean(f)
return f`

[Nic-Ma]

Hi @jwutsetro ,

Thanks for your interest and experiments with MONAI in your challenge.
For this particular case, maybe you can execute GaussianSmooth and ThresholdIntensity transforms on your labels? Then CastToType transform to int number? Or you can write a simple transform to convert it to int label numbers.
@wyli , any other better suggestions?

Thanks.

[jwutsetro]

Hey,

thanks for the fast reply,
I didn’t make myself clear I think. I don’t want to just shrink the ground truth to a smaller size. I rather want the boundary regions to have uncertaity. For example 90% beloning to the lesion and 10% not beloning to it. So that the loss calculated on the boundary is smaller then the original loss, but still bigger than with a smaller binary mask. Here is a reference paper that almost aplies the same strategy ( but they add a fixed value rather than a smoothed one )
www.eng.biu.ac.il/goldbej/files/2019/03/isbi_eytan_2019.pdf

Kind regards,
Joris Wuts

[Nic-Ma]

Hi @wyli ,

Could you please help share some idea as you are building a challenge baseline?

Thanks.

[wyli]

sure, perhaps we'd better discuss this after this year's public challenge, otherwise I feel it's a bit unfair for the other participants

[jwutsetro]

No problem! I ended up using another strategy anyway. But the idea of soft labels is still interesting I think! In another project I am having simular troubles with uncertain boundaries.

I will wait for after the challenge!
Kind regards,
Joris Wuts

[wyli]

would be interesting to try some robust loss functions such as https://arxiv.org/abs/1701.03077?