Surgical-DeSAM: decoupling SAM for instrument segmentation in robotic surgery

SAM [1] is the foundation model for prompt-based image segmentation and is trained on the largest segmentation dataset with over 1 billion high-quality masks. SAM forms a simple-designed transformer and composed of a heavyweight image encoder, a prompt encoder, and a lightweight mask decoder. The image encoder can directly extract image features from input images without the need for a backbone model, while its lightweight prompt encoder can dynamically transform any given prompt into an embedding vector in real-time. These embeddings are then processed by a decoder, generating precise segmentation masks. Prompts have various types, including points, boxes, text, or masks, which limit the SAM’s ability to be directly utilised for real-world applications like surgical instrument segmentation during surgery. It is unrealistic to provide a prompt for each frame of the surgical video.

DETR [3] is the transformer-based detector called DETR (DEtection TRansformer) for object detection. It consists of a CNN backbone, an encoder-decoder transformer and feed-forward networks (FFNs). The CNN backbone is the commonly used ResNet50 [4] which extract the feature (\(\in \Re ^{d\times H\times W}\)) representation from the input image (\(\in \Re ^{3\times H_0\times W_0}\)). The output of the backbone then passes to the transformer encoder with spatial positional encoding and produces object queries and encoder memory. The decoder receives the encoder outputs and predicts the class labels and bounding boxes with centre coordinates, height and width using FFNs.

DETR utilises ResNet50 as the backbone CNN to extract the feature representation. However, as vision-transformer-based networks are showing much better performance than CNN, we replace the backbone network with a recent transformer-based architecture of Swin-transformer [5] and from our Swin-DETR as presented in Fig. 1. The Swin-transformer introduces a shifted window-based hierarchical transformer to add greater efficiency in the self-attention computation. It is important to note that the output of the Swin-transformer can be directly fed to the DETR encoder, where there is an additional step to collapse the spatial dimension of the ResNet50 feature into one dimension to convert it into a sequence of input for the transformer. Overall, SWIN-DETR consists of a Swin-transformer to extract the image feature, which is then passed to the transformer encoder-decoder and FFNs to obtain the final object class predictions and corresponding bounding boxes. More specifically, ResNet5o requires to convert feature map of \(f_\textrm{resnet} \in \Re ^{d\times H\times W}\) into \(f \in \Re ^{d\times HW}\) by collapsing the spatial dimensions where Swin-transformer directly produces output feature map of \(f_\textrm{swin} \in \Re ^{d\times HW}\).

As the image encoder of the SAM and the DETR are performing similar feature extraction, we decouple SAM by removing the image encoder and feeding the DETR encoder output directly to the mask decoder. This facilitates the train end-to-end segmentation model using DETR predicted detection prompt and a decoupled SAM of prompt encoder and mask decoder only. During the training period, we utilise both ground-truths of the detection bounding boxes and segmentation masks to train both models end-to-end. To calculate losses, we adopted box loss \({\mathcal {L}}_\textrm{box}\) combining GIoU [6] and \(l_1\) losses for the detection task following DETR and dice coefficient similarity (DSC) loss \({\mathcal {L}}_\textrm{dsc}\) for the segmentation task. Therefore, total loss \(\textrm{Loss}_\textrm{total}\) can be formulated as: