Introduction: Medical image segmentation has significantly benefitted thanks to deep learning architectures. Furthermore, semi-supervised learning has been a growing trend for improving model performance by leveraging abundant unlabeled data. Moreover, learning multiple tasks within the same model further improves model generalizability. To generate smoother and more accurate segmentation masks from 3D cardiac MR images, we present a Multi-task Cross-task learning approach to enforce the correlation between pixel-level (segmentation) and geometric-level (distance map) tasks.
Methods: We propose a novel semi-supervised module leveraging adversarial learning and task-based consistency regularization for jointly learning multiple tasks in a single backbone module – uncertainty estimation, geometric shape generation, reconstruction, and cardiac structure segmentation. The network inputs a 3D image volume and outputs an uncertainty map, a distance map, and a segmentation map. The distance map is fed to a transformer network to produce a segmentation map that is further used to share the supervisory signal from the predicted segmentation map. To leverage the unlabeled data, the distance map is fed to a regularized adversarial discriminator network to distinguish the predicted distance map from the labeled data. The same encoder backbone also estimates the uncertainty of the predicted segmentation map via Monte Carlo sampling.
Results: We trained and tested our architecture on 100 3D left atrium MRI datasets (80: training; 10: validation; 10: testing). Our model outperforms all the other semi-supervised baseline models for 3D left atrium segmentation according to the Dice (91.8% vs. 89.5%), Jaccard (84.8% vs. 81.2%), and Hausdorff Distance (5.5 mm vs. 8.2 mm) metrics, while using only 20% labeled data.
Conclusion: We presented a framework to improve cardiac feature segmentation accuracy by incorporating spatial uncertainty maps. These uncertainty estimation tools will prompt users if and where a segmentation pipeline may under-perform, enabling their isolation from the segmentation pipeline for improved performance.