EGE-UNet: an Efficient Group Enhanced UNet for skin lesion segmentation

Jiacheng Ruan,Mingye Xie,Jingsheng Gao,Ting Liu,Yuzhuo Fu
2023-07-17
Abstract:Transformer and its variants have been widely used for medical image segmentation. However, the large number of parameter and computational load of these models make them unsuitable for mobile health applications. To address this issue, we propose a more efficient approach, the Efficient Group Enhanced UNet (EGE-UNet). We incorporate a Group multi-axis Hadamard Product Attention module (GHPA) and a Group Aggregation Bridge module (GAB) in a lightweight manner. The GHPA groups input features and performs Hadamard Product Attention mechanism (HPA) on different axes to extract pathological information from diverse perspectives. The GAB effectively fuses multi-scale information by grouping low-level features, high-level features, and a mask generated by the decoder at each stage. Comprehensive experiments on the ISIC2017 and ISIC2018 datasets demonstrate that EGE-UNet outperforms existing state-of-the-art methods. In short, compared to the TransFuse, our model achieves superior segmentation performance while reducing parameter and computation costs by 494x and 160x, respectively. Moreover, to our best knowledge, this is the first model with a parameter count limited to just 50KB. Our code is available at <a class="link-external link-https" href="https://github.com/JCruan519/EGE-UNet" rel="external noopener nofollow">this https URL</a>.
Image and Video Processing,Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
The paper aims to address the challenges in the task of skin lesion segmentation, particularly the need for efficient, low-parameter, and low-computation load models in mobile healthcare applications. Specifically, the paper proposes the Efficient Group Enhanced UNet (EGE-UNet), a lightweight skin lesion segmentation model that achieves excellent segmentation performance while significantly reducing the number of parameters and computational cost. The authors point out that although existing methods based on Transformers and their variants perform well in medical image segmentation tasks, their large number of parameters and computational demands make them unsuitable for mobile health applications. Therefore, the paper proposes a new approach, EGE-UNet, which combines two key modules: the Group Multi-Axis Hadamard Product Attention module (GHPA) and the Group Aggregation Bridging module (GAB). Through these modules, EGE-UNet not only surpasses some existing advanced models in segmentation performance but also reduces the number of parameters and computational complexity by 494 times and 160 times, respectively. Additionally, to the best of the authors' knowledge, this is the first model to limit the number of parameters to approximately 50KB. Experimental results show that EGE-UNet outperforms current state-of-the-art methods on the ISIC2017 and ISIC2018 datasets.