Distillation-SAM: Knowledge Distillation-Based Auto-Prompt Embedding Learning for Surgical Image Segmentation

Surgical image segmentation is vital for various stages of surgical procedures, from preoperative planning to real-time navigation and postoperative assessment. Despite advances in deep learning, current surgical image segmentation methods remain limited. They primarily target instrument segmentation and show poor generalizability across different surgical settings. While the Segment Anything Model (SAM) shows robust generalization capabilities in the segmentation of natural images, adapting SAM to surgical and medical images faces challenges because of its reliance on high-quality user-provided prompts and inherent lack of design for multi-class semantic segmentation. To address these limitations, we propose Distillation-SAM, an effective method that adapts SAM for accurate surgical image segmentation without user-provided prompts while freezing its encoder and decoder. Distillation-SAM introduces a trainable adapter branch that learns both sparse auto-prompt embeddings and enriched image features with dense auto-prompt embeddings, enabling the segmentation of surgical objects such as vessels, instruments, and tissues. We propose a direct knowledge distillation constraint for these auto-prompt embedding learnings by using embeddings derived from ground-truth masks as guidance. To enable multi-class semantic segmentation using SAM, we revise the mask score regression branch in SAM's decoder by incorporating a trainable Multilayer Perceptron to predict mask categories while keeping other parameters frozen. Our experiments in multiple surgical datasets, including IVIS, EndoVis2017, and Cholecseg8k, demonstrate that distillation-SAM outperforms existing methods in vessel, tissue, and instrument segmentation.