This repository contains the code implementation of the IPMI 2021 paper Spatially Varying Label Smoothing: Capturing Uncertainty from Expert Annotations
3D as paper:
class CELossWithSVLS(torch.nn.Module):
def __init__(self, classes=None, sigma=1):
super(CELossWithSVLS, self).__init__()
self.cls = torch.tensor(classes)
self.cls_idx = torch.arange(self.cls).reshape(1, self.cls).cuda()
self.svls_layer, self.svls_kernel = get_svls_filter_3d(sigma=sigma, channels=classes)
self.svls_kernel = self.svls_kernel.cuda()
def forward(self, inputs, labels):
with torch.no_grad():
oh_labels = (labels[...,None] == self.cls_idx).permute(0,4,1,2,3)
b, c, d, h, w = oh_labels.shape
x = oh_labels.view(b, c, d, h, w).repeat(1, 1, 1, 1, 1).float()
x = F.pad(x, (1,1,1,1,1,1), mode='replicate')
svls_labels = self.svls_layer(x)/self.svls_kernel.sum()
return (- svls_labels * F.log_softmax(inputs, dim=1)).sum(dim=1).mean()2D version for 2D image segmentation: you can also find an interactive Colab notebook on toy example for 2D SVLS implementation.
class CELossWithSVLS_2D(torch.nn.Module):
def __init__(self, classes=None, sigma=1, ksize=3):
super(CELossWithSVLS_2D, self).__init__()
self.cls = torch.tensor(classes)
self.svls_layer = get_svls_filter_2d(ksize=3, sigma=sigma, channels=self.cls)
def forward(self, inputs, labels):
oh_labels = F.one_hot(labels.to(torch.int64), num_classes = self.cls).contiguous().permute(0,3,1,2).float()
svls_labels = self.svls_layer(oh_labels)
return (- svls_labels * F.log_softmax(inputs, dim=1)).sum(dim=1).mean()
The implementation of the surface dice is adopted from this repository and rename the folder to "surface_distance" to avoid library importing issue in python.
Train command for SVLS on BraTS 2019
CUDA_VISIBLE_DEVICES=0,1 python main.py --batch_size 2 --data_root /vol/biomedic3/mi615/datasets/BraTS/MICCAI_BraTS_2019_Data_Training/HGG_LGG/ --train_option SVLS
Validation command for SVLS on BraTS 2019
CUDA_VISIBLE_DEVICES=0,1 python deploy.py --batch_size 2 --data_root /vol/biomedic3/mi615/datasets/BraTS/MICCAI_BraTS_2019_Data_Training/HGG_LGG/ --train_option SVLS
To run a demo on all training options and evaluation metrics with reliability diagram with BraTS 2019
CUDA_VISIBLE_DEVICES=0,1 python demo.py --batch_size 2 --data_root /vol/biomedic3/mi615/datasets/BraTS/MICCAI_BraTS_2019_Data_Training/HGG_LGG/ --train_option SVLS
The model architecture is adopted from this repository
If you use this code for your research, please cite our paper.
@article{islam2021spatially,
title={Spatially Varying Label Smoothing: Capturing Uncertainty from Expert Annotations},
author={Islam, Mobarakol and Glocker, Ben},
journal={arXiv preprint arXiv:2104.05788},
year={2021}
}