Improving 3D Medical Image Segmentation at Boundary Regions using Local Self-attention and Global Volume Mixing

Document Type

Article

Publication Title

IEEE Transactions on Artificial Intelligence

Abstract

Volumetric medical image segmentation is a fundamental problem in medical image analysis where the objective is to accurately classify a given 3D volumetric medical image with voxel-level precision. In this work, we propose a novel hierarchical encoder-decoder-based framework that strives to explicitly capture the local and global dependencies for volumetric 3D medical image segmentation. The proposed framework exploits local volume-based self-attention to encode the local dependencies at high resolution and introduces a novel volumetric MLP-Mixer to capture the global dependencies at low-resolution feature representations, respectively. The proposed volumetric MLP-mixer learns better associations among volumetric feature representations. These explicit local and global feature representations contribute to better learning of the shape-boundary characteristics of the organs. Extensive experiments on three different datasets reveal that the proposed method achieves favorable performance compared to state-of-the-art approaches. On the challenging Synapse Multi-organ dataset, the proposed method achieves an absolute 3.82% gain over the state-of-the-art approaches in terms of HD95 evaluation metrics while a similar improvement pattern is exhibited in MSD Liver and Pancreas tumor datasets. We also provide a detailed comparison between recent architectural design choices in the 2D computer vision literature by adapting them for the problem of 3D medical image segmentation. Finally, our experiments on the ZebraFish 3D cell membrane dataset having limited training data demonstrate the superior transfer learning capabilities of the proposed vMixer model on the challenging 3D cell instance segmentation task, where accurate boundary prediction plays a vital role in distinguishing individual cell instances. Our source code is publicly available at https://github.com/Daniyanaj/vMixer.

First Page

1

Last Page

12

DOI

10.1109/TAI.2023.3346833

Publication Date

12-25-2023

Keywords

Attention, Computer architecture, Decoding, Image segmentation, Medical diagnostic imaging, Medical Image Segmentation, MLP-Mixer, Three-dimensional displays, Transfer Learning, Transfer learning, Transformers

Comments

OA version (pathway a) Accepted version

No embargo

When accepted for publication, set statement to accompany deposit (see policy)

Must link to publisher version with DOI

Publisher copyright and source must be acknowledged

Link to Full Text

Share

COinS