# Improving Prostate Gland Segmentation Using Transformer based Architectures

arXiv: https://arxiv.org/abs/2506.14844 · HTML: https://mlsanity.com/paper/2506.14844/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 44 references that are in the map.
- First posted on arXiv 16 Jun 2025.

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation](https://mlsanity.com/paper/1606.06650/) (2016)
- [UNETR: Transformers for 3D Medical Image Segmentation](https://mlsanity.com/paper/2103.10504/) (2021)
- [Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images](https://mlsanity.com/paper/2201.01266/) (2022)
- [Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations](https://mlsanity.com/paper/1707.03237/) (2017)
- [MONAI: An open-source framework for deep learning in healthcare](https://mlsanity.com/paper/2211.02701/) (2022)
- [TransBTS: Multimodal Brain Tumor Segmentation Using Transformer](https://mlsanity.com/paper/2103.04430/) (2021)
- [UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation](https://mlsanity.com/paper/2212.04497/) (2022)
- [Your ViT is Secretly an Image Segmentation Model](https://mlsanity.com/paper/2503.19108/) (2025)
- [A large-scale multicenter breast cancer DCE-MRI benchmark dataset with expert segmentations](https://mlsanity.com/paper/2406.13844/) (2024)
- [TFCNs: A CNN-Transformer Hybrid Network for Medical Image Segmentation](https://mlsanity.com/paper/2207.03450/) (2022)
- [BEFUnet: A Hybrid CNN-Transformer Architecture for Precise Medical Image Segmentation](https://mlsanity.com/paper/2402.08793/) (2024)
- [ConvFormer: Combining CNN and Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2211.08564/) (2022)
- [Fully Automated Organ Segmentation in Male Pelvic CT Images](https://mlsanity.com/paper/1805.12526/) (2018)
- [D-TrAttUnet: Toward Hybrid CNN-Transformer Architecture for Generic and Subtle Segmentation in Medical Images](https://mlsanity.com/paper/2405.04169/) (2024)

## Abstract (arXiv)

Inter reader variability and cross site domain shift challenge the automatic segmentation of prostate anatomy using T2 weighted MRI images. This study investigates whether transformer models can retain precision amid such heterogeneity. We compare the performance of UNETR and SwinUNETR in prostate gland segmentation against our previous 3D UNet model [1], based on 546 MRI (T2weighted) volumes annotated by two independent experts. Three training strategies were analyzed: single cohort dataset, 5 fold cross validated mixed cohort, and gland size based dataset. Hyperparameters were tuned by Optuna. The test set, from an independent population of readers, served as the evaluation endpoint (Dice Similarity Coefficient). In single reader training, SwinUNETR achieved an average dice score of 0.816 for Reader#1 and 0.860 for Reader#2, while UNETR scored 0.8 and 0.833 for Readers #1 and #2, respectively, compared to the baseline UNets 0.825 for Reader #1 and 0.851 for Reader #2. SwinUNETR had an average dice score of 0.8583 for Reader#1 and 0.867 for Reader#2 in cross-validated mixed training. For the gland size-based dataset, SwinUNETR achieved an average dice score of 0.902 for Reader#1 subset and 0.894 for Reader#2, using the five-fold mixed training strategy (Reader#1, n=53; Reader#2, n=87) at larger gland size-based subsets, where UNETR performed poorly. Our findings demonstrate that global and shifted-window self-attention effectively reduces label noise and class imbalance sensitivity, resulting in improvements in the Dice score over CNNs by up to five points while maintaining computational efficiency. This contributes to the high robustness of SwinUNETR for clinical deployment.