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

TL;DR

This paper introduces a hierarchical encoder-decoder framework for 3D medical image segmentation that explicitly captures local and global dependencies using local self-attention and a volumetric MLP-mixer, leading to improved boundary accuracy.

Contribution

The paper proposes a novel combination of local self-attention and volumetric MLP-mixer in a hierarchical framework for enhanced 3D segmentation performance.

Findings

Achieves 3.82% improvement in HD95 on Synapse dataset.

Demonstrates superior transfer learning on ZebraFish dataset.

Outperforms state-of-the-art methods across multiple datasets.

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.