SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zijian Zhou<sup>1</sup> Oluwatosin Alabi<sup>2</sup> Meng Wei<sup>2</sup> Tom Vercauteren<sup>2</sup> Miaojing Shi<sup>3*</sup>
$^{1}$ Department of Informatics, King's College London
$^{2}$ School of Biomedical Engineering & Imaging Sciences, King's College London
<sup>3</sup>College of Electronic and Information Engineering, Tongji University
{first_name}. {last_name}@kcl.ac.uk; mshi@tongji.edu.cn
In this paper, we propose a novel text promptable surgical instrument segmentation approach to overcome challenges associated with diversity and differentiation of surgical instruments in minimally invasive surgeries. We redefine the task as text promptable, thereby enabling a more nuanced comprehension of surgical instruments and adaptability to new instrument types. Inspired by recent advancements in vision-language models, we leverage pretrained image and text encoders as our model backbone and design a text promptable mask decoder consisting of attention- and convolution-based prompting schemes for surgical instrument segmentation prediction. Our model leverages multiple text prompts for each surgical instrument through a new mixture of prompts mechanism, resulting in enhanced segmentation performance. Additionally, we introduce a hard instrument area reinforcement module to improve image feature comprehension and segmentation precision. Extensive experiments on several surgical instrument segmentation datasets demonstrate our model's superior performance and promising generalization capability. To our knowledge, this is the first implementation of a promptable approach to surgical instrument segmentation, offering significant potential for practical application in the field of robotic-assisted surgery. Code is available at https://github.com/franciszzj/TP-SIS.
Minimally invasive surgeries (MIS) have gained widespread attention in various surgical disciplines [16, 48] due to their benefits over traditional open surgery, such as reduced patient discomfort and faster recovery times. Nonetheless, the restricted field of view and indirect vision through endoscopic cameras for existing MIS procedures presents considerable obstacles, making the development of robot-assisted MIS increasingly crucial. Current robot-assisted surgery has to operate under direct control of the surgeon. Enhancing the automatic comprehension of the surgical process through precise instrument segmentation is seen as an essential building block to introduce automation and facilitate robot-surgeon interaction. Despite this recognised importance, existing automatic surgical instrument segmentation faces significant challenges.
First, with fast-paced advances in MIS, there is a surge in the variety of surgical instruments from different vendors. This is however compounded with the lack of a comprehensive and large-scale dataset dedicated to the learning of surgical instrument segmentation. Current methods [43, 20, 12, 55, 5] do not adequately adapt to the ever-changing set of surgical instruments, necessitating re-labelling and re-training of models with the introduction of each new instrument. This has notably hampered the practical application of surgical instrument segmentation within the MIS field.
Second, the segmentation methods face difficulty in distinguishing between different categories of instruments [5], which often have similar appearances. The subtle visual difference across instruments and the often difficult imaging conditions in the confines of the surgical cavity make it harder to differentiate between instrument categories, causing poor segmentation performance.
Overcoming these challenges requires learning-based approaches that are more flexible and robust than the current state of the art. Recent progress in pre-trained vision-language models [33, 36, 47] offer novel opportunities for our research questions as demonstrated by results obtained in diverse downstream computer vision tasks [10, 35, 44, 30, 29, 46, 9]. On the one hand, these vision language models, trained on abundant data, offer robust features which can help compensate for the scarcity of surgical instrument segmentation data. On the other hand, the capability of these models to output aligned image and text features allows for text to supplement information for surgical instrument segmentation, thus simplifying the differentiation between various instrument types.
To address the first challenge, we observe that existing supervised instrument segmentation models [12, 55, 5, 3] are trained on predefined categories, thereby limiting their semantic understanding of instruments outside of the narrow training set. Drawing inspiration from recent advancements in text promptable image segmentation [29, 46] with vision-language models [36], we redefine the task as a text promptable surgical instrument segmentation, thereby enhancing generalization and adaptability to an ever-growing array of new surgical instruments. Figure 1 contrasts the predefined category method with our text-promptable approach. Our model adapts the image and text encoders from a pretrained vision-launguage model, CLIP [36], to extract features from both the surgical image and textural prompt; text promptable mask decoder consisting of attention-based and convolution-based prompting schemes is specifically designed to prompt the surgical instrument in the image from coarse to fine.
Varying text descriptions for each surgical instrument can yield distinct segmentation results. To tackle it, we propose a mixture of prompts mechanism, inspired by mixture of experts (MoE) [31, 41, 37], to leverage diverse prompts. This mechanism inputs multiple source prompts to the model and fuses the resulting predictions via weights generated from a visual-textural gating network. The final segmentation result hence benefits from various prompt information.
To address the second challenge, we aim to enhance the model's capacity to obtain improved image features, thereby reinforcing its ability to segment various surgical instrument categories and delineate precise edge details. We propose a hard instrument area reinforcement module intertwined with the popular image reconstruction approach, masked autoencoder (MAE) [15]. This module performs hard area mining at the site of segmentation error; by masking out the hard area, the module injects an auxiliary image reconstruction to the segmentation model to enhance the understanding of difficult areas in the image features.
To the best of our knowledge, we introduce the first text promptable surgical instrument segmentation approach. Our extensive experiments demonstrate state-of-the-art performance on several surgical instrument segmentation datasets [1, 2, 39, 17]. Additionally, we evaluate the generalization capability of our model through cross-validation between the two datasets, showcasing promising results and emphasizing the significant application potential of our method.
Surgical instrument segmentation. Most of the current works on surgical instrument segmentation are traditional vision-based models. For instance, TernausNet [43] trains an optimized U-Net [38] model for the segmentation of a restricted variety of surgical instruments. ISINet [12] proposes an instance-based segmentation method that utilizes a temporal consistency module to identify instrument candidates. MF-TapNet [20] utilizes motion flow information through an attention pyramid network and is trained independently for binary, part, and instrument type segmentation tasks. While most approaches rely on convolutional networks, a few recent works [55, 3] have investigated the use of transformer-based methods. For example, MATIS [3] utilizes pixel-wise attention, masked attention modules targeting for instrument areas, and video transformers for temporal information handling. Although the segmentation accuracy has been significantly improved by incorporating various architectures and modules, existing works are limited by their reliance on categorical understanding from the annotated training data; when new categories are introduced, they will have to be re-trained, which requires demanding labelling efforts and computational resources.
Vision based Model
Conventional vision-based surgical instrument segmentation method
Ground Truth
Figure 1: Left: input image and its ground truth. Top Right: conventional vision-based surgical instrument segmentation model with predefined categories. Bottom Right: ours leveraging vision-language model for text promptable surgical instrument segmentation.
Our text promptable surgical instrument segmentation method
Vision-language models. Large pretrained models like CLIP [36] has emerged as the most promising approaches for achieving state of the art results on downstream computer vision tasks. CLIP [36] is a large pretrained vision-language model [45] that aligns text with image modalities in the embedding space to solve downstream tasks like zero-shot image classification [36, 27]. It has been successfully applied to various vision tasks, including video captioning [44], video-text retrieval [30, 10], image generation [35], relation learning [57] and text promptable segmentation [46, 29].
Text promptable segmentation. Text promptable segmentation, also conceptualized as referring expression segmentation in the literature, involves the utilization of natural language expressions as prompts for image segmentation, deviating from the traditional reliance solely on class label annotations for training images [18]. Early works use convolution neural networks (CNNs) and recurrent neural networks (RNNs) to extract visual and textual features which are fused via concatenation and used to perform segmentation [18, 25]. Attention mechanisms are also introduced for exploiting relations between visual and textual features [18, 25, 42, 52]. More recent approaches utilize transformers to perform visual and textual feature fusion either in the encoder [11, 49, 51, 22, 54] or decoder [46, 29, 7].
Large pretrained vision-language models have also been utilized for this task [46, 29, 53, 28, 24]. For instance, CRIS [46] uses the CLIP image and text encoders for multimodal knowledge transfer and a transformer decoder for multimodal fusion and segmentation. CLIPSeg [29] extends prompt-based segmentation beyond text prompts to include image prompts processed by the CLIP image encoder. A recent concurrent work, SAM [24], introduces a transformer-based model capable of generating object masks using either image or text prompts. However, existing promptable segmentation methods are not designed to solve highly specific tasks like surgical instrument segmentation.
In this section, we begin by defining the task (Section 3.1), then introduce the image and text encoders inherited from a pretrained vision-language model (CLIP [36]), serving as the feature extractors of our method (Section 3.2). Next, we develop a text promptable mask decoder responsible for decoding the visual features into instrument segments with the help of textual features (Section 3.3). Finally, we elaborate on two key modules: the mixture of prompts module (Section 3.4), designed to enhance segmentation performance through multiple prompts; and the hard instrument area reinforcement module (Section 3.5), focused on improving instrument segmentation on difficult parts.
Given an image  $I \in \mathbb{R}^{H \times W \times 3}$  consisting of surgical instruments for an endoscopic operation, as well as a text description  $T$  describing the name, appearance and function of a given instrument, the objective is to obtain the surgical instrument mask  $M \in \{0,1\}^{H \times W}$  from  $I$ , prompted by the instrument description  $T$ . In cases when multiple instruments need to be segmented, we input the model with
Figure 2: An overview of our method. Our method comprises four key modules: 1) image and text encoders derived from the pretrained vision-language model to obtain visual and textual features; 2) a text promptable mask decoder consisting of attention- and convolution-based prompting schemes for predicting the score map from image features through text prompts; 3) a mixture of prompts mechanism that utilizes a visual-textual gating network to produce pixel-wise weights for merging different score maps; 4) a hard instrument area reinforcement module to reinforce image representation learning specifically on hard-predicted area.
multiple text descriptions corresponding to different categories, and the goal becomes producing the segmentation result  $M \in \{0, \dots, C\}^{H \times W}$ , where  $C$  represents the number of instrument categories.
3.2 Image and Text Encoders
To accomplish text promptable surgical instrument segmentation, we first encode the input endoscopic image  $I$  and instrument description  $T$ , obtaining the visual feature  $F_{I}$  and textual feature  $F_{T}$ , respectively. Importantly, to enforce the semantic congruity between  $F_{I}$  and  $F_{T}$ , we leverage the strength of the image and text encoders from the popular vision-language model, CLIP [36], pretrained on natural images.
Image encoder. Given the input image  $I$ , we utilize and fine tune the image encoder from CLIP [36], which is a vision transformer (ViT) [8], for feature extraction. The ViT-based image encoder comprises 12 layers. We use  $F_{I}^{l}\in \mathbb{R}^{N\times D}$  to denote its  $l$ -th layer output, where  $l$  ranges from 1 to 12,  $N$  is the number of visual tokens, and  $D$  is the dimension of visual features in ViT. Supposing the patch size of the input image is  $p$ , then  $N = \frac{H}{p}\times \frac{W}{p}$ .
Multi-scale feature augmentation. Extracting multi-scale features is important for semantic segmentation in order to segment objects in different scales. Whilst the image encoder in CLIP is not originally designed for multiple scales, we equip it with this ability. Specifically, simulating the multi-scale extraction from CNN-based architecture, we take the output features from the 4-th, 8-th, and 12-th layers of the ViT, i.e.  $F_{I}^{4}, F_{I}^{8}, F_{I}^{12}$ , representing a wide coverage of level-of-detail for the image; and fuse these features by a feature pyramid network (FPN) [26]. Notice in order to feed these features into FPN, we need to first reshape them into the size of  $\frac{H}{p} \times \frac{W}{p} \times D$ , then we up-sample  $F_{I}^{4}$  and down-sample  $F_{I}^{12}$  both by a factor of 2. Finally, the output of the FPN is reshaped to size of  $N \times D$  to obtain the multi-scale image feature  $F_{I} \in \mathbb{R}^{N \times D}$ .
Text encoder. Given the input description  $T$ , we utilize and keep frozen the text encoder from CLIP, which is a transformer encoder [45], for feature extraction. Within the text encoder,  $T$  is tokenized with the inclusion of a [CLS] token at its beginning, representing the global information of the description. Following the common practice in [46, 29], we use the corresponding feature from the [CLS] token as the global textual feature for the description. We denote it by  $F_{T} \in \mathbb{R}^{1 \times D}$ , where  $D$  is dimension of the textual feature, equivalent to that of the visual feature.
3.3 Text Promptable Mask Decoder
After obtaining  $F_{I}$  and  $F_{T}$ , our goal is to decode from  $F_{I}$  a score map  $S \in \mathbb{R}^{H \times W}$  with the help of  $F_{T}$ . Each pixel of  $S$  signifies the probability of this pixel belonging to the instrument class described in  $F_{T}$ . By thresholding  $S$ , it can be transformed into the desired mask  $M$ . To achieve this goal, we present two prompting schemes: attention-based prompting and convolution-based prompting, specified below. For a detailed structure figure of the text promptable mask decoder, please refer to the supplementary material.
Attention-based prompting. Inspired by [45], we employ the attention mechanisms to attend the instrument area in  $F_{I}$  using  $F_{T}$ . First, we compute the self-attention (SA) [45] within  $F_{I}$  to promote the foreground area in the endoscopic image. Layer normalization (LN) [4] is applied to  $F_{I}$  beforehand while skip connection [14] is applied afterwards. We write out this process:  $F_{I_S} = \mathrm{SA}(\mathrm{LN}(F_I)) + F_I$ . Next, we compute the cross-attention (CA) [45] between  $F_{I_S}$  and  $F_{T}$  to localize the instrument area in  $F_{I_S}$  according to  $F_{T}$ . Similarly, this process is written as:  $F_{I_C} = \mathrm{CA}(\mathrm{LN}(F_{I_S}), F_T) + F_{I_S}$ . Subsequently, we apply the feed-forward network, which is basically a fully connected layer to further polishing  $F_{I_C}$ ,  $F_{I_F} = \mathrm{FFN}(\mathrm{LN}(F_{I_C})) + F_{I_C}$ . SA, CA and FFN are organized in one decoding block, we devise three such blocks sequentially to obtain the attention-based prompted feature  $F_{I_A} \in \mathbb{R}^{N \times D}$ .
Convolution-based prompting. We design a novel convolution-based prompting scheme to further refine  $F_{I_A}$  by using  $F_T$  to convolve it locally. To do this,  $F_{I_A} \in \mathbb{R}^{N \times D}$  needs to be firstly reshaped to restore its spatial dimensions, i.e.  $\tilde{F}_{I_A} \in \mathbb{R}^{\frac{H}{p} \times \frac{W}{p} \times D}$ . Our idea is to transform  $F_T$  into the convolution parameters, i.e. kernel weights and bias, and use these parameters to convolve  $\tilde{F}_{I_A}$ , such that the instrument-related positions will be promoted. The transformation is done via a fully connected (FC) layer:  $\tilde{F}_T = FC(F_T)$ .  $F_T$  is a vector of dimension  $D$ , the FC layer reshapes it to the vector  $\tilde{F}_T$  of dimension  $D \times k \times k + 1$ , with " $k \times k$ " representing the convolution kernel size and "+" is the extra dimension accounting for bias. This allows decomposition of  $\tilde{F}_T$  into convolution weights  $w \in \mathbb{R}^{1 \times D \times k \times k}$  and bias  $b \in \mathbb{R}^1$ , which are subsequently used to convolve  $\tilde{F}_{I_A}$ ,
$$
S = \operatorname {S i g m o i d} \left(\operatorname {C o n v} \left(\tilde {F} _ {I _ {A}} | w, b\right)\right). \tag {1}
$$
Sigmoid is applied to transform the final output into probabilities, hence producing the score map  $S$ .
The attention-based prompting computes the attention at each pixel by interacting it with pixels from the whole feature map, while in the convolution-based prompting, each pixel only interacts with its neighboring pixels in the  $k \times k$  convolution kernel. We can see the former prompting as a global decoding process while the latter as a local decoding process to refine the former.
3.4 Mixture of Prompts
The provision of the prompt is crucial for the surgical instrument segmentation accuracy. Recognizing the impact of different prompts on a given model [56, 50], we propose to effectively leverage the strength of various prompts to facilitate the recognition of surgical instruments. Inspired by mixture of experts (MoE) [21], which divides a complex task into multiple sub-tasks and distributes them amongst multiple experts, we introduce a mixture of prompts (MoP) scheme to combine the segmentation outputs from multiple prompts via a visual-textual gating network.
Acquisition of prompts. There exist multiple strategies for generating prompts from simple to more advanced ones. We developed three types of prompt with increasing complexity. First, we use the name of each instrument class as the prompt. Second, inspired by CLIP's prompt template, we design a similar prompt template in the surgical context, i.e. "the surgical instrument area represented by the [class name]". Third, we leverage the recent released large language model GPT-4 [34] to generate prompts. We pose a question template: "Please describe the appearance of [class name] in endoscopic surgery, change the description to a phrase with a subject, and do not use colons." Please see supplementary material for the full list of our prompts.
Fusion of prompt predictions. Given the prompt set  $\mathcal{T} = \{T_1, T_2, T_3\}$  for an instrument to be segmented, we denote the corresponding textual feature for  $i$ -th prompt as  $F_T^i$ . Each  $F_T^i$  is paired the visual feature  $F_I$  to predict a score map  $S_i$  as outlined in Section 3.3, resulting in a set of score maps  $S = \{S_1, S_2, S_3\}$  for all prompts. We have parallelled the prompt processing by reshaping different
prompts to the batch dimension, exploiting the transformer structure to concurrently obtain  $S$ . To combine the score maps in  $S$ , we devise a visual-textual gating network to predict the pixel-wise weight maps.
Visual-textual gating network  $\mathcal{G}$ .  $\mathcal{G}$  consists of a 3-layer residual block [14]. It ingests both the visual feature  $F_{I}$  and textual feature  $F_{T}^{p}$ : we duplicate  $F_{T} \in \mathbb{R}^{1 \times D}$  to  $F_{T}^{p} \in \mathbb{R}^{N \times D}$  to match  $F_{I}$ 's dimension and concatenate them before inputting to  $\mathcal{G}$ . The outputs of  $\mathcal{G}$  are three weight maps corresponding to the three score maps in  $\mathcal{S}$ . These weights are normalized using softmax operations along the prompt dimension. We calculate the weighted sum of scores maps in  $\mathcal{S}$  to derive the final score map  $S^{\mathrm{et}} \in \mathbb{R}^{H \times W}$ . The final mask prediction  $\mathcal{M}$  is obtained from  $S^{\mathrm{et}}$  by thresholding it.
3.5 Hard Instrument Area Reinforcement
Current model falls short on distinguishing the accurate category and boundary of surgical instruments due to the complex surgical conditions (e.g. lighting variance) in endoscopy scenarios. We seek to reinforce its performance on the hard-predicted area in an image by utilizing a MAE-like structure [15] to reconstruct the image especially on hard-predicted area for representation enhancement. We propose a hard instrument area reinforcement module to bolster the model's segmentation accuracy for various instrument classes and their intricate details. Unlike MAE, which pays equal attention to all areas, our objective is to focus the visual encoder on discriminating the challenging instrument area. Therefore, we first introduce a hard area mining scheme to construct a masked image, then feed the visible patches into the decoder to reconstruct the whole image. This module shares the same image encoder with segmentation model.
Hard area mining. Given the estimated segmentation mask  $M^{\mathrm{et}}$ , we compare it with the ground truth mask  $M^{\mathrm{gt}}$  to obtain the falsely predicted area, which is taken as the hard instrument area for segmentation. For the input image  $I$ , we divide the area of hard regions over that of the whole image to obtain a masking ratio  $r \in (0,1)$ : 1) if  $r$  exceeds a pre-defined threshold  $r_t$ , we randomly unmask some masked pixels till the ratio falls to  $r_t$  to create the masked image; 2) if  $r$  is less than  $r_t$ , we randomly mask some more unmasked pixels till the ratio meets  $r_t$  to create the masked image.
Hart area reinforcement. We feed the obtained masked image into the image encoder and employ the same decoder structure as MAE to reconstruct the full image area based on the unmasked area. This helps the model to focus on recognizing the subtle details of the hard-predicted area.
It is worth mentioning that this module is no longer needed during testing.
3.6 Model Training
Our method is based on the pre-trained vision-language model, CLIP. During training, given limited surgical instrument descriptions, we freeze the text encoder and only finetune the image encoder. Two distinct loss functions,  $\mathcal{L}^{seg}$  for image segmentation and  $\mathcal{L}^{rec}$  for hard instrument area reinforcement, are introduced.
For  $\mathcal{L}^{seg}$ , binary cross entropy loss is applied:
$$
\mathcal {L} ^ {\text {s e g}} = - \sum_ {j} \left(M _ {j} ^ {\mathrm {g t}} \log M _ {j} ^ {\mathrm {e t}} + \left(1 - M _ {j} ^ {\mathrm {g t}}\right) \log \left(1 - M _ {j} ^ {\mathrm {e t}}\right)\right) \tag {2}
$$
where  $M_j^{\mathrm{gt}}$  and  $M_j^{\mathrm{et}}$  denote the labels of ground truth mask  $M^{\mathrm{gt}}$  and prediction mask  $M^{\mathrm{et}}$  at the  $j$ -th position, respectively.