MSCT-UNET: multi-scale contrastive transformer within U-shaped network for medical image segmentation.

Journal: Physics in medicine and biology
Published Date:

Abstract

Automatic mutli-organ segmentation from anotomical images is essential in disease diagnosis and treatment planning. The U-shaped neural network with encoder-decoder has achieved great success in various segmentation tasks. However, a pure convolutional neural network (CNN) is not suitable for modeling long-range relations due to limited receptive fields, and a pure transformer is not good at capturing pixel-level features.We propose a new hybrid network named MSCT-UNET which fuses CNN features with transformer features at multi-scale and introduces multi-task contrastive learning to improve the segmentation performance. Specifically, the multi-scale low-level features extracted from CNN are further encoded through several transformers to build hierarchical global contexts. Then the cross fusion block fuses the low-level and high-level features in different directions. The deep-fused features are flowed back to the CNN and transformer branch for the next scale fusion. We introduce multi-task contrastive learning including a self-supervised global contrast learning and a supervised local contrast learning into MSCT-UNET. We also make the decoder stronger by using a transformer to better restore the segmentation map.Evaluation results on ACDC, Synapase and BraTS datasets demonstrate the improved performance over other methods compared. Ablation study results prove the effectiveness of our major innovations.The hybrid encoder of MSCT-UNET can capture multi-scale long-range dependencies and fine-grained detail features at the same time. The cross fusion block can fuse these features deeply. The multi-task contrastive learning of MSCT-UNET can strengthen the representation ability of the encoder and jointly optimize the networks. The source code is publicly available at:https://github.com/msctunet/MSCT_UNET.git.

Authors

  • Heran Xi
    School of Electronic Engineering, Heilongjiang University, Harbin, 150001, People's Republic of China.
  • Haoji Dong
    School of Computer Science and Technology, Heilongjiang University, Harbin, 150000, People's Republic of China.
  • Yue Sheng
    School of Computer Science and Technology, Heilongjiang University, Harbin, 150000, People's Republic of China.
  • Hui Cui
    Shanghai Center for Bioinformation Technology, Shanghai Academy of Science and Technology, 1278 Keyuan Road, Shanghai 201203, PR China; School of Life Science and Technology, ShanghaiTech University, Shanghai 201210, PR China.
  • Chengying Huang
    School of Computer Science and Technology, Heilongjiang University, Harbin, 150000, People's Republic of China.
  • Jinbao Li
    School of Computer Science and Technology, Heilongjiang University, Harbin 150080, China. jbli@hlju.edu.cn.
  • Jinghua Zhu
    School of Computer Science and Technology, Heilongjiang University, Harbin, 150000, People's Republic of China.