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

TL;DR

This paper introduces a multi-task learning framework that combines segmentation and geometric tasks with uncertainty estimation to improve 3D cardiac MRI segmentation accuracy and reliability, especially with limited labeled data.

Contribution

It proposes a novel cross-task consistency approach for joint segmentation and distance map learning, incorporating uncertainty estimates to detect segmentation failures.

Findings

Enhanced segmentation accuracy with limited labeled data

Effective failure detection via uncertainty estimates

Improved model generalizability through multi-task learning

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.