A Multi-Task Cross-Task Learning Architecture for Ad-hoc Uncertainty Estimation in 3D Cardiac MRI Image Segmentation 2021, Oral

Prerequisites

Our code is compatible with python 3.7 or onward.

We depend on some python packages which need to be installed by the user:

• PyTorch
• tqdm
• SimpleITK
• sklearn
• numpy

Authors

• S. M. Kamrul Hasan ([sh3190@rit.edu])
• Cristian A. Linte

• [ Paper ]

• [ Video ]

• [ Code ]

• [ Dataset ]

Abstract

Medical image segmentation has significantly benefitted thanks to deep learning architectures. Furthermore, semi-supervised learning (SSL) has recently been a growing trend for improving a model's overall performance by leveraging abundant unlabeled data. Moreover, learning multiple tasks within the same model further improves model generalizability. To generate smoother and accurate segmentation masks from 3D cardiac MR images, we present a Multi-task Cross-task learning consistency approach to enforce the correlation between the pixel-level (segmentation) and the geometric-level (distance map) tasks. Our extensive experimentation with varied quantities of labeled data in the training sets justifies the effectiveness of our model for the segmentation of the left atrial cavity from Gadolinium-enhanced magnetic resonance (GE-MR) images. With the incorporation of uncertainty estimates to detect failures in the segmentation masks generated by CNNs, our study further showcases the potential of our model to flag low-quality segmentation from a given model.

Video

--> -->

Dataset

The model was trained and tested on the MICCAI STACOM 2018 Atrial Segmentation Challenge datasets featuring 100 3D gadolinium-enhanced MR imaging scans (GE-MRIs) and LA segmentation masks, with an isotropic resolution of 0.625×0.625×0.625mm3. The dimensions of the MRIs may vary depending on each patient, however, all MRIs contain exactly 88 slices in the z axis. All the images were normalized and resized to 112×112×80 before feeding them to the models. We split them into 80 scans for training and 20 scans for validation, and apply the same pre-processing methods. \

Method: U-NetPlus

our proposed MT-CTL model has two distinctive features. First, we combine four different decoders who share the same backbone encoder – V-Net [12]. The uncertainty map generated by the uncertainty decoder is used as the local guidance between the predicted segmentation mask and the mask generated by transforming the distance map. Second, we enforce the correlation between the pixel-level (segmentation) and the geometric-level (distance map) tasks for the generation of smoother and more accurate segmentation masks by introducing the cross-task loss function and include a guidance loss as an uncertainty estimation to smooth out the predicted segmentation mask.

Augmentation

Qualitative Results

Ablation Study

Citation

If you find our work useful in your research, please consider citing:

@article{hasan2021multi,
  title={A Multi-Task Cross-Task Learning Architecture for Ad-hoc Uncertainty Estimation in 3D Cardiac MRI Image Segmentation},
  author={Hasan, SM and Linte, Cristian A},
  journal={arXiv preprint arXiv:2109.07702},
  year={2021}
}

Acknowledgements

Research reported in this publication was supported by the National Institute of General Medical Sciences of the National Institutes of Health under Award No. R35GM128877 and by the Office of Advanced Cyber infrastructure of the National Science Foundation under Award No. 1808530. .

Note

Contact: S. M. Kamrul Hasan (smkamrulhasan.rit@gmail.com)