RepMedSAM: Segment Anything in Medical Images with Lightweight CNN
Keywords: Segment Anything Model; Medical image; Lightweight; Knowledge distillation
TL;DR: We combine the lightweight CNN structure, RepViT, with SAM and apply it to medical image segmentation.
Abstract: Traditional deep learning segmentation models require designing network structures and loss functions specific to different tasks, followed by training dedicated models, which leads to a significant amount of repetitive work. The Segment Anything Model (SAM) provides a unified framework for handling segmentation tasks. However, the current SAM model is mainly applicable to natural images and may require substantial computational resources during inference, posing challenges for widespread clinical implementation. In this work, we utilize RepViT as the Image Encoder to develop a lightweight SAM structure. The training phase consists of two main parts: knowledge distillation and fine-tuning. During the inference phase, reparameterization is employed to optimize inference speed. The proposed method achieves an average DSC of 0.8688 and an average NSD of 0.8746 on the validation set, and it improves inference speed while increasing the number of parameters compared to the baseline.
Submission Number: 13
Loading