ExSwin-UNet: An Unbalanced Weighted UNet with Shifted Window and External Attentions for Fetal Brain MRI Image Segmentation

Abstract

Accurate fetal brain MRI image segmentation is essential for fetal disease diagnosis and treatment. While manual segmentation is laborious, time-consuming, and error-prone, automated segmentation is a challenging task owing to (1) the variations in shape and size of brain structures among patients, (2) the subtle changes caused by congenital diseases, and (3) the complicated anatomy of brain. It is critical to effectively capture the long-range dependencies and correlations among training samples to yield satisfactory results. Recently, some transformer-based models have been proposed and achieved good performance in segmentation tasks. However, the self-attention blocks embedded in transformers often neglect the latent relationships among different samples. Model may have biased results due to the unbalanced data distribution in the training dataset. We propose a novel unbalanced weighted Unet equipped with a new ExSwin transformer block to comprehensively address the above concerns by effectively capturing long-range dependencies and correlations among different samples. We design a deeper encoder to facilitate features extracting and preserving more semantic details. In addition, an adaptive weight adjusting method is implemented to dynamically adjust the loss weight of different classes to optimize learning direction and extract more features from under-learning classes. Extensive experiments on a FeTA dataset demonstrate the effectiveness of our model, achieving better results than state-of-the-art approaches.