A Flexible 2.5D Medical Image Segmentation Approach with In-Slice and Cross-Slice Attention (2405.00130v1)
Abstract: Deep learning has become the de facto method for medical image segmentation, with 3D segmentation models excelling in capturing complex 3D structures and 2D models offering high computational efficiency. However, segmenting 2.5D images, which have high in-plane but low through-plane resolution, is a relatively unexplored challenge. While applying 2D models to individual slices of a 2.5D image is feasible, it fails to capture the spatial relationships between slices. On the other hand, 3D models face challenges such as resolution inconsistencies in 2.5D images, along with computational complexity and susceptibility to overfitting when trained with limited data. In this context, 2.5D models, which capture inter-slice correlations using only 2D neural networks, emerge as a promising solution due to their reduced computational demand and simplicity in implementation. In this paper, we introduce CSA-Net, a flexible 2.5D segmentation model capable of processing 2.5D images with an arbitrary number of slices through an innovative Cross-Slice Attention (CSA) module. This module uses the cross-slice attention mechanism to effectively capture 3D spatial information by learning long-range dependencies between the center slice (for segmentation) and its neighboring slices. Moreover, CSA-Net utilizes the self-attention mechanism to understand correlations among pixels within the center slice. We evaluated CSA-Net on three 2.5D segmentation tasks: (1) multi-class brain MRI segmentation, (2) binary prostate MRI segmentation, and (3) multi-class prostate MRI segmentation. CSA-Net outperformed leading 2D and 2.5D segmentation methods across all three tasks, demonstrating its efficacy and superiority. Our code is publicly available at https://github.com/mirthAI/CSA-Net.
- N. Siddique, S. Paheding, C. P. Elkin, and V. Devabhaktuni, “U-net and its variants for medical image segmentation: A review of theory and applications,” Ieee Access, vol. 9, pp. 82 031–82 057, 2021.
- M. E. Gegundez-Arias, D. Marin-Santos, I. Perez-Borrero, and M. J. Vasallo-Vazquez, “A new deep learning method for blood vessel segmentation in retinal images based on convolutional kernels and modified u-net model,” Computer Methods and Programs in Biomedicine, vol. 205, p. 106081, 2021.
- Y. Al-Kofahi, A. Zaltsman, R. Graves, W. Marshall, and M. Rusu, “A deep learning-based algorithm for 2-d cell segmentation in microscopy images,” BMC bioinformatics, vol. 19, pp. 1–11, 2018.
- M. Kim and B.-D. Lee, “Automatic lung segmentation on chest x-rays using self-attention deep neural network,” Sensors, vol. 21, no. 2, p. 369, 2021.
- E. Gibson, F. Giganti, Y. Hu, E. Bonmati, S. Bandula, K. Gurusamy et al., “Automatic multi-organ segmentation on abdominal ct with dense v-networks,” IEEE transactions on medical imaging, vol. 37, no. 8, pp. 1822–1834, 2018.
- M. Havaei, A. Davy, D. Warde-Farley, A. Biard, A. Courville, Y. Bengio et al., “Brain tumor segmentation with deep neural networks,” Medical image analysis, vol. 35, pp. 18–31, 2017.
- Y. Xia, L. Xie, F. Liu, Z. Zhu, E. K. Fishman, and A. L. Yuille, “Bridging the gap between 2d and 3d organ segmentation with volumetric fusion net,” in Medical Image Computing and Computer Assisted Intervention–MICCAI 2018: 21st International Conference, Granada, Spain, September 16-20, 2018, Proceedings, Part IV 11. Springer, 2018, pp. 445–453.
- K. Hu, C. Liu, X. Yu, J. Zhang, Y. He, and H. Zhu, “A 2.5 d cancer segmentation for mri images based on u-net,” in 2018 5th International Conference on Information Science and Control Engineering (ICISCE). IEEE, 2018, pp. 6–10.
- H. Zhang, A. M. Valcarcel, R. Bakshi, R. Chu, F. Bagnato, R. T. Shinohara et al., “Multiple sclerosis lesion segmentation with tiramisu and 2.5 d stacked slices,” in Medical Image Computing and Computer Assisted Intervention–MICCAI 2019: 22nd International Conference, Shenzhen, China, October 13–17, 2019, Proceedings, Part III 22. Springer, 2019, pp. 338–346.
- L. Han, Y. Chen, J. Li, B. Zhong, Y. Lei, and M. Sun, “Liver segmentation with 2.5 d perpendicular unets,” Computers & Electrical Engineering, vol. 91, p. 107118, 2021.
- S. J. C. Soerensen, R. E. Fan, A. Seetharaman, L. Chen, W. Shao, I. Bhattacharya et al., “Deep learning improves speed and accuracy of prostate gland segmentations on magnetic resonance imaging for targeted biopsy,” The Journal of urology, vol. 206, no. 3, pp. 604–612, 2021.
- A. L. Y. Hung, H. Zheng, Q. Miao, S. S. Raman, D. Terzopoulos, and K. Sung, “Cat-net: A cross-slice attention transformer model for prostate zonal segmentation in mri,” IEEE transactions on medical imaging, vol. 42, no. 1, pp. 291–303, 2023.
- A. L. Y. Hung, H. Zheng, K. Zhao, X. Du, K. Pang, Q. Miao et al., “Csam: A 2.5 d cross-slice attention module for anisotropic volumetric medical image segmentation,” in Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, 2024, pp. 5923–5932.
- X. Wang, R. Girshick, A. Gupta, and K. He, “Non-local neural networks,” in Proceedings of the IEEE conference on computer vision and pattern recognition, 2018, pp. 7794–7803.
- A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner et al., “An image is worth 16x16 words: Transformers for image recognition at scale,” arXiv preprint arXiv:2010.11929, 2020.
- J. Chen, Y. Lu, Q. Yu, X. Luo, E. Adeli, Y. Wang et al., “Transunet: Transformers make strong encoders for medical image segmentation,” arXiv preprint arXiv:2102.04306, 2021.
- H.-Y. Zhou, J. Guo, Y. Zhang, X. Han, L. Yu, L. Wang et al., “nnformer: Volumetric medical image segmentation via a 3d transformer,” IEEE Transactions on Image Processing, 2023.
- Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang et al., “Swin transformer: Hierarchical vision transformer using shifted windows,” in Proceedings of the IEEE/CVF international conference on computer vision, 2021, pp. 10 012–10 022.
- H. Cao, Y. Wang, J. Chen, D. Jiang, X. Zhang, Q. Tian et al., “Swin-unet: Unet-like pure transformer for medical image segmentation,” in European conference on computer vision. Springer, 2022, pp. 205–218.
- A. Hatamizadeh, V. Nath, Y. Tang, D. Yang, H. R. Roth, and D. Xu, “Swin unetr: Swin transformers for semantic segmentation of brain tumors in mri images,” in International MICCAI Brainlesion Workshop. Springer, 2021, pp. 272–284.
- S. Zhang, B. Ren, Z. Yu, H. Yang, X. Han, X. Chen et al., “Tw-net: Transformer weighted network for neonatal brain mri segmentation,” IEEE Journal of Biomedical and Health Informatics, vol. 27, no. 2, pp. 1072–1083, 2022.
- B. Chen, Y. Liu, Z. Zhang, G. Lu, and A. W. K. Kong, “Transattunet: Multi-level attention-guided u-net with transformer for medical image segmentation,” IEEE Transactions on Emerging Topics in Computational Intelligence, 2023.
- S. Yang, X. Li, J. Mei, J. Chen, C. Xie, and Y. Zhou, “3d-transunet for brain metastases segmentation in the brats2023 challenge,” arXiv preprint arXiv:2403.15735, 2024.
- M. Ding, Z. Lin, C. H. Lee, C. H. Tan, and W. Huang, “A multi-scale channel attention network for prostate segmentation,” IEEE Transactions on Circuits and Systems II: Express Briefs, 2023.
- A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez et al., “Attention is all you need,” Advances in neural information processing systems, vol. 30, 2017.
- X. Song, H. Guo, X. Xu, H. Chao, S. Xu, B. Turkbey et al., “Cross-modal attention for mri and ultrasound volume registration,” in Medical Image Computing and Computer Assisted Intervention–MICCAI 2021: 24th International Conference, Strasbourg, France, September 27–October 1, 2021, Proceedings, Part IV 24. Springer, 2021, pp. 66–75.
- J. Zheng, H. Liu, Y. Feng, J. Xu, and L. Zhao, “Casf-net: Cross-attention and cross-scale fusion network for medical image segmentation,” Computer Methods and Programs in Biomedicine, vol. 229, p. 107307, 2023.
- Y. Zhou, H. Chen, Y. Li, X. Cao, S. Wang, and D. Shen, “Cross-model attention-guided tumor segmentation for 3d automated breast ultrasound (abus) images,” IEEE Journal of Biomedical and Health Informatics, vol. 26, no. 1, pp. 301–311, 2021.
- Y. Li, H. Li, and Y. Fan, “Acenet: Anatomical context-encoding network for neuroanatomy segmentation,” Medical image analysis, vol. 70, p. 101991, 2021.
- Q. Yu, L. Xie, Y. Wang, Y. Zhou, E. K. Fishman, and A. L. Yuille, “Recurrent saliency transformation network: Incorporating multi-stage visual cues for small organ segmentation,” in Proceedings of the IEEE conference on computer vision and pattern recognition, 2018, pp. 8280–8289.
- L. Wang, D. Nie, G. Li, É. Puybareau, J. Dolz, Q. Zhang et al., “Benchmark on automatic six-month-old infant brain segmentation algorithms: the iseg-2017 challenge,” IEEE transactions on medical imaging, vol. 38, no. 9, pp. 2219–2230, 2019.
- J. Duan, G. Bello, J. Schlemper, W. Bai, T. J. Dawes, C. Biffi et al., “Automatic 3d bi-ventricular segmentation of cardiac images by a shape-refined multi-task deep learning approach,” IEEE transactions on medical imaging, vol. 38, no. 9, pp. 2151–2164, 2019.
- X. Yang, L. Yu, S. Li, H. Wen, D. Luo, C. Bian et al., “Towards automated semantic segmentation in prenatal volumetric ultrasound,” IEEE transactions on medical imaging, vol. 38, no. 1, pp. 180–193, 2018.
- G. Litjens, R. Toth, W. Van De Ven, C. Hoeks, S. Kerkstra, B. Van Ginneken et al., “Evaluation of prostate segmentation algorithms for mri: the promise12 challenge,” Medical image analysis, vol. 18, no. 2, pp. 359–373, 2014.
- S. G. Armato III, H. Huisman, K. Drukker, L. Hadjiiski, J. S. Kirby, N. Petrick et al., “Prostatex challenges for computerized classification of prostate lesions from multiparametric magnetic resonance images,” Journal of Medical Imaging, vol. 5, no. 4, pp. 044 501–044 501, 2018.
- O. Ronneberger, P. Fischer, and T. Brox, “U-net: Convolutional networks for biomedical image segmentation,” in Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. Springer, 2015, pp. 234–241.
- S. Vesal, I. Gayo, I. Bhattacharya, S. Natarajan, L. S. Marks et al., “Domain generalization for prostate segmentation in transrectal ultrasound images: A multi-center study,” Medical image analysis, vol. 82, p. 102620, 2022.
- C. Hsu, C. Chang, T. W. Chen, H. Tsai, S. Ma, and W. Wang, “Brain tumor segmentation (brats) challenge short paper: Improving three-dimensional brain tumor segmentation using segresnet and hybrid boundary-dice loss,” in International MICCAI Brainlesion Workshop. Springer, 2021, pp. 334–344.
- C. Valdes, P. Nataraj, K. Kisilewicz, A. Simenson, G. Leon, Kang et al., “Impact of nutritional status on total brain tissue volumes in preterm infants,” Children, vol. 11, no. 1, p. 121, 2024.
- J. Romberg, M. Wilke, C. Allgaier, T. Nägele, C. Engel, Poets et al., “Mri-based brain volumes of preterm infants at term: a systematic review and meta-analysis,” Archives of Disease in Childhood-Fetal and Neonatal Edition, vol. 107, no. 5, pp. 520–526, 2022.
- A. Seetharaman, I. Bhattacharya, L. C. Chen, C. A. Kunder, W. Shao, S. J. Soerensen et al., “Automated detection of aggressive and indolent prostate cancer on magnetic resonance imaging,” Medical Physics, vol. 48, no. 6, pp. 2960–2972, 2021.
- I. Bhattacharya, A. Seetharaman, C. Kunder, W. Shao, L. C. Chen, S. J. Soerensen et al., “Selective identification and localization of indolent and aggressive prostate cancers via corrsignia: an mri-pathology correlation and deep learning framework,” Medical image analysis, vol. 75, p. 102288, 2022.