A Lightweight U-Net for Medical Image Segmentation
Yuxing Zhao, Lan Lin · 2024
Nowadays, medical image segmentation is playing a pivotal role in assisting doctors in disease diagnosis and treatment planning. With the rising popularity of bed-side instruments and hand-held medical devices, there is an escalating need for higher segmentation accuracy and faster inference speed in mobile scenarios with constrained power consumption. Many existing model designs focus on improving the accurateness, ignoring considerations of model size and computational complexity. Therefore, it is necessary to design lightweight models for mobile terminals. In this paper, we propose a U-Net model based on a shifted Multilayer Perceptron (MLP), which can segment disease areas from medical images with greater efficiency. Firstly, a novel data preprocessing step is designed to competently suppress noise in images while preserve intricate edge details simultaneously. Subsequently, the convolution module of U-Net is optimized through the implementation of depth separable convolution and cheap linear transformations, generating high-quality feature maps with lower computational costs. Finally, a shift MLP suit for U-shaped architecture is designed to effectively capture global information and acquire spatial information in multiple directions through axial shift. This feature allows the model to proficiently capture additional local dependencies, thereby improving overall performance. The model is deployed on computers and mobile phones to evaluate the precision of segmentation and the speed of inference. Experimental results conducted on diverse datasets provide evidence that this approach significantly enhances the speed of segmentation while maintaining a high level of segmentation accuracy. The solution is expected to facilitate the advancement of research in mobile medical image segmentation.