NVTrans‐UNet: Neighborhood vision transformer based U‐Net for multi‐modal cardiac MR image segmentation

Abstract

With the rapid development of artificial intelligence and image processing technology, medical imaging technology has turned into a critical tool for clinical diagnosis and disease treatment. The extraction and segmentation of the regions of interest in cardiac images are crucial to the diagnosis of cardiovascular diseases. Due to the erratically diastolic and systolic cardiac, the boundaries of Magnetic Resonance (MR) images are quite fuzzy. Moreover, it is hard to provide complete information using a single modality due to the complex structure of the cardiac image. Furthermore, conventional CNN-based segmentation methods are weak in feature extraction. To overcome these challenges, we propose a multi-modal method for cardiac image segmentation, called NVTrans-UNet. Firstly, we employ the Neighborhood Vision Transformer (NVT) module, which takes advantage of Neighborhood Attention (NA) and inductive biases. It can better extract the local information of the cardiac image as well as reduce the computational cost. Secondly, we introduce a Multi-modal Gated Fusion (MGF) network, which can automatically adjust the contributions of different modal feature maps and make full use of multi-modal information. Thirdly, the bottleneck layer with Atrous Spatial Pyramid Pooling (ASPP) is proposed to expand the feature receptive field. Finally, the mixed loss is added to the cardiac image to focus the fuzzy boundary and realize accurate segmentation. We evaluated our model on MyoPS 2020 dataset. The Dice score of myocardial infarction (MI) was 0.642 ± 0.171, and the Dice score of myocardial infarction + edema (MI + ME) was 0.574 ± 0.110. Compared with the baseline, the MI increases by 11.2%, and the MI + ME increases by 12.5%. The results show the effectiveness of the proposed NVTrans-UNet in the segmentation of MI and ME.