Self-Supervised Learning for Organs At Risk and Tumor Segmentation with Uncertainty Quantification

TL;DR

This paper demonstrates that self-supervised pre-training of transformers significantly improves organ and tumor segmentation accuracy while enabling uncertainty quantification, reducing annotation costs.

Contribution

It introduces MC-Swin-U, a transformer-based model with uncertainty quantification, showing improved segmentation performance over fully-supervised methods.

Findings

Self-supervised pre-training enhances segmentation accuracy.

Uncertainty quantification aids in reliable predictions.

Reduces the need for extensive annotated datasets.

Abstract

In this study, our goal is to show the impact of self-supervised pre-training of transformers for organ at risk (OAR) and tumor segmentation as compared to costly fully-supervised learning. The proposed algorithm is called Monte Carlo Transformer based U-Net (MC-Swin-U). Unlike many other available models, our approach presents uncertainty quantification with Monte Carlo dropout strategy while generating its voxel-wise prediction. We test and validate the proposed model on both public and one private datasets and evaluate the gross tumor volume (GTV) as well as nearby risky organs' boundaries. We show that self-supervised pre-training approach improves the segmentation scores significantly while providing additional benefits for avoiding large-scale annotation costs.