Data Adaptive Few-shot Multi Label Segmentation with Foundation Model

Most of radiology tasks necessitate localization of anatomy or landmark, and lesion, which is laborious, repetitive and prone to errors, especially due to variations in patient pose, conditions, and disease. Advances in Artificial Intelligence(AI) has simplified the effort by utilizing various deep learning paradigms (supervised training, transfer learning, active learning) to automate the tasks for localization and landmarking. Most of these approaches are data driven and the burden and associated risks have shifted from end-application user to AI developer for getting data manually annotated for further usage.

There is emerging body of work on performing few shot segmentation aiming to train Deep Learning(DL) network with only few annotated datasets. Recent advancements in the ability of ViT based and Diffusion based models [1] for obtaining point correspondences have shown great promise in leveraging pixel level similarity for region correspondences between pairs of images [2]. A recent work [3] demonstrates the utility of using such feature correspondences for localization of landmarks in medical images as well as obtaining organ segmentation by chaining it with SAM [4].

However, most of these one-shot methods are based on 2D images and expect that the template or exemplar image are well matched to the new test image for correspondence matching. In our experience, such methods do not scale well for 3D imaging volume nor in cases where image representation between exemplars and test data are mismatched due to geometric changes, protocol changes or disease conditions. To overcome these limitations, in this paper, we propose to leverage sub-image level features derived from fine-tuned foundation models and by training a lightweight task adapter, using the derived pixel-features for obtaining multi-label localization which generalizes well across these challenges.

Methods exploring the use of self-supervised learning and transfer learning techniques for few shot learning have been shown to be effective. For instance, the paper [5] introduces double RoI heads into the existing Fast-RCNN to learn more specific features for novel classes. It also uses self-supervised learning to learn more structural and semantic information.

Specialized network architectures play a crucial role in few-shot segmentation by effectively learning from limited data and improving the model’s performance. Prototype networks have been shown to be superior for few-shot learning scenarios [6] [7], using support and query image paradigm, similar to our approach. The key idea is to decompose the holistic class representation into a set of part-aware prototypes, capable of capturing diverse and fine-grained object features. These networks often leverage large set of unlabeled data to enrich the part-aware prototypes [8], which though possible in natural images may be impractical for data-sparse and diverse situations - common in medical imaging domain.

Methods that use text prompts for image segmentation/localization [9], [10] are gaining significance area in computer vision. They leverage the power of natural language to guide the task allowing flexibility, interactivity and the potential to handle a wider range of tasks within the scope provided as part of supervised training. Consequently, these methods do not scale well for out-of-domains images and pointed landmarks identification tasks.

Recent works take advantage of the capability of ViT and diffusion models to extract patch level features [1] and leverage these to pixel correspondence tasks and extend these for localization and segmentation tasks by chaining them with other foundation models to obtain few shot leaning capability [3]. A similar technique [4] leverages latent sub-pixel level features Segment Anything Model (SAM [11]) and train a small Adapter for a single template image in order to obtain segmentation regions on other similar images.

Though the above listed methods are effective for 2D image tasks based on natural images, they do not extend naturally to 3D medical volumes with geometric, protocol and patient variations. Moreover, these methods are based on pixel feature extraction and simple similarity metrics thus requiring user to rely on task-dependent heuristics which need to be manually tuned.

Previous work [1], [3] on few-shot segmentation or localization have relied on template and target images being similar in structure and contrast as captured in feature semantics. To segment similar regions as marked in template, the features obtained at pixel level from ViT models are correlated using similarity metric and further thresholded to localize the region of interest [12]. However, obtaining this threshold is not trivial due to large variations in image intensity (including across slices) (See Fig. 1) and contrast (due to site based intra-protocol changes) observed in imaging data in clinical practice. This impacts the segmentation results in multiple ways: a) For given imaging volume in a subject, multiple templates have to be chosen to get reliable segmentation to accommodate intensity changes across slices and b) Inability to scale segmentation reliably across the diverse pool of imaging data due to various patient conditions or site specific imaging protocols. The need for manual tuning hinders practical adoption of template based few shot region segmentation and a method is acutely needed to overcome this.

In this work, we designed a contrastive learning based adapter which is trained to automatically determine which of the pixel level feature vectors in target image are similar to the template pixel level feature vector for given set of ground-truth label(s); thereby obviating the need for user to fix any threshold. We hypothesize that the proposed contrastive adapter is better at discriminating learnt representation even when intensity appearance/ViT features are almost homogeneous.

The main contribution of our work is to utilize the features provided by foundation models (fine-tuned DINOv2) to enable accurate segmentation/localization using one or few template images. To enable this in clinical settings we do the following a) train a Siamese architecture based contrastive learning adapter to learn robust feature matching metric b) ability to enable multi distinct region localization simultaneously using a single model, c) Ability to generate automatic prompts for SAM model to obtain refined segmentation and d) by virtue of this, ability to annotate 3D volumes with minimal templates.

Transformer architectures enable deriving patch level features for images, which can be interpolated to pixel level features [3]. Next, we utilize these pixel level features to correlate the region marked in template image with similar region in target image.

Basic adapter: A straightforward approach to obtain this would be to correlate template image pixel feature vector with the feature vectors for each of the pixels in the target image using cosine similarity measure and threshold the correlation map at a fixed heuristic threshold (= 0.5) to obtain the binary localization mask. This is referred to as basic adapter and involves no training. The basic adapter is not well suited to account for variations in clinical data or multi-label tasks and hence we design two more sophisticated trained adapters to handle these limitations, as described below.

Classification adapter: The Classification adapter is a lightweight feed-forward three layer neural network, taking as input the pixel features and trained with a cross-entropy loss function. During inference, pixel level features for all pixels in the target image are fed to the model to obtain the pixel-wise binary or multi-label localization mask. The localization mask is subject to connected component analysis to remove any stray/isolated pixels.We find that following this simple classification approach leads to sub-optimal results (especially in MRI data) as discussed in the experimental section.

We have trained both binary and multi-label localization classifiers as described below: a) Binary localization: We train a binary classifier to label pixels within an ROI as foreground and randomly selected pixels outside the ROI (matching the number of pixels in foreground to maintain class balance) as background shown in Fig. 2. b) Multi-label localization: We define a background as set of pixels not belonging to any of the multi-label ROI. We then train a multi-class classifier to label pixels belonging to different ROI labels and background (i.e number of classes = number of labels+1).

Contrastive adapter: We follow the Siamese architecture with a contrastive loss (cross-entropy loss) as described in Eq [8] from [13]. The network design is shown in Fig. 4(a) for binary localization and in Fig 4(b) for multi-label localization. It consists of two fully connected lightweight feed-forward neural network (Fig. 4) with shared weights. The outputs of the two sub-networks are concatenated and fed to the softmax layer.

a) Binary localization: We consider pairs of pixels to be positive if both belong to the same RoI. We create negative pair pixels by choosing one pixel from RoI and another random pixel from outside the RoI as shown in Fig. 3. Feature vectors corresponding to pairs of positive and negative pixels are passed through the sub-networks, as shown in the Fig. 4 (a), to predict how similar the input pair feature vectors are with each other. This is achieved by minimizing the distance between positive pairs and maximizing distance between negative pairs. We train the model with cross-entropy loss function to obtain the binary mask as localization map.

b) Multi-label localization (Fig. 4 (b)): For a given label RoI, we consider pairs of pixels to be positive if both belong to the same RoI. Negative pairs are created by choosing one-pixel in chosen label ROI and one outside the chosen label ROI. This process is repeated for all the labels in consideration. We train the contrastive model with these paired pixel features using contrastive loss to obtain the softmax probability distribution. The index of the softmax distribution with highest probability is associated with the corresponding label.

During inference for Contrastive adapter, as shown in Fig. 5, we choose reference pixels from each label RoI and compute the feature vectors and treat them as template feature vectors for that RoI. Similarly, for the target slice, we compute the feature vectors for all the pixels. We pair each template feature vector with all the feature vectors from the target image and pass the pairs to the appropriate contrastive model (binary or multi-label), which then assigns label(s) to each pixel in target image according to the highest probability value from the softmax distribution. Note from Fig. 5 that reference pixel feature vector is fed to the upper Feed Forward Neural network (FFN) while target pixel feature vector is fed to the lower FFN. The localization mask is subject to connected component analysis to remove any stray/isolated pixels. If the task is landmark localization, this is the final output. If the task is segmentation of structure, we randomly choose ten pixels from localized region and prompt the SAM to refine the segmentation of the object of interest. The output of SAM is the final segmentation task output.

The simultaneous multi-label localization using a single model using the proposed contrastive method is shown in the Fig. 6, using a single adapter by comparison with pixel from RoI in source image.

We evaluate the efficacy of the proposed few-shot method for five different medical 3D image tasks, two of them for segmentation and three for localization. To benchmark the performance, we compare the results with recently proposed and popular few shot methods such as UniverSeg [14] and PerSAM [4]. Both of these methods enable users to tackle new segmentation task without the need to train or fine-tune a model, adapting to a new segmentation task at inference based on few sample examples.

For PerSAM and basic adapter based localization methods (which are designed for binary tasks), for multi-label tasks, we aggregate the predictions using each label in the template pool.

We introduce a method for few-shot multi-label localization and segmentation of medical image volumes, utilizing robust contextual feature vectors extracted from transformer models trained on medical data. These models use a few template image for localization of all volumes to produce segmentation or localization outputs. To account for variability within the data and consequently feature vectors which impact heuristic threshold based methods, for initial localization of images, we proposed a contrastive similarity metric learning model, extending it by adopting a multi-label contrastive strategy for enabling tasks with multiple labels. As compared to the state-of-the-art methods for few-shot segmentation/localization not requiring in-domain training, the proposed contrastive adapter approach is superior for both segmentation and localization tasks.