Probabilistic Feature Matching for Fast Scalable Visual Prompting

1 Introduction

Foundation Models ushered in a significant shift in how machine learning models are developed and deployed, pivoting from a paradigm centered on training use case-tailored models on task-specific data to a paradigm where single generalist models are pretrained on diverse large-scale data, then finetuned for a wide range of tasks [Bommasani et al., 2022]. Specifically in computer vision, models such as SAM [Kirillov et al., 2023], CLIP [Radford et al., 2021], and selfsupervised backbones such as DINO [Caron et al., 2021] and DINOv2 [Oquab et al., 2023] have unlocked powerful and versatile visual functionalities like object detection, semantic segmentation and expressive embeddings that are at the core of a multitude of diverse applications. In particular, the possibility of using and combining these models in novel ways to address specific challenges in applied computer vision has been a topic of recent interest, including as a means to design new workflows in technical domains such as visual inspection (see e.g. [Rigotti et al., 2023]).

† corresponding authors

In this work we take inspiration from the recent advancements driven by the approach of compositionally combining multiple Foundation Models to address sophisticated computer vision tasks. Specifically, we focus on the problem of image segmentation, which is a fundamental task in computer vision with a wide range of applications, including medical imaging, autonomous driving, and visual inspection, with a particular focus in developing a human-computer interaction workflow to facilitate open-world segmentation of images by visual prompting through sparse user annotations. For that, we largely build upon a previous architecture named Matcher, which was designed to perform training-free fewshot segmentation using in-context examples by means of offthe-shelf vision Foundation Models [Liu et al., 2023]. Our framework enhances this approach's interactivity in two crucial ways: 1) we integrate a pretrained segmentation model to translate user scribbles on a representative sample of the object class to be segmented into reference masks, which are then passed to the few-shot segmentation architecture; 2) we develop a scalable probabilistic feature soft-matching procedure whose efficiency and low-latency allows us to embed few-shot segmentation in a real-time interactive workflow.

2 RelatedWork

The Segment Anything Model (SAM) [Kirillov et al., 2023] has popularized the prompting paradigm in computer vision by enabling fine-grained image segmentation through interactive prompts in the form of points and/or bounding boxes.

Both Visual Prompting via Inpainting [Bar et al., 2022] and SegGPT/Painter [Wang et al., 2023] presented visual prompting models trained on few-shot image segmentation datasets. These models operate on a reference image and corresponding segmentation masks and generate a segmentation mask for a target image based on the reference.

[Zhang et al., 2023] introduced a training-free method for one-shot segmentation leveraging pretrained image encoders in conjunction with SAM. The labeled pixels within the annotated mask on a reference image are assigned to pixels on target images thanks to a cosine similarity matrix of their corresponding encoded patches. The target patch of maximum similarity is then utilized by SAM to generate a segmentation mask for the target object.

[Gupta and Kembhavi, 2022] presented a neuro-symbolic approach for solving complex visual tasks given natural language instructions by leveraging the in-context learning ability of LLMs to generate modular programs that combine pretrained models leveraging their compositionality, a feature that has received recent interest for enabling flexible generalization (see e.g. [Ito et al., 2022]).

[Liu et al., 2023] introduced Matcher, an approach that uses a bidirectional matching procedure to match the encoded reference and target image patches using the Hungarian algorithm, an accurate but slow assignment algorithm with worstcase complexity cubic in the size of the problem [Crouse, 2016]. Similarly to [Zhang et al., 2023], one-shot (or fewshot) segmentation is implemented by assigning annotated encoded pixels on reference images to encoded target pixels, which then serve as prompts for SAM to produce segmentation mask proposals on the target images. The set of mask proposals is finally scored and either accepted or rejected.

[Janouskova et al., 2023] proposed a framework for modelassisted labeling of visual inspection defects through an interactive annotation process leveraging gradient-based explainability to improve the efficiency of the provided labels.

3 Visual Prompting Framework

System architecture. Figure 1 presents our Sofmatcher framework for interactive image segmentation guided by visual prompting on a reference image. This consists of 3 steps: 1) Prompting & reference segmentation, where a user provides scribbles on the reference image indicating the object class to be labeled on the target images, and where the scribbles are used as sparse prompt for SAM which then is used to output a reference mask; 2) Matching, where soft probabilistic matching (detailed below) outputs a probability map over pixels of each target image quantifying their match to pixels in the reference mask; points are then sampled from the probability map, clustered and used for 3) Mask generation, where clustered points are used as sparse prompts to SAM to generate mask proposals; these are filtered based on SAM's IoU predictions and aggregated into the mask output.

The key innovations of our framework compared to previous approaches like Matcher [Liu et al., 2023] are aimed at producing an architecture that is amenable to being embedded in an interactive object segmentation workflow where users can provide visual prompts by coarsely annotating reference images through scribbles and interact in real-time with the resulting segmentation masks, possibly by correcting or complementing them with additional annotations.

Our first innovation for this is the Prompting & reference segmentation step in Fig. 1, which, while conceptually simple, provides a way for the user to directly and intuitively prompt the segmentation pipeline with coarse visual prompts (scribbles) instead of requiring detailed segmentation masks.

Our second major innovation is a computationally efficient version of the Matching step in Fig. 1, and was dictated by the requirement of low-latency segmentation and the observation that feature matching procedure used in the past, like the Hungarian algorithm (see e.g. [Liu et al., 2023]), display a worst-case computational complexity that scales cubically with image sizes (number of patches) [Crouse, 2016], making them unpractical for an interactive workflow. Instead of using (Hungarian) bipartite matching based on the cosine similarity between reference and target features, we opt for an Optimal Transport (OT) approach based on the quadratic cosine similarity matrix as a cost matrix. While very related, this method allows us to motivate a sequence of approximations for an efficient implementation of the matching procedure: we first introduce an entropic regularization, then consider the case of large regularization limit where the solution to the OT problem converges to the geometric mean of softmaxed cosine similarity maps between individual reference features and target feature maps (where the averaging is conducting across reference features) [Dognin et al., 2019], an operation which only has quadratic complexity in the number of image patches complexity and results in our Softmatcher procedure. Moreover, it affords an even more scalable implementation by approximating the softmax computation of reference-target feature similarities through Random Fourier Features [Rahimi and Recht, 2007; Choromanski et al., 2020], which we call Softmatcher RFF.

Figure 2 compares the timing of matching reference and target image features with the Hungarian algorithm, compared to our proposed soft matching methods as a function of image size assuming a featurization based on a VIT encoder with patch size of 14, feature size of 768. Softmatcher is around 6x faster than the Hungarian algorithm at image size 448, and this discrepancy quickly increases with image size due to its better computation complexity scaling. Softmatcher RFF is slightly faster and displays even better scalability.

We evaluate our visual prompting pipeline on FSS-1000 [Li et al., 2020], which consists of 1000 object classes with pixel-wise annotations. FSS-1000 contains many objects that are not part of any previously annotated dataset (e.g., tiny daily objects, merchandise, and cartoon characters). As this disentangles previous knowledge from pretrained models to a certain degree, it lends itself well as a few-shot benchmark.

We integrate this improved matching pipeline into an interactive Visual Prompting platform that allows users to segment object classes of interest by merely highlighting representative objects in one or more reference images with scribbles. Given the improved computation complexity, our method allows the user to iterate in real-time with the segmentation outputs, adding additional scribbles on additional references to improve segmentation in case the model missed something, resulting in an intuitive and seamlessly interactive workflow.s Deployed service and front-end. The interactive web interface is designed to provide seamless interaction between the user and the Softmatcher pipeline. It consists of a frontend built with Angular, a Python API back-end, and an inference service using Torch Serve. Users add scribbles to any image to mark objects of interest. The visual prompting pipeline then highlights similar objects with precise segmentation masks the target images. If the user is not satisfied with the initial results, they can refine the outputs by iteratively adding or deleting scribbles. Alternatively, instead of adding more scribbles, users can add additional prompts by converting output segmentation masks from a previous run into reference masks. These reference masks will skip step 1 of the pipeline (see Fig. 1). The system also allows for scribbles to be classified into different categories, enabling the creation of segmentation masks for multiple classes.

The process of repeatedly adding and adjusting scribbles provides users with a deeper understanding of how the model operates. By understanding the model's capabilities and limitations, users learn to collaborate with the model more effectively, leading to better outcomes. We've also started to enhance our framework's interactivity with vision-language models like CLIP, enabling the use of text prompts in addition to reference scribbles. This opens up the possibility of combining visual and text prompts to refine masks mutually and address scenarios where scribbling alone is not enough.

Demonstration. We illustrate how users typically engage with our web interface and the visual prompting pipeline through three sample projects. The first two projects illustrate a general use case on everyday objects, while the third shows a domain-specific proprietary defect detection dataset. Our demonstration covers the interactive process of adding scribbles to images, executing the pipeline to receive segmentation masks, and then enhancing the results by adding additional scribbles. Furthermore, we showcase the capability for users to process images with references from various classes.

Acknowledgments

We would like to thank Trafikverket for their collaboration, specifically for collecting image data of civil infrastructures. We are also grateful to Youssef Mroueh for insightful discussions. This work is partly funded by the European Union's Horizon Europe research and innovation program under grant agreements No. 101070408 (SustainML) and the Swiss State Secretariat for Education, Research and Innovation (SERI) under contract number 22.00295.